diff --git a/.buildkite/ci_config.yaml b/.buildkite/ci_config.yaml index b199e554a73..a60a4194e9b 100644 --- a/.buildkite/ci_config.yaml +++ b/.buildkite/ci_config.yaml @@ -8,8 +8,8 @@ run_all_patterns: - "CMakeLists.txt" - "requirements/common.txt" - "requirements/cuda.txt" - - "requirements/build.txt" - - "requirements/test.txt" + - "requirements/build/cuda.txt" + - "requirements/test/cuda.txt" - "setup.py" - "csrc/" - "cmake/" diff --git a/.buildkite/ci_config_intel.yaml b/.buildkite/ci_config_intel.yaml index 375be84a396..a1c0091e0f1 100644 --- a/.buildkite/ci_config_intel.yaml +++ b/.buildkite/ci_config_intel.yaml @@ -6,8 +6,8 @@ run_all_patterns: - "CMakeLists.txt" - "requirements/common.txt" - "requirements/xpu.txt" - - "requirements/build.txt" - - "requirements/test.txt" + - "requirements/build/cuda.txt" + - "requirements/test/cuda.txt" - "setup.py" - "csrc/" - "cmake/" diff --git a/.buildkite/image_build/image_build_torch_nightly.sh b/.buildkite/image_build/image_build_torch_nightly.sh new file mode 100755 index 00000000000..a23c658d46b --- /dev/null +++ b/.buildkite/image_build/image_build_torch_nightly.sh @@ -0,0 +1,68 @@ +#!/bin/bash +set -euo pipefail + +# Build a vLLM test image with PyTorch nightly installed. +# Called by the pipeline generator's "vLLM Against PyTorch Nightly" group. + +if [[ $# -lt 5 ]]; then + echo "Usage: $0 " + exit 1 +fi + +REGISTRY=$1 +REPO=$2 +BUILDKITE_COMMIT=$3 +BRANCH=$4 +IMAGE_TAG=$5 + +# --- Arguments --- +echo "--- :mag: Arguments" +echo "REGISTRY: ${REGISTRY}" +echo "REPO: ${REPO}" +echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}" +echo "BRANCH: ${BRANCH}" +echo "IMAGE_TAG: ${IMAGE_TAG}" + +# --- ECR login --- +echo "--- :key: ECR login" +aws ecr-public get-login-password --region us-east-1 \ + | docker login --username AWS --password-stdin "$REGISTRY" +aws ecr get-login-password --region us-east-1 \ + | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com + +# --- Set up buildx --- +echo "--- :docker: Setting up buildx" +docker buildx create --name vllm-builder --driver docker-container --use || true +docker buildx inspect --bootstrap +docker buildx ls + +# --- Skip if image already exists --- +echo "--- :mag: Checking if image already exists" +if docker manifest inspect "$IMAGE_TAG" >/dev/null 2>&1; then + echo "Image found: $IMAGE_TAG — skipping build" + exit 0 +fi +echo "Image not found, proceeding with build..." + +# --- CUDA 13.0 for nightly builds --- +# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9 +NIGHTLY_CUDA_VERSION="13.0.0" +NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04" +NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04" + +echo "--- :docker: Building torch nightly image (CUDA ${NIGHTLY_CUDA_VERSION})" +docker buildx build --file docker/Dockerfile \ + --build-arg max_jobs=16 \ + --build-arg buildkite_commit="$BUILDKITE_COMMIT" \ + --build-arg USE_SCCACHE=1 \ + --build-arg PYTORCH_NIGHTLY=1 \ + --build-arg CUDA_VERSION="${NIGHTLY_CUDA_VERSION}" \ + --build-arg BUILD_BASE_IMAGE="${NIGHTLY_BUILD_BASE_IMAGE}" \ + --build-arg FINAL_BASE_IMAGE="${NIGHTLY_FINAL_BASE_IMAGE}" \ + --build-arg torch_cuda_arch_list="8.0 8.9 9.0 10.0 12.0" \ + --tag "$IMAGE_TAG" \ + --push \ + --target test \ + --progress plain . + +echo "--- :white_check_mark: Torch nightly image build complete: $IMAGE_TAG" diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 295840440ad..c59be699502 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.buildkite/intel_jobs/test-intel.yaml @@ -35,6 +35,7 @@ steps: python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN && python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 && + python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8 && python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 && python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 && python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel' @@ -61,4 +62,4 @@ steps: pytest -v -s v1/structured_output && pytest -v -s v1/test_serial_utils.py && pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py && - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py' + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py' diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index 45b2996f7ea..b3a6bb8ed4c 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -98,8 +98,15 @@ steps: commands: - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + - block: "Unblock to build release Docker images" + depends_on: ~ + key: block-build-release-images + if: build.env("NIGHTLY") != "1" + - group: "Build release Docker images" key: "build-release-images" + depends_on: block-build-release-images + allow_dependency_failure: true steps: - label: "Build release image - x86_64 - CUDA 12.9" depends_on: ~ @@ -617,6 +624,8 @@ steps: - label: ":docker: Build release image - x86_64 - ROCm" id: build-rocm-release-image depends_on: + - step: block-build-release-images + allow_failure: true - step: build-rocm-base-wheels allow_failure: false agents: diff --git a/.buildkite/scripts/generate-and-upload-nightly-index.sh b/.buildkite/scripts/generate-and-upload-nightly-index.sh index fa6eb979af5..7cef252c607 100755 --- a/.buildkite/scripts/generate-and-upload-nightly-index.sh +++ b/.buildkite/scripts/generate-and-upload-nightly-index.sh @@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) if [[ "$has_new_python" -eq 0 ]]; then # use new python from docker docker pull python:3-slim - PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3" + PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3" fi echo "Using python interpreter: $PYTHON" diff --git a/.buildkite/scripts/hardware_ci/run-hpu-test.sh b/.buildkite/scripts/hardware_ci/run-hpu-test.sh index 10df07b2000..0b5d0af4b6f 100644 --- a/.buildkite/scripts/hardware_ci/run-hpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-hpu-test.sh @@ -42,7 +42,7 @@ WORKDIR /workspace/vllm ENV no_proxy=localhost,127.0.0.1 ENV PT_HPU_ENABLE_LAZY_COLLECTIVES=true -RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build.txt)' +RUN bash -c 'pip install -r <(sed "/^torch/d" requirements/build/cuda.txt)' RUN VLLM_TARGET_DEVICE=empty pip install --no-build-isolation -e . RUN pip install git+https://github.com/vllm-project/vllm-gaudi.git diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh index 0467da7c230..6579810e982 100644 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -50,6 +50,6 @@ docker run \ pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py pytest -v -s v1/structured_output pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py - pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py -k "not (test_register_kv_caches and FLASH_ATTN and True)" + pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py pytest -v -s v1/test_serial_utils.py ' diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c7e9e13f2fd..dda5d4064c3 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -123,7 +123,7 @@ steps: soft_fail: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - requirements/nightly_torch_test.txt + - requirements/test/nightly-torch.txt - vllm/platforms/rocm.py commands: - bash standalone_tests/pytorch_nightly_dependency.sh @@ -532,28 +532,6 @@ steps: - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: V1 Speculative Decoding (slow) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/model_executor/models/ - - vllm/v1/attention/ - - vllm/model_executor/layers/ - - tests/v1/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py - - - label: V1 attention (H100-MI250) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] @@ -1073,7 +1051,8 @@ steps: - tests/models/multimodal/test_mapping.py commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing + - 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 @@ -1878,28 +1857,6 @@ steps: - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: V1 Speculative Decoding (slow) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325] - agent_pool: mi325_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/model_executor/models/ - - vllm/v1/attention/ - - vllm/model_executor/layers/ - - tests/v1/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py - - label: Acceptance Length Test (Large Models) # TBD timeout_in_minutes: 180 @@ -1914,7 +1871,7 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test + - pytest -v -s v1/spec_decode/test_acceptance_length.py - label: V1 attention (H100-MI325) # 14.5m @@ -2298,7 +2255,8 @@ steps: - 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/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -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 @@ -3186,28 +3144,6 @@ steps: - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine -- label: V1 Speculative Decoding (slow) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] - agent_pool: mi355_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/model_executor/models/ - - vllm/v1/attention/ - - vllm/model_executor/layers/ - - tests/v1/spec_decode/ - - vllm/platforms/rocm.py - commands: - - pytest -v -s -m 'slow_test' v1/spec_decode/test_eagle.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_extract_hidden_states.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_max_len.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_mtp.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_ngram.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_speculators_eagle3.py - - pytest -v -s -m 'slow_test' v1/spec_decode/test_tree_attention.py - - - label: V1 attention (B200-MI355) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] @@ -3452,7 +3388,8 @@ steps: - 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/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest -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 diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 19815e14b4e..56e528ffb37 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -268,6 +268,20 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh + - label: NixlConnector PD + Spec Decode acceptance (2 GPUs) timeout_in_minutes: 30 device: a100 @@ -281,6 +295,20 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt - bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh +- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) + timeout_in_minutes: 30 + working_dir: "/vllm-workspace/tests" + num_devices: 2 + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - tests/v1/kv_connector/nixl_integration/ + commands: + - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt + - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + - label: Pipeline + Context Parallelism (4 GPUs) timeout_in_minutes: 60 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 362c733c8d7..8b9765130ae 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -20,7 +20,20 @@ steps: - tests/kernels/core - tests/kernels/test_concat_mla_q.py commands: - - pytest -v -s kernels/core kernels/test_concat_mla_q.py + - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py + +- label: Kernels MiniMax Reduce RMS Test (2 GPUs) + timeout_in_minutes: 15 + num_devices: 2 + device: h100 + source_file_dependencies: + - csrc/minimax_reduce_rms_kernel.cu + - csrc/minimax_reduce_rms_kernel.h + - vllm/model_executor/layers/mamba/linear_attn.py + - vllm/model_executor/layers/mamba/lamport_workspace.py + - tests/kernels/core/test_minimax_reduce_rms.py + commands: + - pytest -v -s kernels/core/test_minimax_reduce_rms.py - label: Kernels Attention Test %N timeout_in_minutes: 35 @@ -187,7 +200,14 @@ steps: timeout_in_minutes: 90 device: h100 num_devices: 2 - optional: true + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/device_communicators/ + - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py @@ -196,6 +216,13 @@ steps: timeout_in_minutes: 90 device: b200 num_devices: 2 - optional: true + source_file_dependencies: + - csrc/quantization/cutlass_w8a8/moe/ + - csrc/moe/ + - tests/kernels/moe + - vllm/model_executor/layers/fused_moe/ + - vllm/model_executor/layers/quantization/ + - vllm/distributed/device_communicators/ + - vllm/config commands: - pytest -v -s kernels/moe/test_moe_layer.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 39029efe9cd..a07d702cf3c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -91,6 +91,16 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt +- label: LM Eval TurboQuant KV Cache + timeout_in_minutes: 75 + source_file_dependencies: + - vllm/model_executor/layers/quantization/turboquant/ + - vllm/v1/attention/backends/turboquant_attn.py + - vllm/v1/attention/ops/triton_turboquant_decode.py + - vllm/v1/attention/ops/triton_turboquant_store.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt + - label: GPQA Eval (GPT-OSS) (H100) timeout_in_minutes: 120 device: h100 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 02c42568872..97906c81c29 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -224,7 +224,8 @@ steps: - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_MLA - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - + - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py + - label: Acceptance Length Test (Large Models) # optional timeout_in_minutes: 25 gpu: h100 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 8ba4484fb02..10b038d8b8a 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -1,5 +1,5 @@ group: Models - Basic -depends_on: +depends_on: - image-build steps: - label: Basic Models Tests (Initialization) @@ -13,10 +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 timeout_in_minutes: 45 - torch_nightly: true source_file_dependencies: - vllm/model_executor/models/ - tests/models/test_initialization.py @@ -27,6 +28,8 @@ steps: # test.) Also run if model initialization test file is modified - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Basic Models Tests (Other) timeout_in_minutes: 45 @@ -42,10 +45,10 @@ steps: device: mi325_1 depends_on: - image-build-amd - + - label: Basic Models Test (Other CPU) # 5min - depends_on: + depends_on: - image-build-cpu timeout_in_minutes: 10 source_file_dependencies: diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 7eac9e30193..c13371e25f1 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -1,10 +1,9 @@ group: Models - Language -depends_on: +depends_on: - image-build steps: - label: Language Models Tests (Standard) timeout_in_minutes: 25 - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/language @@ -12,10 +11,11 @@ steps: # Test standard language models, excluding a subset of slow tests - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and (not slow_test)' + mirror: + torch_nightly: {} - label: Language Models Tests (Extra Standard) %N timeout_in_minutes: 45 - torch_nightly: true source_file_dependencies: - vllm/model_executor/models/ - tests/models/language/pooling/test_embedding.py @@ -27,10 +27,11 @@ steps: - pip freeze | grep -E 'torch' - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Language Models Tests (Hybrid) %N timeout_in_minutes: 75 - torch_nightly: true source_file_dependencies: - vllm/ - tests/models/language/generation @@ -42,6 +43,8 @@ steps: # Shard hybrid language model tests - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB parallelism: 2 + mirror: + torch_nightly: {} - label: Language Models Test (Extended Generation) # 80min timeout_in_minutes: 110 @@ -62,7 +65,7 @@ steps: - image-build-amd commands: - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr' - - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.5.2' + - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0' - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)' - label: Language Models Test (PPL) diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index 3bf907bb6c5..ff0fd2e7a62 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -56,7 +56,8 @@ steps: - tests/models/multimodal commands: - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git - - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/processing + - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing + - pytest models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work mirror: amd: diff --git a/.buildkite/test_areas/pytorch.yaml b/.buildkite/test_areas/pytorch.yaml index ad538e91918..a3648219d89 100644 --- a/.buildkite/test_areas/pytorch.yaml +++ b/.buildkite/test_areas/pytorch.yaml @@ -64,6 +64,6 @@ steps: device: h200_18gb soft_fail: true source_file_dependencies: - - requirements/nightly_torch_test.txt + - requirements/test/nightly-torch.txt commands: - bash standalone_tests/pytorch_nightly_dependency.sh diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index 0a395ea5588..a42d59b021c 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -9,7 +9,7 @@ steps: - vllm/model_executor/layers/quantization - tests/quantization commands: - # temporary install here since we need nightly, will move to requirements/test.in + # temporary install here since we need nightly, will move to requirements/test/cuda.in # after torchao 0.12 release, and pin a working version of torchao nightly here # since torchao nightly is only compatible with torch nightly currently diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index a0b73096867..76cc887ed0a 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -42,3 +42,16 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference" + +- label: DFlash Speculators Correctness + timeout_in_minutes: 30 + device: h100 + optional: true + num_devices: 1 + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - tests/v1/spec_decode/test_speculators_dflash.py + commands: + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index a9c1c17128b..e1a2a582083 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,7 +3,7 @@ # This lists cover the "core" components of vLLM that require careful review /vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy -/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery +/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @@ -120,16 +120,16 @@ mkdocs.yaml @hmellor /tools/pre_commit @hmellor # CPU -/vllm/v1/worker/cpu* @bigPYJ1151 +/vllm/v1/worker/cpu* @bigPYJ1151 @xuechendi /csrc/cpu @bigPYJ1151 -/vllm/platforms/cpu.py @bigPYJ1151 +/vllm/platforms/cpu.py @bigPYJ1151 @xuechendi /cmake/cpu_extension.cmake @bigPYJ1151 -/docker/Dockerfile.cpu @bigPYJ1151 +/docker/Dockerfile.cpu @bigPYJ1151 @xuechendi # Intel GPU -/vllm/v1/worker/xpu* @jikunshang -/vllm/platforms/xpu.py @jikunshang -/docker/Dockerfile.xpu @jikunshang +/vllm/v1/worker/xpu* @jikunshang @xuechendi +/vllm/platforms/xpu.py @jikunshang @xuechendi +/docker/Dockerfile.xpu @jikunshang @xuechendi # Nemotron-specific files /vllm/model_executor/models/*nemotron* @tomeras91 diff --git a/.github/mergify.yml b/.github/mergify.yml index e8ef4d49dd0..baf65e14a88 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -18,7 +18,7 @@ pull_request_rules: - name: comment-pre-commit-failure description: Comment on PR when pre-commit check fails conditions: - - status-failure=pre-commit + - check-failure=pre-commit - -closed - -draft actions: @@ -51,7 +51,7 @@ pull_request_rules: - name: comment-dco-failure description: Comment on PR when DCO check fails conditions: - - status-failure=dco + - check-failure=dco - -closed - -draft actions: @@ -83,8 +83,8 @@ pull_request_rules: - or: - files~=^examples/.*deepseek.*\.py - files~=^tests/.*deepseek.*\.py - - files~=^vllm/entrypoints/openai/tool_parsers/.*deepseek.*\.py - files~=^vllm/model_executor/models/.*deepseek.*\.py + - files~=^vllm/tool_parsers/.*deepseek.*\.py - files~=^vllm/reasoning/.*deepseek.*\.py - files~=^vllm/transformers_utils/.*deepseek.*\.py - title~=(?i)DeepSeek @@ -110,9 +110,10 @@ pull_request_rules: - or: - files~=^examples/.*llama.*\.py - files~=^tests/.*llama.*\.py - - files~=^vllm/entrypoints/openai/tool_parsers/llama.*\.py - files~=^vllm/model_executor/models/.*llama.*\.py - - files~=^vllm/transformers_utils/configs/.*llama.*\.py + - files~=^vllm/reasoning/.*llama.*\.py + - files~=^vllm/tool_parsers/.*llama.*\.py + - files~=^vllm/transformers_utils/.*llama.*\.py - title~=(?i)llama actions: label: @@ -133,6 +134,23 @@ pull_request_rules: add: - multi-modality +- name: label-mistral + description: Automatically apply mistral label + conditions: + - label != stale + - or: + - files~=^examples/.*mistral.*\.py + - files~=^tests/.*mistral.*\.py + - files~=^vllm/model_executor/models/.*mistral.*\.py + - files~=^vllm/reasoning/.*mistral.*\.py + - files~=^vllm/tool_parsers/.*mistral.*\.py + - files~=^vllm/transformers_utils/.*mistral.*\.py + - title~=(?i)Mistral + actions: + label: + add: + - mistral + - name: label-new-model description: Automatically apply new-model label conditions: @@ -167,7 +185,9 @@ pull_request_rules: - files~=^examples/.*qwen.*\.py - files~=^tests/.*qwen.*\.py - files~=^vllm/model_executor/models/.*qwen.*\.py + - files~=^vllm/tool_parsers/.*qwen.*\.py - files~=^vllm/reasoning/.*qwen.*\.py + - files~=^vllm/transformers_utils/.*qwen.*\.py - title~=(?i)Qwen actions: label: @@ -244,6 +264,7 @@ pull_request_rules: - files=\.buildkite/ci_config_intel.yaml - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py + - files=vllm/model_executor/kernels/linear/mxfp8/xpu.py - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py - files=vllm/distributed/device_communicators/xpu_communicator.py - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py @@ -251,6 +272,7 @@ pull_request_rules: - files=vllm/v1/worker/xpu_worker.py - files=vllm/v1/worker/xpu_model_runner.py - files=vllm/_xpu_ops.py + - files=vllm/kernels/xpu_ops.py - files~=^vllm/lora/ops/xpu_ops - files=vllm/lora/punica_wrapper/punica_xpu.py - files=vllm/platforms/xpu.py @@ -258,7 +280,6 @@ pull_request_rules: - title~=(?i)XPU - title~=(?i)Intel - title~=(?i)BMG - - title~=(?i)Arc actions: label: add: @@ -378,17 +399,18 @@ pull_request_rules: add: - tool-calling -- name: auto-rebase if approved, ready, and 40 commits behind main +- name: auto-rebase to keep merge candidate within 1 day behind main conditions: - base = main - label=ready - "#approved-reviews-by >= 1" - - "#commits-behind >= 40" + - "#commits-behind >= 50" + - "#check-failure = 0" - -closed - -draft - -conflict actions: - rebase: {} + update: {} - name: ping author on conflicts and add 'needs-rebase' label conditions: diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 2cb5c176ae0..3efa582f670 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -320,20 +320,25 @@ jobs: script: | // Configuration: Map labels to GitHub users to CC // You can add multiple users per label, and multiple label configurations + // {users} will be replaced with @mentions const ccConfig = { rocm: { - users: ['hongxiayang', 'tjtanaa', 'vllmellm'], // Add more users as needed: ['user1', 'user2', 'user3'] - message: 'CC {users} for ROCm-related issue' // {users} will be replaced with @mentions + users: ['hongxiayang', 'tjtanaa', 'vllmellm'], + message: 'CC {users} for ROCm-related issue', + }, + mistral: { + users: ['patrickvonplaten', 'juliendenize', 'andylolu2'], + message: 'CC {users} for Mistral-related issue', }, // Add more label -> user mappings here // Example: // cuda: { // users: ['user1', 'user2'], - // message: 'CC {users} for CUDA-related issue' + // message: 'CC {users} for CUDA-related issue', // }, // performance: { // users: ['perfexpert'], - // message: 'CC {users} for performance issue' + // message: 'CC {users} for performance issue', // }, }; diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index 3c1a50bf808..b8e7cebf2dc 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -32,7 +32,7 @@ jobs: - name: Install dependencies and build vLLM run: | - uv pip install -r requirements/cpu-build.txt --index-strategy unsafe-best-match + uv pip install -r requirements/build/cpu.txt --index-strategy unsafe-best-match uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match uv pip install -e . --no-build-isolation env: diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index bbc55e9532b..8ab8d3e7035 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -2,6 +2,7 @@ name: pre-commit on: pull_request: + types: [opened, synchronize, reopened, labeled] push: branches: [main] @@ -15,7 +16,11 @@ permissions: jobs: pre-run-check: - if: github.event_name == 'pull_request' + if: >- + github.event_name == 'pull_request' && + (github.event.action != 'labeled' || + github.event.label.name == 'ready' || + github.event.label.name == 'verified') runs-on: ubuntu-latest steps: - name: Check PR label and author merge count @@ -44,7 +49,12 @@ jobs: pre-commit: needs: pre-run-check - if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') + if: >- + always() && + (github.event.action != 'labeled' || + github.event.label.name == 'ready' || + github.event.label.name == 'verified') && + (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped') runs-on: ubuntu-latest steps: - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 diff --git a/.github/workflows/scripts/build.sh b/.github/workflows/scripts/build.sh index c69ebbb42da..30c63ee0da9 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 -$python_executable -m pip install -r requirements/build.txt -r requirements/cuda.txt +$python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt # Limit the number of parallel jobs to avoid OOM export MAX_JOBS=1 diff --git a/.gitignore b/.gitignore index 7b822165d3e..134bbc5cc89 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,9 @@ vllm/third_party/triton_kernels/* # FlashMLA interface copied from source vllm/third_party/flashmla/flash_mla_interface.py +# DeepGEMM vendored package built from source +vllm/third_party/deep_gemm/ + # triton jit .triton @@ -26,6 +29,7 @@ __pycache__/ # Distribution / packaging .Python build/ +!requirements/build/ cmake-build-*/ CMakeUserPresets.json develop-eggs/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index f55df24bc6f..33b1db69dec 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -39,15 +39,24 @@ repos: rev: 0.11.1 hooks: - id: pip-compile - args: [requirements/test.in, -c, requirements/common.txt, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu130, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"] - files: ^requirements/test\.(in|txt)$ + args: [ + requirements/test/cuda.in, + -c, requirements/cuda.txt, + -o, requirements/test/cuda.txt, + --index-strategy, unsafe-best-match, + --torch-backend, cu130, + --python-platform, x86_64-manylinux_2_28, + --python-version, "3.12", + ] + files: ^requirements/(common|cuda|test/cuda)\.(in|txt)$ - id: pip-compile alias: pip-compile-rocm name: pip-compile-rocm args: [ - requirements/rocm-test.in, -o, requirements/rocm-test.txt, - --index-strategy, unsafe-best-match, + requirements/test/rocm.in, -c, requirements/rocm.txt, + -o, requirements/test/rocm.txt, + --index-strategy, unsafe-best-match, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12", # Exclude torch and CUDA/NVIDIA packages @@ -59,30 +68,76 @@ repos: --no-emit-package, cuda-pathfinder, --no-emit-package, cuda-toolkit, --no-emit-package, cupy-cuda12x, + # nvidia packages (unsuffixed / unified naming) --no-emit-package, nvidia-cublas, --no-emit-package, nvidia-cuda-cupti, --no-emit-package, nvidia-cuda-nvrtc, --no-emit-package, nvidia-cuda-runtime, - --no-emit-package, nvidia-cudnn-cu13, + --no-emit-package, nvidia-cudnn, --no-emit-package, nvidia-cufft, --no-emit-package, nvidia-cufile, --no-emit-package, nvidia-curand, --no-emit-package, nvidia-cusolver, --no-emit-package, nvidia-cusparse, + --no-emit-package, nvidia-cusparselt, + --no-emit-package, nvidia-nccl, + --no-emit-package, nvidia-nvjitlink, + --no-emit-package, nvidia-nvshmem, + --no-emit-package, nvidia-nvtx, + # nvidia cu12 packages + --no-emit-package, nvidia-cublas-cu12, + --no-emit-package, nvidia-cuda-cupti-cu12, + --no-emit-package, nvidia-cuda-nvrtc-cu12, + --no-emit-package, nvidia-cuda-runtime-cu12, + --no-emit-package, nvidia-cudnn-cu12, + --no-emit-package, nvidia-cufft-cu12, + --no-emit-package, nvidia-cufile-cu12, + --no-emit-package, nvidia-curand-cu12, + --no-emit-package, nvidia-cusolver-cu12, + --no-emit-package, nvidia-cusparse-cu12, + --no-emit-package, nvidia-cusparselt-cu12, + --no-emit-package, nvidia-nccl-cu12, + --no-emit-package, nvidia-nvjitlink-cu12, + --no-emit-package, nvidia-nvshmem-cu12, + --no-emit-package, nvidia-nvtx-cu12, + # nvidia cu13 packages + --no-emit-package, nvidia-cublas-cu13, + --no-emit-package, nvidia-cuda-cupti-cu13, + --no-emit-package, nvidia-cuda-nvrtc-cu13, + --no-emit-package, nvidia-cuda-runtime-cu13, + --no-emit-package, nvidia-cudnn-cu13, + --no-emit-package, nvidia-cufft-cu13, + --no-emit-package, nvidia-cufile-cu13, + --no-emit-package, nvidia-curand-cu13, + --no-emit-package, nvidia-cusolver-cu13, + --no-emit-package, nvidia-cusparse-cu13, --no-emit-package, nvidia-cusparselt-cu13, --no-emit-package, nvidia-nccl-cu13, - --no-emit-package, nvidia-nvjitlink, + --no-emit-package, nvidia-nvjitlink-cu13, --no-emit-package, nvidia-nvshmem-cu13, - --no-emit-package, nvidia-nvtx, + --no-emit-package, nvidia-nvtx-cu13, ] - files: ^requirements/rocm-test\.(in|txt)$ + files: ^requirements/(common|rocm|test/rocm)\.(in|txt)$ + - id: pip-compile + alias: pip-compile-xpu + name: pip-compile-xpu + args: [ + requirements/test/xpu.in, + -c, requirements/xpu.txt, + -o, requirements/test/xpu.txt, + --index-strategy, unsafe-best-match, + --torch-backend, xpu, + --python-platform, x86_64-manylinux_2_39, + --python-version, "3.12", + ] + files: ^requirements/(common|xpu|test/xpu)\.(in|txt)$ - repo: local hooks: - id: format-torch-nightly-test - name: reformat nightly_torch_test.txt to be in sync with test.in + name: reformat test/nightly-torch.txt to be in sync with test/cuda.in language: python entry: python tools/pre_commit/generate_nightly_torch_test.py - files: ^requirements/test\.(in|txt)$ + files: ^requirements/test/cuda\.(in|txt)$ - id: mypy-local name: Run mypy locally for lowest supported Python version entry: python tools/pre_commit/mypy.py 0 "3.10" diff --git a/AGENTS.md b/AGENTS.md index 61312b29ef7..215e4195eb4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -72,11 +72,11 @@ uv pip install -e . --torch-backend=auto ```bash # Install test dependencies. -# requirements/test.txt is pinned to x86_64; on other platforms, use the +# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the # unpinned source file instead: -uv pip install -r requirements/test.in # resolves for current platform +uv pip install -r requirements/test/cuda.in # resolves for current platform # Or on x86_64: -uv pip install -r requirements/test.txt +uv pip install -r requirements/test/cuda.txt # Run a specific test file (use .venv/bin/python directly; # `source activate` does not persist in non-interactive shells): diff --git a/CMakeLists.txt b/CMakeLists.txt index 2d3b76d2f04..f24c12eff83 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -307,6 +307,8 @@ set(VLLM_EXT_SRC "csrc/torch_bindings.cpp") if(VLLM_GPU_LANG STREQUAL "CUDA") + list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu") + SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. @@ -1222,6 +1224,7 @@ endif() # For CUDA we also build and ship some external projects. if (VLLM_GPU_LANG STREQUAL "CUDA") + include(cmake/external_projects/deepgemm.cmake) include(cmake/external_projects/flashmla.cmake) include(cmake/external_projects/qutlass.cmake) diff --git a/benchmarks/kernels/benchmark_moe_align_block_size.py b/benchmarks/kernels/benchmark_moe_align_block_size.py index 5f9a131f79b..a340500379a 100644 --- a/benchmarks/kernels/benchmark_moe_align_block_size.py +++ b/benchmarks/kernels/benchmark_moe_align_block_size.py @@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) from vllm.triton_utils import triton +from vllm.utils.torch_utils import set_random_seed def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor: @@ -44,7 +45,7 @@ configs = list( def benchmark(num_tokens, num_experts, topk, ep_size, provider): """Benchmark function for Triton.""" block_size = 256 - torch.cuda.manual_seed_all(0) + set_random_seed(0) topk_ids = get_topk_ids(num_tokens, num_experts, topk) e_map = None diff --git a/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py b/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py index 4384d3e5682..c9aaef284d7 100644 --- a/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py +++ b/benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py @@ -16,6 +16,7 @@ from vllm.utils.deep_gemm import ( fp8_gemm_nt, per_block_cast_to_fp8, ) +from vllm.utils.torch_utils import set_random_seed def benchmark_shape( @@ -235,9 +236,7 @@ def run_benchmarks(verbose: bool = False): torch.backends.cudnn.allow_tf32 = True # Set seeds for reproducibility - torch.manual_seed(42) - torch.cuda.manual_seed(42) - + set_random_seed(42) # Define benchmark shapes (m, n, k) shapes = [ (8, 4096, 7168), diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index e23f6b923f1..881039f43f0 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/benchmarks/multi_turn/benchmark_serving_multi_turn.py @@ -1439,6 +1439,12 @@ async def main() -> None: action="store_true", help="Export summary to Excel file (optional)", ) + parser.add_argument( + "--stats-json-output", + type=str, + default=None, + help="Export per-request stats (ttft_ms, tpot_ms, etc.) to a JSON file", + ) parser.add_argument( "-v", "--verbose", @@ -1651,6 +1657,19 @@ async def main() -> None: warmup_runtime_sec=warmup_runtime_sec, ) + if args.stats_json_output is not None: + # Export per-request metrics as a JSON array for downstream analysis. + stats_data = [s._asdict() for s in client_metrics] + logger.info( + f"{Color.GREEN}Writing per-request stats JSON: " + f"{args.stats_json_output}{Color.RESET}" + ) + os.makedirs( + os.path.dirname(os.path.abspath(args.stats_json_output)), exist_ok=True + ) + with open(args.stats_json_output, "w") as f: + json.dump(stats_data, f, indent=2) + if args.output_file is not None: # Write a JSON file with the updated conversations # The "assistant" content will contain the answers from the tested LLM diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 1b3f0d5ad83..0389ce1299a 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -349,6 +349,7 @@ endif() set(VLLM_EXT_SRC "csrc/cpu/activation.cpp" "csrc/cpu/utils.cpp" + "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" @@ -383,6 +384,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/cpu_wna16.cpp" "csrc/cpu/cpu_fused_moe.cpp" "csrc/cpu/utils.cpp" + "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/dnnl_kernels.cpp" "csrc/cpu/torch_bindings.cpp" @@ -395,6 +397,7 @@ if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_AVX2 "csrc/cpu/utils.cpp" + "csrc/cpu/spec_decode_utils.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake new file mode 100644 index 00000000000..c3a48a64fc7 --- /dev/null +++ b/cmake/external_projects/deepgemm.cmake @@ -0,0 +1,151 @@ +include(FetchContent) + +# If DEEPGEMM_SRC_DIR is set, DeepGEMM is built from that directory +# instead of downloading. +# It can be set as an environment variable or passed as a cmake argument. +# The environment variable takes precedence. +if (DEFINED ENV{DEEPGEMM_SRC_DIR}) + set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR}) +endif() + +if(DEEPGEMM_SRC_DIR) + FetchContent_Declare( + deepgemm + SOURCE_DIR ${DEEPGEMM_SRC_DIR} + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +else() + # This ref should be kept in sync with tools/install_deepgemm.sh + FetchContent_Declare( + deepgemm + GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git + GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03 + GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_PROGRESS TRUE + CONFIGURE_COMMAND "" + BUILD_COMMAND "" + ) +endif() + +# Use FetchContent_Populate (not MakeAvailable) to avoid processing +# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls. +FetchContent_GetProperties(deepgemm) +if(NOT deepgemm_POPULATED) + FetchContent_Populate(deepgemm) +endif() +message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}") + +# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 +set(DEEPGEMM_SUPPORT_ARCHS) +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a") +endif() +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") +elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") +endif() + +cuda_archs_loose_intersection(DEEPGEMM_ARCHS + "${DEEPGEMM_SUPPORT_ARCHS}" "${CUDA_ARCHS}") + +if(DEEPGEMM_ARCHS) + message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}") + + find_package(CUDAToolkit REQUIRED) + + # + # Build the _C pybind11 extension from DeepGEMM's C++ source. + # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime. + # + Python_add_library(_deep_gemm_C MODULE WITH_SOABI + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp") + + # The pybind11 module name must be _C to match DeepGEMM's Python imports. + set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C") + + target_compile_definitions(_deep_gemm_C PRIVATE + "-DTORCH_EXTENSION_NAME=_C") + + target_include_directories(_deep_gemm_C PRIVATE + "${deepgemm_SOURCE_DIR}/csrc" + "${deepgemm_SOURCE_DIR}/deep_gemm/include" + "${deepgemm_SOURCE_DIR}/third-party/cutlass/include" + "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include" + "${deepgemm_SOURCE_DIR}/third-party/fmt/include") + + target_compile_options(_deep_gemm_C PRIVATE + $<$:-std=c++17> + $<$:-O3> + $<$:-Wno-psabi> + $<$:-Wno-deprecated-declarations>) + + # torch_python is required because DeepGEMM uses pybind11 type casters + # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which + # use torch::Library custom ops. + find_library(TORCH_PYTHON_LIBRARY torch_python + PATHS "${TORCH_INSTALL_PREFIX}/lib" + REQUIRED) + + target_link_libraries(_deep_gemm_C PRIVATE + torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}" + CUDA::cudart CUDA::nvrtc) + + # Install the shared library into the vendored package directory + install(TARGETS _deep_gemm_C + LIBRARY DESTINATION vllm/third_party/deep_gemm + COMPONENT _deep_gemm_C) + + # + # Vendor DeepGEMM Python package files + # + install(FILES + "${deepgemm_SOURCE_DIR}/deep_gemm/__init__.py" + DESTINATION vllm/third_party/deep_gemm + COMPONENT _deep_gemm_C) + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/utils/" + DESTINATION vllm/third_party/deep_gemm/utils + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/testing/" + DESTINATION vllm/third_party/deep_gemm/testing + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/legacy/" + DESTINATION vllm/third_party/deep_gemm/legacy + COMPONENT _deep_gemm_C + FILES_MATCHING PATTERN "*.py") + + # Generate envs.py (normally generated by DeepGEMM's setup.py build step) + file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py" + "# Pre-installed environment variables\npersistent_envs = dict()\n") + install(FILES "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py" + DESTINATION vllm/third_party/deep_gemm + RENAME envs.py + COMPONENT _deep_gemm_C) + + # + # Install include files needed for JIT compilation at runtime. + # The JIT compiler finds these relative to the package directory. + # + + # DeepGEMM's own CUDA headers + install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/include/" + DESTINATION vllm/third_party/deep_gemm/include + COMPONENT _deep_gemm_C) + + # CUTLASS and CuTe headers (vendored for JIT, separate from vLLM's CUTLASS) + install(DIRECTORY "${deepgemm_SOURCE_DIR}/third-party/cutlass/include/" + DESTINATION vllm/third_party/deep_gemm/include + COMPONENT _deep_gemm_C) + +else() + message(STATUS "DeepGEMM will not compile: " + "unsupported CUDA architecture ${CUDA_ARCHS}") + # Create empty target so setup.py doesn't fail on unsupported systems + add_custom_target(_deep_gemm_C) +endif() diff --git a/csrc/async_util.cuh b/csrc/async_util.cuh new file mode 100644 index 00000000000..392d78c53fd --- /dev/null +++ b/csrc/async_util.cuh @@ -0,0 +1,100 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +namespace vllm { +namespace cuda_async { + +__device__ __forceinline__ void cp_async_shared_global_16_cg( + void* smem_ptr, const void* glob_ptr) { +#if defined(USE_ROCM) + *reinterpret_cast(smem_ptr) = *reinterpret_cast(glob_ptr); +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" + : + : "r"(smem), "l"(glob_ptr)); +#elif defined(__CUDA_ARCH__) + *reinterpret_cast(smem_ptr) = *reinterpret_cast(glob_ptr); +#else + (void)smem_ptr; + (void)glob_ptr; +#endif +} + +__device__ __forceinline__ void cp_async_shared_global_ca(void* smem_ptr, + const void* glob_ptr, + int size_bytes) { +#if defined(USE_ROCM) + if (size_bytes == 4) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else if (size_bytes == 8) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); + if (size_bytes == 4) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n" + : + : "r"(smem), "l"(glob_ptr)); + } else if (size_bytes == 8) { + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" + : + : "r"(smem), "l"(glob_ptr)); + } else { + asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" + : + : "r"(smem), "l"(glob_ptr)); + } +#elif defined(__CUDA_ARCH__) + if (size_bytes == 4) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else if (size_bytes == 8) { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } else { + *reinterpret_cast(smem_ptr) = + *reinterpret_cast(glob_ptr); + } +#else + (void)smem_ptr; + (void)glob_ptr; + (void)size_bytes; +#endif +} + +__device__ __forceinline__ void cp_async_commit_group() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM) + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +template +__device__ __forceinline__ void cp_async_wait_group() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM) + asm volatile("cp.async.wait_group %0;\n" : : "n"(n)); +#endif +} + +} // namespace cuda_async +} // namespace vllm diff --git a/csrc/attention/dtype_fp8.cuh b/csrc/attention/dtype_fp8.cuh index e714e321b0b..1afec3c3997 100644 --- a/csrc/attention/dtype_fp8.cuh +++ b/csrc/attention/dtype_fp8.cuh @@ -17,6 +17,22 @@ enum class Fp8KVCacheDataType { kFp8E5M2 = 2, }; +inline Fp8KVCacheDataType get_fp8_kv_cache_data_type( + const std::string& dtype_str) { + // dtype_str refers to CacheDType at vllm.config.cache.CacheDType + if (dtype_str == "auto" || dtype_str == "float16" || + dtype_str == "bfloat16") { + // unquantized kv cache + return Fp8KVCacheDataType::kAuto; + } else if (dtype_str == "fp8" || dtype_str == "fp8_ds_mla" || + dtype_str == "fp8_e4m3") { + return Fp8KVCacheDataType::kFp8E4M3; + } else if (dtype_str == "fp8_e5m2") { + return Fp8KVCacheDataType::kFp8E5M2; + } + TORCH_CHECK(false, "Unsupported fp8 kv cache data type: ", dtype_str); +} + // fp8 vector types for quantization of kv cache template <> struct Vec { diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu index a7e5fdabf64..d1cfbaeb05d 100644 --- a/csrc/cache_kernels.cu +++ b/csrc/cache_kernels.cu @@ -104,37 +104,49 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs, static_assert(sizeof(CUdeviceptr) == sizeof(int64_t)); static_assert(sizeof(size_t) == sizeof(int64_t)); #if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 - CUmemcpyAttributes attr = {}; - attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; - size_t attrs_idx = 0; - #if defined(CUDA_VERSION) && CUDA_VERSION >= 13000 - CUresult result = cuMemcpyBatchAsync( - reinterpret_cast(dst_data), - reinterpret_cast(src_data), - reinterpret_cast(size_data), static_cast(n), &attr, - &attrs_idx, 1, static_cast(stream)); - TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed with error ", - result); - #else - size_t fail_idx = 0; - CUresult result = cuMemcpyBatchAsync( - reinterpret_cast(dst_data), - reinterpret_cast(src_data), - reinterpret_cast(size_data), static_cast(n), &attr, - &attrs_idx, 1, &fail_idx, static_cast(stream)); - TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", - fail_idx, " with error ", result); - #endif -#else - // Fallback for CUDA < 12.8 and ROCm: individual async copies. - // cudaMemcpyDefault lets the driver infer direction from pointer types. - for (int64_t i = 0; i < n; i++) { - cudaMemcpyAsync(reinterpret_cast(dst_data[i]), - reinterpret_cast(src_data[i]), - static_cast(size_data[i]), cudaMemcpyDefault, - stream); - } + // Resolve cuMemcpyBatchAsync at runtime via cuGetProcAddress so that + // binaries compiled with CUDA 12.8+ still work on older drivers, and + // we avoid the CUDA 13.0 header remapping (#define to _v2 signature). + // The function pointer is cached after the first call. + using BatchFn = + CUresult (*)(CUdeviceptr*, CUdeviceptr*, size_t*, size_t, + CUmemcpyAttributes*, size_t*, size_t, size_t*, CUstream); + static BatchFn batch_fn = []() -> BatchFn { + CUdriverProcAddressQueryResult sym_status; + void* fn_ptr = nullptr; + CUresult res = cuGetProcAddress("cuMemcpyBatchAsync", &fn_ptr, 12080, + CU_GET_PROC_ADDRESS_DEFAULT, &sym_status); + if (res != CUDA_SUCCESS || fn_ptr == nullptr) { + return nullptr; + } + return reinterpret_cast(fn_ptr); + }(); + + if (batch_fn != nullptr) { + CUmemcpyAttributes attr = {}; + attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM; + size_t attrs_idx = 0; + size_t fail_idx = 0; + CUresult result = batch_fn(reinterpret_cast(dst_data), + reinterpret_cast(src_data), + reinterpret_cast(size_data), + static_cast(n), &attr, &attrs_idx, 1, + &fail_idx, static_cast(stream)); + TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ", + fail_idx, " with error ", result); + } else #endif + { + // Fallback for CUDA < 12.8, older drivers, and ROCm: + // individual async copies. + // cudaMemcpyDefault lets the driver infer direction from pointer types. + for (int64_t i = 0; i < n; i++) { + cudaMemcpyAsync(reinterpret_cast(dst_data[i]), + reinterpret_cast(src_data[i]), + static_cast(size_data[i]), cudaMemcpyDefault, + stream); + } + } } namespace vllm { diff --git a/csrc/cpu/spec_decode_utils.cpp b/csrc/cpu/spec_decode_utils.cpp new file mode 100644 index 00000000000..a76b8bc6937 --- /dev/null +++ b/csrc/cpu/spec_decode_utils.cpp @@ -0,0 +1,409 @@ +#include "cpu_types.hpp" + +#include + +namespace cpu_utils { + +void eagle_prepare_inputs_padded_kernel_impl( + const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& valid_sampled_tokens_count, + const torch::Tensor& query_start_loc_gpu, + torch::Tensor& token_indices_to_sample, + torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs) { + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* valid_count_ptr = + valid_sampled_tokens_count.data_ptr(); + const int32_t* query_loc_ptr = query_start_loc_gpu.data_ptr(); + int32_t* indices_out_ptr = token_indices_to_sample.data_ptr(); + int64_t* rejected_out_ptr = num_rejected_tokens_gpu.data_ptr(); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t num_draft_tokens = cu_draft_ptr[req_idx] - start_idx; + int64_t num_valid_tokens = valid_count_ptr[req_idx]; + + int64_t num_rejected = 0; + if (num_draft_tokens > 0) { + num_rejected = num_draft_tokens + 1 - num_valid_tokens; + } + + int32_t q_last_tok_idx = query_loc_ptr[req_idx + 1] - 1; + int32_t index_to_sample = q_last_tok_idx - num_rejected; + + indices_out_ptr[req_idx] = index_to_sample; + rejected_out_ptr[req_idx] = num_rejected; + } +} + +void eagle_prepare_next_token_padded_kernel_impl( + const torch::Tensor& sampled_token_ids, + const torch::Tensor& discard_request_mask, + const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids, + torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size, + const int64_t num_sampled_tokens_per_req, const int64_t num_reqs) { + const int64_t* sampled_ids_ptr = sampled_token_ids.data_ptr(); + const bool* discard_mask_ptr = discard_request_mask.data_ptr(); + const int64_t* backup_ids_ptr = backup_next_token_ids.data_ptr(); + int64_t* next_ids_out_ptr = next_token_ids.data_ptr(); + int64_t* valid_count_out_ptr = valid_sampled_tokens_count.data_ptr(); + + const int64_t stride = sampled_token_ids.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + const int64_t* row_ptr = sampled_ids_ptr + req_idx * stride; + int64_t valid_count = 0; + int64_t last_valid_token = -1; + + for (int64_t pos = 0; pos < num_sampled_tokens_per_req; ++pos) { + int64_t token = row_ptr[pos]; + if (token != -1 && token < vocab_size) { + valid_count++; + last_valid_token = token; + } + } + + bool discard = discard_mask_ptr[req_idx]; + if (discard) { + next_ids_out_ptr[req_idx] = backup_ids_ptr[req_idx]; + valid_count_out_ptr[req_idx] = 0; + } else { + next_ids_out_ptr[req_idx] = + (valid_count > 0) ? last_valid_token : backup_ids_ptr[req_idx]; + valid_count_out_ptr[req_idx] = valid_count; + } + } +} + +void eagle_step_slot_mapping_metadata_kernel_impl( + const torch::Tensor& positions, const torch::Tensor& block_table, + torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions, + torch::Tensor& out_slot_mapping, const int64_t block_size, + const int64_t max_model_len, const int64_t PAD_ID) { + const int64_t batch_size = positions.size(0); + const int64_t input_batch_size = out_slot_mapping.size(0); + + const int64_t* pos_ptr = positions.data_ptr(); + const int32_t* bt_ptr = block_table.data_ptr(); + int32_t* seq_lens_ptr = seq_lens.data_ptr(); + int64_t* out_clamped_ptr = out_clamped_positions.data_ptr(); + int64_t* out_slot_ptr = out_slot_mapping.data_ptr(); + + const int64_t bt_stride = block_table.stride(0); + const int64_t n_blocks_per_req = block_table.size(1); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < input_batch_size; ++req_idx) { + if (req_idx >= batch_size) { + out_slot_ptr[req_idx] = PAD_ID; + continue; + } + + int64_t position = pos_ptr[req_idx]; + int64_t new_position = position + 1; + bool exceeds_max = new_position >= max_model_len; + int64_t clamped_position = exceeds_max ? 0 : new_position; + + out_clamped_ptr[req_idx] = clamped_position; + + int64_t block_number = clamped_position / block_size; + block_number = std::min(block_number, n_blocks_per_req - 1); + int32_t block_id = bt_ptr[req_idx * bt_stride + block_number]; + int64_t slot_id = block_id * block_size + (clamped_position % block_size); + out_slot_ptr[req_idx] = exceeds_max ? PAD_ID : slot_id; + + int32_t seq_len = seq_lens_ptr[req_idx]; + int32_t new_seq_len = exceeds_max ? 1 : (seq_len + 1); + new_seq_len = std::min(new_seq_len, static_cast(max_model_len)); + seq_lens_ptr[req_idx] = new_seq_len; + } +} + +void copy_and_expand_eagle_inputs_kernel_impl( + const torch::Tensor& target_token_ids, + const torch::Tensor& target_positions, const torch::Tensor& next_token_ids, + torch::Tensor& out_input_ids, torch::Tensor& out_positions, + torch::Tensor& out_is_rejected_token_mask, + torch::Tensor& out_is_masked_token_mask, + torch::Tensor& out_new_token_indices, + torch::Tensor& out_hidden_state_mapping, + const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc, + const int64_t padding_token_id, const int64_t parallel_drafting_token_id, + const int64_t total_input_tokens, + const int64_t num_padding_slots_per_request, const bool shift_input_ids) { + const int64_t num_reqs = query_end_loc.size(0); + + const int64_t* target_ids_ptr = target_token_ids.data_ptr(); + const int64_t* target_pos_ptr = target_positions.data_ptr(); + const int64_t* next_ids_ptr = next_token_ids.data_ptr(); + const int32_t* query_start_ptr = query_start_loc.data_ptr(); + const int32_t* query_end_ptr = query_end_loc.data_ptr(); + + int64_t* out_ids_ptr = out_input_ids.data_ptr(); + int64_t* out_pos_ptr = out_positions.data_ptr(); + bool* out_rej_mask_ptr = out_is_rejected_token_mask.data_ptr(); + bool* out_mask_ptr = out_is_masked_token_mask.data_ptr(); + int32_t* out_new_idx_ptr = out_new_token_indices.data_ptr(); + int32_t* out_hidden_map_ptr = out_hidden_state_mapping.data_ptr(); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + int32_t q_start = query_start_ptr[req_idx]; + int32_t next_q_start = query_start_ptr[req_idx + 1]; + int32_t q_end = query_end_ptr[req_idx]; + + int64_t num_valid_tokens = + shift_input_ids ? (q_end - q_start) : (q_end - q_start + 1); + int64_t input_offset = shift_input_ids ? 1 : 0; + + int64_t out_start = q_start + req_idx * (num_padding_slots_per_request - + (shift_input_ids ? 1 : 0)); + int64_t num_rejected = next_q_start - q_end - 1; + int64_t total_output_tokens = + num_valid_tokens + num_padding_slots_per_request + num_rejected; + + int64_t start_pos = target_pos_ptr[q_start]; + int64_t bonus_token = next_ids_ptr[req_idx]; + + for (int64_t j = 0; j < total_output_tokens; ++j) { + int64_t out_idx = out_start + j; + bool is_valid = j < num_valid_tokens; + bool is_bonus = j == num_valid_tokens; + bool is_parallel = (j > num_valid_tokens) && + (j < num_valid_tokens + num_padding_slots_per_request); + bool is_rejected = j >= num_valid_tokens + num_padding_slots_per_request; + + int64_t in_idx = + std::min(static_cast(q_start + input_offset + j), + total_input_tokens - 1); + + int64_t token_id = padding_token_id; + if (is_valid) + token_id = target_ids_ptr[in_idx]; + else if (is_bonus) + token_id = bonus_token; + else if (is_parallel) + token_id = parallel_drafting_token_id; + + out_ids_ptr[out_idx] = token_id; + out_pos_ptr[out_idx] = is_rejected ? 0 : (start_pos + j); + out_rej_mask_ptr[out_idx] = is_rejected; + out_mask_ptr[out_idx] = is_parallel; + + if (is_bonus || is_parallel) { + int64_t new_token_local_idx = j - num_valid_tokens; + int64_t new_token_out_idx = + req_idx * num_padding_slots_per_request + new_token_local_idx; + out_new_idx_ptr[new_token_out_idx] = out_idx; + } + } + + if (shift_input_ids) { + int64_t n_input = next_q_start - q_start; + for (int64_t j = 0; j < n_input; ++j) { + out_hidden_map_ptr[q_start + j] = out_start + j; + } + } + } +} + +void rejection_greedy_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, + const torch::Tensor& bonus_token_ids, + const std::optional& is_greedy, const int64_t max_spec_len) { + const int64_t batch_size = cu_num_draft_tokens.size(0); + + int64_t* out_ptr = output_token_ids.data_ptr(); + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr(); + const int64_t* target_argmax_ptr = target_argmax.data_ptr(); + const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr(); + const bool* greedy_ptr = + is_greedy.has_value() ? is_greedy.value().data_ptr() : nullptr; + + const int64_t out_stride = output_token_ids.stride(0); + const int64_t bonus_stride = bonus_token_ids.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + if (greedy_ptr && !greedy_ptr[req_idx]) continue; + + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t end_idx = cu_draft_ptr[req_idx]; + int64_t num_draft_tokens = end_idx - start_idx; + + bool rejected = false; + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { + int64_t target_id = target_argmax_ptr[start_idx + pos]; + out_ptr[req_idx * out_stride + pos] = target_id; + + if (draft_ids_ptr[start_idx + pos] != target_id) { + rejected = true; + break; + } + } + + if (!rejected) { + out_ptr[req_idx * out_stride + num_draft_tokens] = + bonus_ids_ptr[req_idx * bonus_stride]; + } + } +} + +void rejection_random_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids, + const torch::Tensor& recovered_token_ids, + const torch::Tensor& uniform_probs, + const std::optional& is_greedy, const int64_t max_spec_len, + const int64_t vocab_size, const bool no_draft_probs) { + const int64_t batch_size = cu_num_draft_tokens.size(0); + + int64_t* out_ptr = output_token_ids.data_ptr(); + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr(); + const float* draft_probs_ptr = + no_draft_probs ? nullptr : draft_probs.value().data_ptr(); + const float* target_probs_ptr = target_probs.data_ptr(); + const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr(); + const int64_t* recovered_ids_ptr = recovered_token_ids.data_ptr(); + const float* uniform_probs_ptr = uniform_probs.data_ptr(); + const bool* greedy_ptr = + is_greedy.has_value() ? is_greedy.value().data_ptr() : nullptr; + + const int64_t out_stride = output_token_ids.stride(0); + const int64_t bonus_stride = bonus_token_ids.stride(0); + const int64_t target_stride = target_probs.stride(0); + const int64_t draft_probs_stride = + no_draft_probs ? 0 : draft_probs.value().stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + if (greedy_ptr && greedy_ptr[req_idx]) continue; + + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t end_idx = cu_draft_ptr[req_idx]; + int64_t num_draft_tokens = end_idx - start_idx; + + bool rejected = false; + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { + int64_t token_idx = start_idx + pos; + int64_t draft_id = draft_ids_ptr[token_idx]; + + float p = target_probs_ptr[token_idx * target_stride + draft_id]; + float q = + no_draft_probs + ? 1.0f + : draft_probs_ptr[token_idx * draft_probs_stride + draft_id]; + float uniform_p = uniform_probs_ptr[token_idx]; + + float ratio = (q > 0.0f) ? (p / q) : 0.0f; + + if (ratio >= uniform_p) { + out_ptr[req_idx * out_stride + pos] = draft_id; + } else { + out_ptr[req_idx * out_stride + pos] = recovered_ids_ptr[token_idx]; + rejected = true; + break; + } + } + + if (!rejected) { + out_ptr[req_idx * out_stride + num_draft_tokens] = + bonus_ids_ptr[req_idx * bonus_stride]; + } + } +} + +void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& cu_num_tokens, + const int64_t replace_from, const int64_t replace_to) { + const int64_t batch_size = cu_num_tokens.size(0); + const int64_t* cu_tokens_ptr = cu_num_tokens.data_ptr(); + + int64_t* out_ptr = output.data_ptr(); + const int64_t* in_ptr = input.data_ptr(); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + int64_t start_idx = req_idx == 0 ? 0 : cu_tokens_ptr[req_idx - 1]; + int64_t end_idx = cu_tokens_ptr[req_idx]; + int64_t val = in_ptr[req_idx]; + + if (val == replace_from) { + val = replace_to; + } + + for (int64_t i = start_idx; i < end_idx; ++i) { + out_ptr[i] = val; + } + } +} + +void sample_recovered_tokens_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& inv_q, + const int64_t vocab_size, const bool no_draft_probs) { + const int64_t batch_size = cu_num_draft_tokens.size(0); + + int64_t* out_ptr = output_token_ids.data_ptr(); + const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr(); + const int64_t* draft_ids_ptr = draft_token_ids.data_ptr(); + const float* draft_probs_ptr = + no_draft_probs ? nullptr : draft_probs.value().data_ptr(); + const float* target_probs_ptr = target_probs.data_ptr(); + const float* inv_q_ptr = inv_q.data_ptr(); + + const int64_t target_stride = target_probs.stride(0); + const int64_t draft_probs_stride = + no_draft_probs ? 0 : draft_probs.value().stride(0); + const int64_t inv_q_stride = inv_q.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { + int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; + int64_t end_idx = cu_draft_ptr[req_idx]; + int64_t num_draft_tokens = end_idx - start_idx; + + const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride; + + for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { + int64_t token_idx = start_idx + pos; + int64_t draft_id = draft_ids_ptr[token_idx]; + + const float* token_target_probs = + target_probs_ptr + token_idx * target_stride; + const float* token_draft_probs = + no_draft_probs ? nullptr + : (draft_probs_ptr + token_idx * draft_probs_stride); + + int64_t best_id = 0; + float best_val = -1.0f; + + for (int64_t v = 0; v < vocab_size; ++v) { + float prob = token_target_probs[v]; + if (no_draft_probs) { + if (v == draft_id) prob = 0.0f; + } else { + float diff = prob - token_draft_probs[v]; + prob = diff > 0.0f ? diff : 0.0f; + } + + float val = prob * req_inv_q[v]; + if (val > best_val) { + best_val = val; + best_id = v; + } + } + out_ptr[token_idx] = best_id; + } + } +} + +} // namespace cpu_utils diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index fbc9c65241b..fcf7064f606 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -138,6 +138,61 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, torch::Tensor slot_mapping, const int64_t block_size); +namespace cpu_utils { +void eagle_prepare_inputs_padded_kernel_impl( + const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& valid_sampled_tokens_count, + const torch::Tensor& query_start_loc_gpu, + torch::Tensor& token_indices_to_sample, + torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs); +void eagle_prepare_next_token_padded_kernel_impl( + const torch::Tensor& sampled_token_ids, + const torch::Tensor& discard_request_mask, + const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids, + torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size, + const int64_t num_sampled_tokens_per_req, const int64_t num_reqs); +void eagle_step_slot_mapping_metadata_kernel_impl( + const torch::Tensor& positions, const torch::Tensor& block_table, + torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions, + torch::Tensor& out_slot_mapping, const int64_t block_size, + const int64_t max_model_len, const int64_t PAD_ID); +void copy_and_expand_eagle_inputs_kernel_impl( + const torch::Tensor& target_token_ids, + const torch::Tensor& target_positions, const torch::Tensor& next_token_ids, + torch::Tensor& out_input_ids, torch::Tensor& out_positions, + torch::Tensor& out_is_rejected_token_mask, + torch::Tensor& out_is_masked_token_mask, + torch::Tensor& out_new_token_indices, + torch::Tensor& out_hidden_state_mapping, + const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc, + const int64_t padding_token_id, const int64_t parallel_drafting_token_id, + const int64_t total_input_tokens, + const int64_t num_padding_slots_per_request, const bool shift_input_ids); +void rejection_greedy_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, + const torch::Tensor& bonus_token_ids, + const std::optional& is_greedy, const int64_t max_spec_len); +void rejection_random_sample_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids, + const torch::Tensor& recovered_token_ids, + const torch::Tensor& uniform_probs, + const std::optional& is_greedy, const int64_t max_spec_len, + const int64_t vocab_size, const bool no_draft_probs); +void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input, + const torch::Tensor& cu_num_tokens, + const int64_t replace_from, const int64_t replace_to); +void sample_recovered_tokens_kernel_impl( + torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, + const torch::Tensor& draft_token_ids, + const std::optional& draft_probs, + const torch::Tensor& target_probs, const torch::Tensor& inv_q, + const int64_t vocab_size, const bool no_draft_probs); +} // namespace cpu_utils + TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // vLLM custom ops @@ -363,6 +418,70 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt " "block_size) -> ()", &compute_slot_mapping_kernel_impl); + + // Speculative decoding kernels + ops.def( + "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, " + "Tensor valid_sampled_tokens_count, Tensor query_start_loc_gpu, " + "Tensor(a3!) token_indices_to_sample, " + "Tensor(a4!) num_rejected_tokens_gpu, " + "SymInt num_reqs) -> ()", + &cpu_utils::eagle_prepare_inputs_padded_kernel_impl); + ops.def( + "eagle_prepare_next_token_padded_kernel_impl(" + "Tensor sampled_token_ids, Tensor discard_request_mask, " + "Tensor backup_next_token_ids, Tensor(a3!) next_token_ids, " + "Tensor(a4!) valid_sampled_tokens_count, SymInt vocab_size, " + "SymInt num_sampled_tokens_per_req, SymInt num_reqs) -> ()", + &cpu_utils::eagle_prepare_next_token_padded_kernel_impl); + ops.def( + "eagle_step_slot_mapping_metadata_kernel_impl(" + "Tensor positions, Tensor block_table, Tensor(a2!) seq_lens, " + "Tensor(a3!) out_clamped_positions, Tensor(a4!) out_slot_mapping, " + "SymInt block_size, SymInt max_model_len, SymInt PAD_ID) -> ()", + &cpu_utils::eagle_step_slot_mapping_metadata_kernel_impl); + ops.def( + "copy_and_expand_eagle_inputs_kernel_impl(" + "Tensor target_token_ids, Tensor target_positions, " + "Tensor next_token_ids, Tensor(a3!) out_input_ids, " + "Tensor(a4!) out_positions, " + "Tensor(a5!) out_is_rejected_token_mask, " + "Tensor(a6!) out_is_masked_token_mask, " + "Tensor(a7!) out_new_token_indices, " + "Tensor(a8!) out_hidden_state_mapping, " + "Tensor query_start_loc, Tensor query_end_loc, " + "SymInt padding_token_id, SymInt parallel_drafting_token_id, " + "SymInt total_input_tokens, SymInt num_padding_slots_per_request, " + "bool shift_input_ids) -> ()", + &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl); + ops.def( + "rejection_greedy_sample_kernel_impl(" + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " + "Tensor draft_token_ids, Tensor target_argmax, " + "Tensor bonus_token_ids, Tensor? is_greedy, " + "SymInt max_spec_len) -> ()", + &cpu_utils::rejection_greedy_sample_kernel_impl); + ops.def( + "rejection_random_sample_kernel_impl(" + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " + "Tensor draft_token_ids, Tensor? draft_probs, " + "Tensor target_probs, Tensor bonus_token_ids, " + "Tensor recovered_token_ids, Tensor uniform_probs, " + "Tensor? is_greedy, SymInt max_spec_len, SymInt vocab_size, " + "bool no_draft_probs) -> ()", + &cpu_utils::rejection_random_sample_kernel_impl); + ops.def( + "expand_kernel_impl(Tensor(a0!) output, Tensor input, " + "Tensor cu_num_tokens, SymInt replace_from, " + "SymInt replace_to) -> ()", + &cpu_utils::expand_kernel_impl); + ops.def( + "sample_recovered_tokens_kernel_impl(" + "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " + "Tensor draft_token_ids, Tensor? draft_probs, " + "Tensor target_probs, Tensor inv_q, SymInt vocab_size, " + "bool no_draft_probs) -> ()", + &cpu_utils::sample_recovered_tokens_kernel_impl); } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp index 8aa99b3e03a..0a3c9e9cc7f 100644 --- a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp @@ -389,20 +389,28 @@ struct Sm90ColOrScalarBroadcastArray { CUTLASS_DEVICE void begin() { - cute::Tensor pred = make_tensor(shape(tCgCol)); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = get<0>(tCcCol(i)) < m; - } - if (!params.col_broadcast) { fill(tCrCol, *(params.ptr_col_array[group])); return; } - // Filter so we don't issue redundant copies over stride-0 modes - // (only works if 0-strides are in same location, which is by construction) - copy_if(pred, filter(tCgCol), filter(tCrCol)); + // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and + // EPI_N are stride-0 for the column broadcast. Slice those modes at + // index 0 to avoid redundant copies AND ensure pred/data consistency + static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast"); + static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast"); + + auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + + cute::Tensor pred = make_tensor(shape(tCgCol_s)); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = get<0>(tCcCol_s(i)) < m; + } + + copy_if(pred, tCgCol_s, tCrCol_s); } template diff --git a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp index 8203d8930ce..29e6ec41e2a 100644 --- a/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp +++ b/csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp @@ -382,20 +382,28 @@ struct Sm90ColOrScalarBroadcast { CUTLASS_DEVICE void begin() { - cute::Tensor pred = make_tensor(shape(tCgCol)); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(pred); ++i) { - pred(i) = get<0>(tCcCol(i)) < m; - } - if (!params.col_broadcast) { fill(tCrCol, *(params.ptr_col)); return; } - // Filter so we don't issue redundant copies over stride-0 modes - // (only works if 0-strides are in same location, which is by construction) - copy_if(pred, filter(tCgCol), filter(tCrCol)); + // tCgCol has layout (CPY,CPY_M,CPY_N,EPI_M,EPI_N) where CPY_N and + // EPI_N are stride-0 for the column broadcast. Slice those modes at + // index 0 to avoid redundant copies AND ensure pred/data consistency + static_assert(decltype(stride<2>(tCgCol))::value == 0, "Expected stride-0 CPY_N for col broadcast"); + static_assert(decltype(stride<4>(tCgCol))::value == 0, "Expected stride-0 EPI_N for col broadcast"); + + auto tCgCol_s = tCgCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCrCol_s = tCrCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + auto tCcCol_s = tCcCol(_,_,0,_,0); // (CPY,CPY_M,EPI_M) + + cute::Tensor pred = make_tensor(shape(tCgCol_s)); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(pred); ++i) { + pred(i) = get<0>(tCcCol_s(i)) < m; + } + + copy_if(pred, tCgCol_s, tCrCol_s); } template diff --git a/csrc/fused_qknorm_rope_kernel.cu b/csrc/fused_qknorm_rope_kernel.cu index a51e1a347e1..0bf48fd3e83 100644 --- a/csrc/fused_qknorm_rope_kernel.cu +++ b/csrc/fused_qknorm_rope_kernel.cu @@ -19,8 +19,10 @@ #include #include +#include #include +#include "async_util.cuh" #include "cuda_compat.h" #include "dispatch_utils.h" #include "type_convert.cuh" @@ -86,6 +88,9 @@ inline __device__ __host__ T divUp(T m, T n) { } // namespace tensorrt_llm::common namespace tensorrt_llm::kernels { + +using namespace vllm::cuda_async; + // NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation, // with added support for passing the cos_sin_cache as an input. // https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -301,6 +306,237 @@ __global__ void fusedQKNormRopeKernel( #endif } +// Multi-token-head kernel: one warp processes HEADS_PER_WARP token-heads for +// the same token, sharing cos/sin from shared memory via cp.async. +// When HEADS_PER_WARP > 1 the warp reuses the loaded cos/sin across all heads, +// hiding global-memory latency and improving occupancy for large batches. +template +__global__ void fusedQKNormRopeKernelNTokenHeads( + void* qkv_void, int const num_heads_q, int const num_heads_k, + int const num_heads_v, float const eps, void const* q_weight_void, + void const* k_weight_void, void const* cos_sin_cache_void, + int64_t const* position_ids, int const num_tokens, int const rotary_dim) { +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + if constexpr ((std::is_same_v) || + std::is_same_v) { + return; + } else { +#endif + + using Converter = vllm::_typeConvert; + static_assert(Converter::exists, + "Input QKV data type is not supported for this CUDA " + "architecture or toolkit version."); + using T_in = typename Converter::hip_type; + using T2_in = typename Converter::packed_hip_type; + + using CacheConverter = vllm::_typeConvert; + static_assert(CacheConverter::exists, + "Cache data type is not supported for this CUDA architecture " + "or toolkit version."); + using T_cache = typename CacheConverter::hip_type; + + extern __shared__ char smem_storage[]; + // Shared memory layout: + // [0, cos_sin_bytes) : cos/sin for each warp (warpsPerBlock * + // rotary_dim * sizeof(T_cache)) + // [cos_sin_bytes, ...) : QKV tiles + // per warp (warpsPerBlock * HEADS_PER_WARP * 32 * elemSizeBytes) + T_cache* const smem = reinterpret_cast(smem_storage); + + T_in* qkv = reinterpret_cast(qkv_void); + T_in const* q_weight = reinterpret_cast(q_weight_void); + T_in const* k_weight = reinterpret_cast(k_weight_void); + T_cache const* cos_sin_cache = + reinterpret_cast(cos_sin_cache_void); + + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + + int const total_qk_heads = num_heads_q + num_heads_k; + int const num_heads = num_heads_q + num_heads_k + num_heads_v; + int const head_chunks_per_token = + (total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP; + + int const warp_global = blockIdx.x * warpsPerBlock + warpId; + int const tokenIdx = warp_global / head_chunks_per_token; + int const headChunk = warp_global % head_chunks_per_token; + int const first_head = headChunk * HEADS_PER_WARP; + int const num_heads_this_warp = + (first_head + HEADS_PER_WARP <= total_qk_heads) + ? HEADS_PER_WARP + : (total_qk_heads - first_head); + + if (tokenIdx >= num_tokens) return; + + static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64"); + constexpr int numElemsPerThread = head_dim / 32; + constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16); + static_assert(elemSizeBytes % 4 == 0, + "elemSizeBytes must be a multiple of 4"); + constexpr int vecSize = elemSizeBytes / 4; + using vec_T = typename tensorrt_llm::common::packed_as::type; + + int const cos_sin_bytes = + warpsPerBlock * rotary_dim * static_cast(sizeof(T_cache)); + int const qkv_tile_bytes = 32 * elemSizeBytes; + char* const this_warp_head_smem = + smem_storage + cos_sin_bytes + + warpId * (HEADS_PER_WARP * qkv_tile_bytes); + + // === Group 0: async load all heads' QKV into smem (issued first). === + for (int k = 0; k < num_heads_this_warp; ++k) { + int const localHeadIdx = first_head + k; + bool const isQ = localHeadIdx < num_heads_q; + int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; + int offWarp; + if (isQ) { + offWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; + } else { + offWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + + headIdx * head_dim; + } + int const offThread = offWarp + laneId * numElemsPerThread; + char* smem_dst = + this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes; + cp_async_shared_global_ca(smem_dst, + reinterpret_cast(&qkv[offThread]), + elemSizeBytes); + } + cp_async_commit_group(); // commit group 0 (QKV) + + // === Group 1: async load cos/sin into smem (issued second). === + int64_t const pos_id = position_ids[tokenIdx]; + T_cache const* const cache_ptr = cos_sin_cache + pos_id * rotary_dim; + int const copy_bytes = rotary_dim * static_cast(sizeof(T_cache)); + int const num_copies = (copy_bytes + 15) / 16; + for (int copyId = laneId; copyId < num_copies; copyId += 32) { + char* smem_ptr = + reinterpret_cast(&smem[warpId * rotary_dim]) + copyId * 16; + const char* glob_ptr = + reinterpret_cast(cache_ptr) + copyId * 16; + cp_async_shared_global_16_cg(smem_ptr, glob_ptr); + } + cp_async_commit_group(); // commit group 1 (cos/sin) + + // wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done. + cp_async_wait_group<1>(); + + float elements[numElemsPerThread]; + float elements2[numElemsPerThread]; + int const rotary_lanes = rotary_dim / numElemsPerThread; + int const embed_dim = rotary_dim / 2; + T_cache const* const cos_smem = &smem[warpId * rotary_dim]; + T_cache const* const sin_smem = &smem[warpId * rotary_dim + embed_dim]; + + // Preload weights into registers once, reused across all heads. + float q_w[numElemsPerThread]; + float k_w[numElemsPerThread]; +#pragma unroll + for (int i = 0; i < numElemsPerThread; i++) { + int const dim = laneId * numElemsPerThread + i; + q_w[i] = Converter::convert(q_weight[dim]); + k_w[i] = Converter::convert(k_weight[dim]); + } + + for (int k = 0; k < num_heads_this_warp; ++k) { + int const localHeadIdx = first_head + k; + bool const isQ = localHeadIdx < num_heads_q; + int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; + + int offsetWarp; + if (isQ) { + offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; + } else { + offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + + headIdx * head_dim; + } + int const offsetThread = offsetWarp + laneId * numElemsPerThread; + + // === Part 1: QK Norm (read from smem; group 0 already done). === + float sumOfSquares = 0.0f; + { + char const* smem_src = + this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes; + vec_T vec = *reinterpret_cast(smem_src); + constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); +#pragma unroll + for (int i = 0; i < num_packed_elems; i++) { + T2_in packed_val = *(reinterpret_cast(&vec) + i); + float2 vals = Converter::convert(packed_val); + sumOfSquares += vals.x * vals.x; + sumOfSquares += vals.y * vals.y; + elements[2 * i] = vals.x; + elements[2 * i + 1] = vals.y; + } + } + + sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares); + float rms_rcp = rsqrtf(sumOfSquares / static_cast(head_dim) + eps); + +#pragma unroll + for (int i = 0; i < numElemsPerThread; i++) { + elements[i] *= rms_rcp * (isQ ? q_w[i] : k_w[i]); + } + + // On first head: wait for group 1 (cos/sin) before RoPE. + if (k == 0) cp_async_wait_group<0>(); + + // === Part 2: RoPE using cos/sin from shared memory. === + if (laneId < rotary_lanes) { + if constexpr (interleave) { +#pragma unroll + for (int i = 0; i < numElemsPerThread / 2; ++i) { + int const idx0 = 2 * i; + int const idx1 = 2 * i + 1; + int const dim_idx = laneId * numElemsPerThread + idx0; + float const val0 = elements[idx0]; + float const val1 = elements[idx1]; + int const half_dim = dim_idx / 2; + float const cos_val = CacheConverter::convert(cos_smem[half_dim]); + float const sin_val = CacheConverter::convert(sin_smem[half_dim]); + elements[idx0] = val0 * cos_val - val1 * sin_val; + elements[idx1] = val0 * sin_val + val1 * cos_val; + } + } else { + __syncwarp(); + int const pairOffset = (rotary_dim / 2) / numElemsPerThread; +#pragma unroll + for (int i = 0; i < numElemsPerThread; i++) { + elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset); + if (laneId < pairOffset) elements2[i] = -elements2[i]; + int dim_idx = laneId * numElemsPerThread + i; + dim_idx = (dim_idx * 2) % rotary_dim; + int const half_dim = dim_idx / 2; + float const cos_val = CacheConverter::convert(cos_smem[half_dim]); + float const sin_val = CacheConverter::convert(sin_smem[half_dim]); + elements[i] = elements[i] * cos_val + elements2[i] * sin_val; + } + __syncwarp(); + } + } + + // Store. + { + vec_T vec; + constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); +#pragma unroll + for (int i = 0; i < num_packed_elems; i++) { + T2_in packed_val = Converter::convert( + make_float2(elements[2 * i], elements[2 * i + 1])); + *(reinterpret_cast(&vec) + i) = packed_val; + } + *reinterpret_cast(&qkv[offsetThread]) = vec; + } + } + +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) + } +#endif +} + // Borrowed from // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568 #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ @@ -321,15 +557,12 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, void const* cos_sin_cache, bool const interleave, int64_t const* position_ids, cudaStream_t stream) { constexpr int blockSize = 256; - int const warpsPerBlock = blockSize / 32; int const totalQKHeads = num_heads_q + num_heads_k; int const totalWarps = num_tokens * totalQKHeads; - int const gridSize = common::divUp(totalWarps, warpsPerBlock); dim3 gridDim(gridSize); dim3 blockDim(blockSize); - switch (head_dim) { case 64: DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { @@ -360,6 +593,118 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, "Unsupported head dimension for fusedQKNormRope: ", head_dim); } } + +// Launch: one warp processes token_heads_per_warp token-heads (1, 2, 4, or 8). +// When token_heads_per_warp == 1, delegates to the 1-head baseline above. +template +void launchFusedQKNormRopeNTokenHeads( + void* qkv, int const num_tokens, int const num_heads_q, + int const num_heads_k, int const num_heads_v, int const head_dim, + int const rotary_dim, float const eps, void const* q_weight, + void const* k_weight, void const* cos_sin_cache, bool const interleave, + int64_t const* position_ids, int const token_heads_per_warp, + cudaStream_t stream) { + TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 || + token_heads_per_warp == 4 || token_heads_per_warp == 8, + "token_heads_per_warp must be 1, 2, 4, or 8, got ", + token_heads_per_warp); + + // token_heads_per_warp == 1: delegate to the 1-head baseline kernel. + if (token_heads_per_warp == 1) { + launchFusedQKNormRope( + qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, + rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave, + position_ids, stream); + return; + } + + // NTokenHeads kernel uses cp.async to load cos/sin in 16-byte chunks. + // If rotary_dim * sizeof(cache_dtype) is not a multiple of 16, the last + // cp.async would write past the shared memory allocation. + // Fall back to the base kernel instead of failing. + { + size_t const rotary_bytes = + static_cast(rotary_dim) * + (std::is_same_v ? sizeof(float) : 2u); + if (rotary_bytes % 16 != 0) { + launchFusedQKNormRope( + qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, + rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave, + position_ids, stream); + return; + } + } + + constexpr int blockSize = 256; + int const warpsPerBlock = blockSize / 32; + int const totalQKHeads = num_heads_q + num_heads_k; + // Grid: one warp per (token, head_chunk); same token → reuse cos/sin in smem. + int const head_chunks_per_token = + (totalQKHeads + token_heads_per_warp - 1) / token_heads_per_warp; + int const total_warps = num_tokens * head_chunks_per_token; + int const gridSize = common::divUp(total_warps, warpsPerBlock); + dim3 gridDim(gridSize); + dim3 blockDim(blockSize); + // Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same + // layout). + size_t const cache_elem_size = + std::is_same_v ? sizeof(float) : 2u; + // QKV smem: token_heads_per_warp tiles per warp, each tile 32*(head_dim/32*2) + // = 2*head_dim bytes. + size_t const qkv_smem_per_warp = static_cast(token_heads_per_warp) * + 2u * static_cast(head_dim); + size_t const smem_bytes = + warpsPerBlock * static_cast(rotary_dim) * cache_elem_size + + warpsPerBlock * qkv_smem_per_warp; + +#define LAUNCH_N_TOKEN_HEADS(N) \ + do { \ + switch (head_dim) { \ + case 64: \ + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ + fusedQKNormRopeKernelNTokenHeads \ + <<>>( \ + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \ + k_weight, cos_sin_cache, position_ids, num_tokens, \ + rotary_dim); \ + }); \ + break; \ + case 128: \ + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ + fusedQKNormRopeKernelNTokenHeads \ + <<>>( \ + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \ + k_weight, cos_sin_cache, position_ids, num_tokens, \ + rotary_dim); \ + }); \ + break; \ + case 256: \ + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ + fusedQKNormRopeKernelNTokenHeads \ + <<>>( \ + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \ + k_weight, cos_sin_cache, position_ids, num_tokens, \ + rotary_dim); \ + }); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \ + } \ + } while (0) + + if (token_heads_per_warp == 2) { + LAUNCH_N_TOKEN_HEADS(2); + } else if (token_heads_per_warp == 4) { + LAUNCH_N_TOKEN_HEADS(4); + } else if (token_heads_per_warp == 8) { + LAUNCH_N_TOKEN_HEADS(8); + } +#undef LAUNCH_N_TOKEN_HEADS +} + } // namespace tensorrt_llm::kernels void fused_qk_norm_rope( @@ -374,7 +719,8 @@ void fused_qk_norm_rope( torch::Tensor& k_weight, // RMSNorm weights for key [head_dim] torch::Tensor& cos_sin_cache, // Cos/sin cache [max_position, head_dim] bool is_neox, // Whether RoPE is applied in Neox style - torch::Tensor& position_ids // Position IDs for RoPE [num_tokens] + torch::Tensor& position_ids, // Position IDs for RoPE [num_tokens] + int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value ) { // Input validation CHECK_INPUT(qkv); @@ -414,15 +760,48 @@ void fused_qk_norm_rope( qkv.size(1) == total_heads * head_dim, "QKV tensor size must match total number of heads and head dimension"); - auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + auto device_id = qkv.get_device(); + auto stream = at::cuda::getCurrentCUDAStream(device_id); + + // Select token_heads_per_warp: forced value if >0, else auto-select. + // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures, + // fall back to token_heads_per_warp=1 (base kernel) until profiled. + int token_heads_per_warp; + if (forced_token_heads_per_warp > 0) { // only support SM80+ + token_heads_per_warp = static_cast(forced_token_heads_per_warp); + } else { + token_heads_per_warp = 1; + auto* dev_prop = at::cuda::getDeviceProperties(device_id); + int sm_version = dev_prop->major * 10 + dev_prop->minor; + int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k); + if (sm_version == 90) { + if (head_dim >= 256) { + if (total_qk_units < 4096LL) { + token_heads_per_warp = 1; + } else if (total_qk_units < 8192LL) { + token_heads_per_warp = 2; + } else { + token_heads_per_warp = 4; + } + } else { + if (total_qk_units < 10240LL) { + token_heads_per_warp = 1; + } else if (total_qk_units < 40960LL) { + token_heads_per_warp = 4; + } else { + token_heads_per_warp = 8; + } + } + } + } VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] { using qkv_scalar_t = scalar_t; VLLM_DISPATCH_FLOATING_TYPES( cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] { using cache_scalar_t = scalar_t; - tensorrt_llm::kernels::launchFusedQKNormRope( + tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads< + qkv_scalar_t, cache_scalar_t>( qkv.data_ptr(), static_cast(num_tokens), static_cast(num_heads_q), static_cast(num_heads_k), static_cast(num_heads_v), static_cast(head_dim), @@ -430,7 +809,7 @@ void fused_qk_norm_rope( q_weight.data_ptr(), k_weight.data_ptr(), cos_sin_cache.data_ptr(), !is_neox, reinterpret_cast(position_ids.data_ptr()), - stream); + token_heads_per_warp, stream); }); }); -} \ No newline at end of file +} diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index 69b6564be75..8b0356815bb 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -240,8 +240,9 @@ template __global__ void per_token_group_quant_8bit_packed_kernel( const T* __restrict__ input, void* __restrict__ output_q, unsigned int* __restrict__ output_s_packed, const int group_size, - const int num_groups, const int groups_per_block, const int groups_per_row, - const int mn, const int tma_aligned_mn, const float eps, + const int num_groups_padded, const int groups_per_block, + const int padded_groups_per_row, const int groups_per_row, const int mn, + const int tma_aligned_mn, const int num_scale_elems, const float eps, const float min_8bit, const float max_8bit) { const int threads_per_group = 16; const int64_t local_group_id = threadIdx.x / threads_per_group; @@ -249,51 +250,62 @@ __global__ void per_token_group_quant_8bit_packed_kernel( const int64_t block_group_id = blockIdx.x * groups_per_block; const int64_t global_group_id = block_group_id + local_group_id; - if (global_group_id >= num_groups) { + if (global_group_id >= num_groups_padded) { return; } - const int64_t block_group_offset = global_group_id * group_size; + // map flat group id to 2D indices (mn_idx, sf_k_idx) + const int sf_k_idx = + static_cast(global_group_id % padded_groups_per_row); + const int mn_idx = static_cast(global_group_id / padded_groups_per_row); - const T* group_input = input + block_group_offset; - DST_DTYPE* group_output = - static_cast(output_q) + block_group_offset; + // whether it is a valid group (not padding) + const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row); // shared memory to cache each group's data to avoid double DRAM reads. extern __shared__ __align__(16) char smem_raw[]; T* smem = reinterpret_cast(smem_raw); T* smem_group = smem + local_group_id * group_size; - const float y_s = - ComputeGroupScale(group_input, smem_group, group_size, lane_id, - threads_per_group, eps, max_8bit); - // pack 4 scales into a uint32 + // compute scale for valid groups + float y_s = 0.f; + if (is_valid_group) { + const T* group_input = + input + static_cast(mn_idx) * groups_per_row * group_size + + sf_k_idx * group_size; + y_s = ComputeGroupScale(group_input, smem_group, group_size, + lane_id, threads_per_group, eps, max_8bit); + } + + // pack 4 scales into a uint32 exponent if (lane_id == 0) { - // map flat group id to 2D indices (mn_idx, sf_k_idx) - const int sf_k_idx = static_cast(global_group_id % groups_per_row); - const int mn_idx = static_cast(global_group_id / groups_per_row); - - if (mn_idx < mn) { - // each uint32 in output_s_packed stores 4 packed scales - const int sf_k_pack_idx = sf_k_idx / 4; - const int pos = sf_k_idx % 4; + // each uint32 in output_s_packed stores 4 packed scales + const int sf_k_pack_idx = sf_k_idx / 4; + const int pos = sf_k_idx % 4; + const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; + if (is_valid_group) { // reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit // exponent, and place it into the correct byte of the 32-bit word. const unsigned int bits = __float_as_uint(y_s); - const unsigned int exponent = (bits >> 23u) & 0xffu; - const unsigned int contrib = exponent << (pos * 8u); - - const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; - // atomically OR 8-bit exponent into the packed scales buffer - atomicOr(output_s_packed + out_idx, contrib); + const uint8_t exponent = static_cast((bits >> 23u) & 0xffu); + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = exponent; + } else if (out_idx < num_scale_elems) { + // write zero for padding groups if within bounds of output_s_packed + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = 0; } } __syncthreads(); - QuantizeGroup(smem_group, group_output, group_size, lane_id, - threads_per_group, y_s, min_8bit, max_8bit); + if (is_valid_group) { + DST_DTYPE* group_output = + static_cast(output_q) + + static_cast(mn_idx) * groups_per_row * group_size + + sf_k_idx * group_size; + QuantizeGroup(smem_group, group_output, group_size, lane_id, + threads_per_group, y_s, min_8bit, max_8bit); + } } void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, @@ -310,7 +322,6 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, const int64_t mn = input.numel() / k; const int64_t groups_per_row = k / group_size; - const int64_t num_groups = mn * groups_per_row; STD_TORCH_CHECK(output_s_packed.dim() == 2, "output_s_packed must be 2D, got dim=", output_s_packed.dim(), @@ -330,36 +341,46 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk, "], but got [", output_s_packed.size(0), ", ", output_s_packed.size(1), "]."); + // Verify column-major TMA-aligned layout + STD_TORCH_CHECK(output_s_packed.stride(0) == 1 && + output_s_packed.stride(1) == tma_aligned_mn, + "output_s_packed must have strides [1, ", tma_aligned_mn, + "], but got [", output_s_packed.stride(0), ", ", + output_s_packed.stride(1), "]."); cudaStream_t stream = get_current_cuda_stream(); constexpr int THREADS_PER_GROUP = 16; - const int groups_per_block = GetGroupsPerBlock(num_groups); + // Expand the grid to cover MN and K padding so every byte in + // output_s_packed is written (padding bytes get zeroed by the kernel). + const int64_t padded_groups_per_row = k_num_packed_sfk * 4; + const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row; + // Number of elements in output_s_packed. + const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn; + + const int groups_per_block = GetGroupsPerBlock(num_groups_padded); auto dst_type = output_q.scalar_type(); - const int num_blocks = num_groups / groups_per_block; + const int num_blocks = num_groups_padded / groups_per_block; const int num_threads = groups_per_block * THREADS_PER_GROUP; - // zero-initialize packed scales, since we use atomicOr to accumulate - // exponents from different groups. - torch::stable::zero_(output_s_packed); - -#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \ - do { \ - dim3 grid(num_blocks); \ - dim3 block(num_threads); \ - size_t smem_bytes = \ - static_cast(groups_per_block) * group_size * sizeof(T); \ - per_token_group_quant_8bit_packed_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - reinterpret_cast(output_s_packed.data_ptr()), \ - static_cast(group_size), static_cast(num_groups), \ - groups_per_block, static_cast(groups_per_row), \ - static_cast(mn), static_cast(tma_aligned_mn), \ - static_cast(eps), static_cast(min_8bit), \ - static_cast(max_8bit)); \ +#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \ + do { \ + dim3 grid(num_blocks); \ + dim3 block(num_threads); \ + size_t smem_bytes = \ + static_cast(groups_per_block) * group_size * sizeof(T); \ + per_token_group_quant_8bit_packed_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + reinterpret_cast(output_s_packed.data_ptr()), \ + static_cast(group_size), static_cast(num_groups_padded), \ + groups_per_block, static_cast(padded_groups_per_row), \ + static_cast(groups_per_row), static_cast(mn), \ + static_cast(tma_aligned_mn), \ + static_cast(num_scale_elems), static_cast(eps), \ + static_cast(min_8bit), static_cast(max_8bit)); \ } while (0) VLLM_STABLE_DISPATCH_FLOATING_TYPES( diff --git a/csrc/minimax_reduce_rms_kernel.cu b/csrc/minimax_reduce_rms_kernel.cu new file mode 100644 index 00000000000..6245b02d6e9 --- /dev/null +++ b/csrc/minimax_reduce_rms_kernel.cu @@ -0,0 +1,879 @@ + +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include +#include +#include + +#include "cuda_compat.h" +#include "cuda_utils.h" +#include "core/registration.h" +#include "minimax_reduce_rms_kernel.h" + +#include + +#define FINAL_MASK 0xffffffff +#define MINIMAX_REDUCE_RMS_WARP_SIZE 32 + +namespace vllm { +namespace tensorrt_llm { + +template +struct LamportComm { + __device__ __forceinline__ LamportComm(void** workspace, int rank) { + counter_ptr = &reinterpret_cast(workspace[NRanks * 3])[0]; + flag_ptr = &reinterpret_cast(workspace[NRanks * 3])[2]; + clear_ptr = &reinterpret_cast(workspace[NRanks * 3 + 1])[0]; + flag_value = *flag_ptr; + auto comm_size = reinterpret_cast(workspace[NRanks * 3 + 1])[1]; + clear_size = *clear_ptr; + int data_offset = flag_value % 3; + int clear_offset = (flag_value + 2) % 3; + for (int r = 0; r < NRanks; ++r) { + data_bufs[r] = reinterpret_cast(workspace[2 * NRanks + r]) + + data_offset * comm_size; + } + clear_buf = reinterpret_cast(workspace[2 * NRanks + rank]) + + clear_offset * comm_size; + __syncthreads(); + if (threadIdx.x == 0) { + atomicAdd(counter_ptr, 1); + } + } + + __device__ __forceinline__ void update(int64_t new_clear_size) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + while (*reinterpret_cast(counter_ptr) != gridDim.x) { + } + *flag_ptr = (flag_value + 1) % 3; + *clear_ptr = new_clear_size; + *counter_ptr = 0; + } + } + + int* counter_ptr; + int* flag_ptr; + int64_t* clear_ptr; + uint8_t* data_bufs[NRanks]; + uint8_t* clear_buf; + int64_t clear_size; + int flag_value; +}; + +__device__ __forceinline__ bool is_neg_zero(float v) { + return *reinterpret_cast(&v) == 0x80000000; +} + +__device__ __forceinline__ bool is_neg_zero(float4 v) { + return is_neg_zero(v.x) || is_neg_zero(v.y) || is_neg_zero(v.z) || + is_neg_zero(v.w); +} + +__device__ __forceinline__ float4 get_neg_zero() { + float4 vec; +#pragma unroll + for (int i = 0; i < 4; ++i) { + reinterpret_cast(&vec)[i] = 0x80000000; + } + return vec; +} + +template +__device__ __forceinline__ float rms_rsqrt(float& v, float eps) { + constexpr float kInvDim = 1.0F / static_cast(Dim); + v = rsqrtf((v * kInvDim) + eps); + return v; +} + +template +__device__ __forceinline__ float4 rms_rsqrt(float4& v, float eps) { + constexpr float kInvDim = 1.0F / static_cast(Dim); + v.x = rsqrtf((v.x * kInvDim) + eps); + v.y = rsqrtf((v.y * kInvDim) + eps); + v.z = rsqrtf((v.z * kInvDim) + eps); + v.w = rsqrtf((v.w * kInvDim) + eps); + return v; +} +__device__ __forceinline__ float4 ld_global_volatile(float4* addr) { + float4 val; + asm volatile("ld.volatile.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w) + : "l"(addr)); + return val; +} + +__device__ __forceinline__ float ld_global_volatile(float* addr) { + float val; + asm volatile("ld.volatile.global.f32 %0, [%1];" : "=f"(val) : "l"(addr)); + return val; +} + +// Used by the scalar (non-float4) kernel only +template +__inline__ __device__ T warpReduceSumV2(T* val) { +#pragma unroll + for (int i = 0; i < NUM; i++) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32); + } + return (T)(0.0f); +} + +template +__inline__ __device__ T blockReduceSumV2(T* val) { + static __shared__ T shared[NUM][33]; + int lane = threadIdx.x & 0x1f; + int wid = threadIdx.x >> 5; + + warpReduceSumV2(val); + + if (lane == 0) { +#pragma unroll + for (int i = 0; i < NUM; i++) { + shared[i][wid] = val[i]; + } + } + + __syncthreads(); + + bool is_mask = threadIdx.x < (blockDim.x / 32.f); +#pragma unroll + for (int i = 0; i < NUM; i++) { + val[i] = is_mask ? shared[i][lane] : (T)(0.0f); + } + warpReduceSumV2(val); + return (T)0.0f; +} + +// for float4 version +template +__device__ __forceinline__ void local_warp_reduce_sum_array( + T* value_ptr, uint32_t active_mask = 0xffffffffu) { + static_assert(kNumThreads >= 1 && + kNumThreads <= MINIMAX_REDUCE_RMS_WARP_SIZE); +#pragma unroll + for (int i = 0; i < ArraySize; ++i) { +#pragma unroll + for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { + value_ptr[i] += __shfl_xor_sync(active_mask, value_ptr[i], mask, + MINIMAX_REDUCE_RMS_WARP_SIZE); + } + } +} + +constexpr int next_pow2(int val) { + int result = 1; + while (result < val) { + result <<= 1; + } + return result; +} + +// --------------------------------------------------------------------------- + +template +class IndexHelper { + public: + __device__ __forceinline__ IndexHelper(MiniMaxReduceRMSParams const& params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + token_id = grid.cluster_rank(); + access_id_in_token = cluster.thread_rank(); + token_stride = grid.num_clusters(); +#else + token_id = blockIdx.x; + access_id_in_token = threadIdx.x; + token_stride = gridDim.x; +#endif + access_id = token_id * params.hidden_dim / kElemsPerAccess + + access_id_in_token; + access_stride = token_stride * params.hidden_dim / kElemsPerAccess; + tot_access = params.size_q / kElemsPerAccess; + } + + int token_id; + int access_id_in_token; + int token_stride; + int access_id; + int access_stride; + int tot_access; +}; + +/** +* this kernel is used to for minimax attention module +* input tensor [total_tokens, hidden_dim / tp_size], fp32 +* rms weight [hidden_dim / tp_size], bf16 +step 1: reduce from single rank to get the variance sum (reduce(input^2, +dim=-1)) step 2: reduce from all ranks to get the variance sum +(all_reduce(variance_sum)) step 3: calculate the rms norm (input * +rsqrt(variance + eps)) in this case, max hidden_dim is 6144 (float data), for +each token, we only need 6144 / 4 / tp_size = (1536 / tp_size) threads so we can +assume cluster size is 1 (tp_size >= 2) + */ +template +__global__ void __launch_bounds__(1024) + minimax_reduce_rms_kernel_lamport(MiniMaxReduceRMSParams params) { + IndexHelper index_helper(params); + int token_id = index_helper.token_id; + int access_id_in_token = index_helper.access_id_in_token; + int token_stride = index_helper.token_stride; + int access_id = index_helper.access_id; + int access_stride = index_helper.access_stride; + int tot_access = index_helper.tot_access; + int tot_tokens = params.size_q / params.hidden_dim; + float4 clear_vec = get_neg_zero(); + + LamportComm comm(params.workspace, params.rank); + int clear_access = comm.clear_size / kElemsPerAccess; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + for (int idx = access_id; idx < tot_access; + idx += access_stride, token_id += token_stride) { + alignas(16) DType vals[kElemsPerAccess]; + float sum_variance = 0.F; + *reinterpret_cast(vals) = + reinterpret_cast(params.allreduce_in)[idx]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + sum_variance += static_cast(vals[i]) * static_cast(vals[i]); + } + blockReduceSumV2(&sum_variance); + if (is_neg_zero(sum_variance)) { + sum_variance = 0.F; + } + if (threadIdx.x == 0) { + for (int r = 0; r < NRanks; ++r) { + reinterpret_cast( + comm.data_bufs[r])[(params.rank * tot_tokens) + token_id] = + (sum_variance); + } + } + + bool done = false; + float vars_all_ranks[NRanks]; + while (!done) { + done = true; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + vars_all_ranks[r] = ld_global_volatile(&reinterpret_cast( + comm.data_bufs[params.rank])[(r * tot_tokens) + token_id]); + done &= !is_neg_zero(vars_all_ranks[r]); + } + } + sum_variance = 0.F; +#pragma unroll + for (int r = 0; r < NRanks; ++r) { + sum_variance += vars_all_ranks[r]; + } + + DType norm_weight[kElemsPerAccess]; + *reinterpret_cast::vec_type*>(norm_weight) = + reinterpret_cast::vec_type*>( + params.rms_gamma)[access_id_in_token]; + +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[i] = static_cast( + static_cast(vals[i]) * + rsqrtf( + (sum_variance / static_cast(params.hidden_dim) / NRanks) + + params.rms_eps) * + static_cast(norm_weight[i])); + } + + reinterpret_cast(params.rms_norm_out)[idx] = + *reinterpret_cast(vals); + } + for (int idx = access_id; idx < clear_access; idx += access_stride) { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + comm.update(params.size_q * NRanks); +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif +} + +/** + * Float4 variant: process 4 rows at once, allreduce variance sums as float4 for + * better memory coalescing. sum_variance is always float; applies to all DTypes + * (half, bf16, float). When tot_tokens % 4 != 0, the last group pads rows with + * zeros; padded rows are not written to rms_norm_out. IsQK: when true, process + * Q+K in one loop with doubled comm buffer; when false, single-matrix (Q only). + */ +template +__global__ void __launch_bounds__(1024) + minimax_reduce_qk_rms_kernel_lamport_float4(MiniMaxReduceRMSParams params) { + // Compile-time per-rank dimensions + constexpr int RankQDim = OriginQDim / NRanks; + constexpr int RankKDim = OriginKDim / NRanks; + // Threads needed to cover one row of Q / K with float4 accesses + constexpr int ThreadsPerRowQ = RankQDim / kElemsPerAccess; + constexpr int ThreadsPerRowK = RankKDim / kElemsPerAccess; + // Number of warps dedicated to Q / K + constexpr int NumWarpQ = (ThreadsPerRowQ + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) / + MINIMAX_REDUCE_RMS_WARP_SIZE; + constexpr int NumWarpK = (ThreadsPerRowK + MINIMAX_REDUCE_RMS_WARP_SIZE - 1) / + MINIMAX_REDUCE_RMS_WARP_SIZE; + + int tot_tokens = params.size_q / RankQDim; + int tot_groups = (tot_tokens + 3) / 4; // ceiling; last group may be partial + + // Memory strides for strided qkv tensors (elements -> float4-access units) + int access_stride_q = (params.stride_q > 0 ? params.stride_q : RankQDim) / + kElemsPerAccess; + int access_stride_k = (params.stride_k > 0 ? params.stride_k : RankKDim) / + kElemsPerAccess; + // Output strides: default to contiguous (hidden_dim / hidden_dim_k) + int access_stride_q_out = + (params.stride_q_out > 0 ? params.stride_q_out : params.hidden_dim) / + kElemsPerAccess; + int access_stride_k_out = + (params.stride_k_out > 0 ? params.stride_k_out : params.hidden_dim_k) / + kElemsPerAccess; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + namespace cg = cooperative_groups; + cg::cluster_group cluster = cg::this_cluster(); + cg::grid_group grid = cg::this_grid(); + int group_id = grid.cluster_rank(); + int access_id_in_token = cluster.thread_rank(); + int group_stride = grid.num_clusters(); +#else + int group_id = blockIdx.x; + int access_id_in_token = threadIdx.x; + int group_stride = gridDim.x; +#endif + + bool is_q = (access_id_in_token < NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE); + int k_thread_idx = + access_id_in_token - (NumWarpQ * MINIMAX_REDUCE_RMS_WARP_SIZE); + bool is_valid_q = (access_id_in_token < ThreadsPerRowQ); + bool is_valid_k = (k_thread_idx >= 0 && k_thread_idx < ThreadsPerRowK); + float4 clear_vec = get_neg_zero(); + + // Shared memory for two-level block reduction and scale broadcast + __shared__ float block_reduce_sum[4][MINIMAX_REDUCE_RMS_WARP_SIZE + 1]; + __shared__ float global_scale_q[4]; + __shared__ float global_scale_k[4]; + + LamportComm comm(params.workspace, params.rank); + + DType norm_weight[kElemsPerAccess]{}; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.wait;"); +#endif + if (is_q) { + if (is_valid_q) { + *reinterpret_cast::vec_type*>( + norm_weight) = + reinterpret_cast::vec_type const*>( + params.rms_gamma)[access_id_in_token]; + } + } else { + if (is_valid_k) { + *reinterpret_cast::vec_type*>( + norm_weight) = + reinterpret_cast::vec_type const*>( + params.rms_gamma_k)[k_thread_idx]; + } + } + + // Main loop: process one group of 4 tokens per iteration. + for (int g = group_id; g < tot_groups; g += group_stride) { + alignas(16) DType vals[4][kElemsPerAccess]{}; + float warp_sum_variance[4]{0.F, 0.F, 0.F, 0.F}; + + if (is_q) { +#pragma unroll + for (int row = 0; row < 4; ++row) { + int token_r = g * 4 + row; + if (token_r >= tot_tokens || !is_valid_q) { + continue; + } + int idx_r = token_r * access_stride_q + access_id_in_token; + *reinterpret_cast(&vals[row][0]) = + reinterpret_cast(params.allreduce_in)[idx_r]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + float x = static_cast(vals[row][i]); + warp_sum_variance[row] += x * x; + } + } + } else { +#pragma unroll + for (int row = 0; row < 4; ++row) { + int token_r = g * 4 + row; + if (token_r >= tot_tokens || !is_valid_k) { + continue; + } + int idx_r = token_r * access_stride_k + k_thread_idx; + *reinterpret_cast(&vals[row][0]) = + reinterpret_cast(params.allreduce_in_k)[idx_r]; +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + float x = static_cast(vals[row][i]); + warp_sum_variance[row] += x * x; + } + } + } + + local_warp_reduce_sum_array( + warp_sum_variance); + // Warp lane 0 writes its warp's partial sum to shared memory + int lane = threadIdx.x & (MINIMAX_REDUCE_RMS_WARP_SIZE - 1); + if (lane == 0) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + block_reduce_sum[t][threadIdx.x / MINIMAX_REDUCE_RMS_WARP_SIZE] = + warp_sum_variance[t]; + } + } + __syncthreads(); + + int tid = threadIdx.x; + + if (tid < MINIMAX_REDUCE_RMS_WARP_SIZE) { + constexpr int kNumWarpQPow2 = + (next_pow2(NumWarpQ) > NRanks) ? next_pow2(NumWarpQ) : NRanks; + float local_sum[4]; +#pragma unroll + for (int t = 0; t < 4; ++t) { + local_sum[t] = (tid < NumWarpQ) ? block_reduce_sum[t][tid] : 0.F; + } + // After this, all kNumWarpQPow2 lanes (including tid 0..NRanks-1) have + // the total Q sum-of-squares for all 4 tokens. + local_warp_reduce_sum_array(local_sum); + + if (tid < NRanks) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + if (is_neg_zero(local_sum[t])) { + local_sum[t] = 0.F; + } + } + // Parallel push: thread tid writes this rank's Q sum to rank tid's buf + reinterpret_cast( + comm.data_bufs[tid])[(params.rank * tot_groups * 2) + (2 * g)] = + *reinterpret_cast(local_sum); + + // Parallel pull: thread tid reads rank tid's contribution from + // this rank's (params.rank's) buffer + bool done = false; + float4 var_all_ranks; + while (!done) { + done = true; + var_all_ranks = ld_global_volatile(&reinterpret_cast( + comm.data_bufs[params.rank])[(tid * tot_groups * 2) + (2 * g)]); + done &= !is_neg_zero(var_all_ranks); + } + + // Warp-level allreduce: each of the NRanks threads holds one rank's + // partial sum; after this all NRanks threads have the global total. + constexpr uint32_t kQActiveMask = (1u << NRanks) - 1u; + local_warp_reduce_sum_array( + reinterpret_cast(&var_all_ranks), kQActiveMask); + + // Thread 0 computes rsqrt with compile-time Dim and writes to smem + if (tid == 0) { + *reinterpret_cast(global_scale_q) = + rms_rsqrt(var_all_ranks, params.rms_eps); + } + } + } else if (tid >= MINIMAX_REDUCE_RMS_WARP_SIZE * NumWarpQ && + tid < MINIMAX_REDUCE_RMS_WARP_SIZE * (NumWarpQ + 1)) { + // --- K leader warp --- + constexpr int kNumWarpKPow2 = + (next_pow2(NumWarpK) > NRanks) ? next_pow2(NumWarpK) : NRanks; + float local_sum[4]; +#pragma unroll + for (int t = 0; t < 4; ++t) { + local_sum[t] = (k_thread_idx < NumWarpK) + ? block_reduce_sum[t][NumWarpQ + k_thread_idx] + : 0.F; + } + local_warp_reduce_sum_array(local_sum); + + if (k_thread_idx < NRanks) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + if (is_neg_zero(local_sum[t])) { + local_sum[t] = 0.F; + } + } + reinterpret_cast( + comm.data_bufs[k_thread_idx])[(params.rank * tot_groups * 2) + + (2 * g + 1)] = + *reinterpret_cast(local_sum); + + bool done = false; + float4 var_all_ranks; + while (!done) { + done = true; + var_all_ranks = ld_global_volatile(&reinterpret_cast( + comm.data_bufs[params.rank])[(k_thread_idx * tot_groups * 2) + + (2 * g + 1)]); + done &= !is_neg_zero(var_all_ranks); + } + + constexpr uint32_t kKActiveMask = (1u << NRanks) - 1u; + local_warp_reduce_sum_array( + reinterpret_cast(&var_all_ranks), kKActiveMask); + + if (k_thread_idx == 0) { + *reinterpret_cast(global_scale_k) = + rms_rsqrt(var_all_ranks, params.rms_eps); + } + } + } + __syncthreads(); + + if (is_q) { +#pragma unroll + for (int t = 0; t < 4; ++t) { + warp_sum_variance[t] = global_scale_q[t]; + } +#pragma unroll + for (int r = 0; r < 4; ++r) { +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[r][i] = static_cast(static_cast(vals[r][i]) * + warp_sum_variance[r] * + static_cast(norm_weight[i])); + } + int token_r = g * 4 + r; + if (token_r >= tot_tokens || !is_valid_q) { + continue; + } + int idx_out = token_r * access_stride_q_out + access_id_in_token; + reinterpret_cast(params.rms_norm_out)[idx_out] = + *reinterpret_cast(&vals[r][0]); + } + } else { +#pragma unroll + for (int t = 0; t < 4; ++t) { + warp_sum_variance[t] = global_scale_k[t]; + } +#pragma unroll + for (int r = 0; r < 4; ++r) { +#pragma unroll + for (int i = 0; i < kElemsPerAccess; ++i) { + vals[r][i] = static_cast(static_cast(vals[r][i]) * + warp_sum_variance[r] * + static_cast(norm_weight[i])); + } + int token_r = g * 4 + r; + if (token_r >= tot_tokens || !is_valid_k) { + continue; + } + int idx_out = token_r * access_stride_k_out + k_thread_idx; + reinterpret_cast(params.rms_norm_out_k)[idx_out] = + *reinterpret_cast(&vals[r][0]); + } + } + } // end group loop +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + + int clear_access = static_cast(comm.clear_size / kElemsPerAccess); + int clear_stride = group_stride * blockDim.x; + for (int idx = group_id * blockDim.x + threadIdx.x; idx < clear_access; + idx += clear_stride) { + reinterpret_cast(comm.clear_buf)[idx] = clear_vec; + } + + comm.update(static_cast(2) * tot_groups * kElemsPerAccess * + NRanks); +} + +int get_sm_count() { + static int sm_count = 0; + if (sm_count == 0) { + int device_id; + CUDA_CHECK(cudaGetDevice(&device_id)); + cudaDeviceProp device_prop; + cudaGetDeviceProperties(&device_prop, device_id); + sm_count = device_prop.multiProcessorCount; + } + return sm_count; +} + +inline int getSMVersion(bool queryRealSmArch = false) { + int device{-1}; + CUDA_CHECK(cudaGetDevice(&device)); + int sm_major = 0; + int sm_minor = 0; + CUDA_CHECK(cudaDeviceGetAttribute(&sm_major, + cudaDevAttrComputeCapabilityMajor, device)); + CUDA_CHECK(cudaDeviceGetAttribute(&sm_minor, + cudaDevAttrComputeCapabilityMinor, device)); + int sm = sm_major * 10 + sm_minor; + if (sm == 121 && !queryRealSmArch) { + return 120; + } + return sm; +} + +template +int get_max_active_blocks(KernelFunc kernel, int block_size, + int dynamic_smem = 0) { + int max_active = 0; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active, kernel, block_size, dynamic_smem)); + return std::max(max_active, 1); +} + +template +void minimax_reduce_rms_kernel_launcher(MiniMaxReduceRMSParams const& params) { + static int SM = getSMVersion(); + int token_num = params.size_q / params.hidden_dim; + int sm_count = get_sm_count(); + int cluster_size = 1; + int cluster_num = token_num; + int threads_per_token = params.hidden_dim / kElemsPerAccess; + int block_size = threads_per_token; + + int max_blocks_per_sm = get_max_active_blocks( + minimax_reduce_rms_kernel_lamport, block_size); + int max_grid = max_blocks_per_sm * sm_count; + + int grid_size = + (std::min(max_grid, cluster_num * cluster_size) / cluster_size) * + cluster_size; + + cudaLaunchConfig_t cfg; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + + cudaLaunchAttribute attribute[2]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = SM >= 90 ? 2 : 0; + + CUDA_CHECK(cudaLaunchKernelEx( + &cfg, minimax_reduce_rms_kernel_lamport, params)); +} + +template +void minimax_reduce_rms_kernel_launcher_float4( + MiniMaxReduceRMSParams const& params) { + TORCH_CHECK(params.size_q % params.hidden_dim == 0); + TORCH_CHECK(params.hidden_dim % kElemsPerAccess == 0); + if (params.stride_q > 0) { + TORCH_CHECK(params.stride_q % kElemsPerAccess == 0); + } + TORCH_CHECK(params.allreduce_in_k != nullptr, + "float4 QK kernel requires K input"); + TORCH_CHECK(params.hidden_dim >= params.hidden_dim_k); + TORCH_CHECK(params.size_k % params.hidden_dim_k == 0); + TORCH_CHECK(params.hidden_dim_k % kElemsPerAccess == 0); + TORCH_CHECK(params.size_q / params.hidden_dim == + params.size_k / params.hidden_dim_k); + if (params.stride_k > 0) { + TORCH_CHECK(params.stride_k % kElemsPerAccess == 0); + } + + int token_num = params.size_q / params.hidden_dim; + int tot_groups = (token_num + 3) / 4; + if (tot_groups == 0) { + return; + } + + static int SM = getSMVersion(); + int sm_count = get_sm_count(); + int cluster_size = 1; + int cluster_num = tot_groups; + + int access_per_row_q = params.hidden_dim / kElemsPerAccess; + int access_per_row_k = params.hidden_dim_k / kElemsPerAccess; + + // Round each section up to a warp boundary + auto divUp = [](int a, int b) { return (a + b - 1) / b * b; }; + int block_size = divUp(access_per_row_q, MINIMAX_REDUCE_RMS_WARP_SIZE) + + divUp(access_per_row_k, MINIMAX_REDUCE_RMS_WARP_SIZE); + + auto kfn = + minimax_reduce_qk_rms_kernel_lamport_float4; + + int max_blocks_per_sm = get_max_active_blocks(kfn, block_size); + int max_grid = max_blocks_per_sm * sm_count; + int grid_size = + (std::min(max_grid, cluster_num * cluster_size) / cluster_size) * + cluster_size; + + cudaLaunchConfig_t cfg; + cfg.gridDim = grid_size; + cfg.blockDim = block_size; + cfg.dynamicSmemBytes = 0; + cfg.stream = params.stream; + + cudaLaunchAttribute attribute[2]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = 1; + attribute[1].id = cudaLaunchAttributeClusterDimension; + attribute[1].val.clusterDim.x = cluster_size; + attribute[1].val.clusterDim.y = 1; + attribute[1].val.clusterDim.z = 1; + cfg.attrs = attribute; + cfg.numAttrs = SM >= 90 ? 2 : 0; + + CUDA_CHECK(cudaLaunchKernelEx(&cfg, kfn, params)); +} + +template +void dispatch_dtype(MiniMaxReduceRMSParams const& params) { + // Use the optimized QK float4 kernel when: + // - K input is present, AND + // - the full (NRanks * per-rank) dimensions match the MiniMax M2 shape. + // Otherwise fall back to the scalar kernel. + bool use_float4 = (params.allreduce_in_k != nullptr) && + (params.hidden_dim * params.nranks == 6144) && + (params.hidden_dim_k * params.nranks == 1024); + + if (params.dtype == at::ScalarType::Half) { + if (use_float4) { + minimax_reduce_rms_kernel_launcher_float4( + params); + } else { + minimax_reduce_rms_kernel_launcher(params); + } + } else if (params.dtype == at::ScalarType::BFloat16) { + if (use_float4) { + minimax_reduce_rms_kernel_launcher_float4<__nv_bfloat16, NRanks, 6144, + 1024>(params); + } else { + minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params); + } + } else if (params.dtype == at::ScalarType::Float) { + if (use_float4) { + minimax_reduce_rms_kernel_launcher_float4( + params); + } else { + minimax_reduce_rms_kernel_launcher(params); + } + } else { + TORCH_CHECK(false, "Unsupported data type for minimax_reduce_rms_op"); + } +} + +void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params) { + if (params.nranks == 2) { + dispatch_dtype<2>(params); + } else if (params.nranks == 4) { + dispatch_dtype<4>(params); + } else if (params.nranks == 8) { + dispatch_dtype<8>(params); + } else if (params.nranks == 16) { + dispatch_dtype<16>(params); + } else { + TORCH_CHECK(false, "minimax_reduce_rms_op: unsupported ranks number!"); + } +} +} // namespace tensorrt_llm +} // namespace vllm + +torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, + torch::Tensor const& norm_weight, + torch::Tensor workspace, int64_t const rank, + int64_t const nranks, double const eps) { + auto allreduce_params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); + + allreduce_params.nranks = static_cast(nranks); + allreduce_params.rank = static_cast(rank); + allreduce_params.dtype = input.scalar_type(); + allreduce_params.size_q = static_cast(input.numel()); + allreduce_params.hidden_dim = static_cast(input.size(-1)); + allreduce_params.stride_q = allreduce_params.hidden_dim; + allreduce_params.workspace = + reinterpret_cast(workspace.mutable_data_ptr()); + allreduce_params.allreduce_in = input.data_ptr(); + allreduce_params.rms_gamma = norm_weight.data_ptr(); + allreduce_params.rms_eps = static_cast(eps); + allreduce_params.stream = at::cuda::getCurrentCUDAStream(input.get_device()); + + torch::Tensor rms_norm_out = torch::empty_like(input); + allreduce_params.rms_norm_out = rms_norm_out.mutable_data_ptr(); + + vllm::tensorrt_llm::minimax_reduce_rms_op(allreduce_params); + + return rms_norm_out; +} + +std::tuple minimax_allreduce_rms_qk( + torch::Tensor qkv, torch::Tensor const& norm_weight_q, + torch::Tensor const& norm_weight_k, torch::Tensor workspace, + int64_t const q_size, int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps) { + TORCH_CHECK(qkv.dim() == 2, "minimax_allreduce_rms_qk: qkv must be 2D"); + TORCH_CHECK(qkv.is_contiguous(), + "minimax_allreduce_rms_qk: qkv must be contiguous"); + int64_t qkv_dim = qkv.size(-1); + TORCH_CHECK(qkv_dim == q_size + 2 * kv_size, + "minimax_allreduce_rms_qk: qkv last dim must equal " + "q_size + 2 * kv_size"); + TORCH_CHECK(rank < nranks, + "minimax_allreduce_rms_qk: rank must be less than nranks"); + + int64_t num_tokens = qkv.size(0); + int elem_bytes = qkv.element_size(); + + torch::Tensor q_out = torch::empty({num_tokens, q_size}, qkv.options()); + torch::Tensor k_out = torch::empty({num_tokens, kv_size}, qkv.options()); + + auto params = vllm::tensorrt_llm::MiniMaxReduceRMSParams(); + params.nranks = static_cast(nranks); + params.rank = static_cast(rank); + params.dtype = qkv.scalar_type(); + params.size_q = static_cast(num_tokens * q_size); + params.hidden_dim = static_cast(q_size); + params.size_k = static_cast(num_tokens * kv_size); + params.hidden_dim_k = static_cast(kv_size); + params.stride_q = static_cast(qkv_dim); + params.stride_k = static_cast(qkv_dim); + params.stride_q_out = 0; // q_out is contiguous; kernel uses hidden_dim + params.stride_k_out = 0; // k_out is contiguous; kernel uses hidden_dim_k + params.workspace = reinterpret_cast(workspace.mutable_data_ptr()); + + uint8_t* base = static_cast(qkv.data_ptr()); + params.allreduce_in = base; + params.allreduce_in_k = base + q_size * elem_bytes; + params.rms_gamma = norm_weight_q.data_ptr(); + params.rms_gamma_k = norm_weight_k.data_ptr(); + params.rms_eps = static_cast(eps); + params.stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + + params.rms_norm_out = q_out.mutable_data_ptr(); + params.rms_norm_out_k = k_out.mutable_data_ptr(); + + vllm::tensorrt_llm::minimax_reduce_rms_op(params); + return {q_out, k_out}; +} diff --git a/csrc/minimax_reduce_rms_kernel.h b/csrc/minimax_reduce_rms_kernel.h new file mode 100644 index 00000000000..e8c2d012247 --- /dev/null +++ b/csrc/minimax_reduce_rms_kernel.h @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include + +namespace vllm { +namespace tensorrt_llm { + +template +struct ElemsPerAccess; + +template <> +struct ElemsPerAccess { + static constexpr int value = 8; + using vec_type = float4; +}; + +template <> +struct ElemsPerAccess { + static constexpr int value = 8; + using vec_type = float4; +}; + +template <> +struct ElemsPerAccess { + static constexpr int value = 4; + using vec_type = float4; +}; + +template +static constexpr int kElemsPerAccess = ElemsPerAccess::value; + +struct MiniMaxReduceRMSParams { + int nranks{}; + int rank{}; + at::ScalarType dtype{at::ScalarType::Undefined}; + int size_q{}; + int hidden_dim{}; + int size_k{}; + int hidden_dim_k{}; + int stride_q{}; // row stride for q input (elements); when > hidden_dim, + // q is part of a wider qkv tensor + int stride_k{}; // row stride for k input (elements); when > hidden_dim_k, + // k is part of a wider qkv tensor + int stride_q_out{}; // row stride for q output (elements); 0 = contiguous + int stride_k_out{}; // row stride for k output (elements); 0 = contiguous + void** workspace{}; + void* allreduce_in{}; + void* rms_norm_out{}; + void* rms_gamma{}; + void* allreduce_in_k{}; + void* rms_norm_out_k{}; + void* rms_gamma_k{}; + float rms_eps{}; + cudaStream_t stream{}; +}; + +void minimax_reduce_rms_op(MiniMaxReduceRMSParams const& params); + +} // namespace tensorrt_llm +} // namespace vllm diff --git a/csrc/ops.h b/csrc/ops.h index b5e9ff36a59..0a0b6c2d7d0 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -96,7 +96,8 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q, int64_t num_heads_k, int64_t num_heads_v, int64_t head_dim, double eps, torch::Tensor& q_weight, torch::Tensor& k_weight, torch::Tensor& cos_sin_cache, - bool is_neox, torch::Tensor& position_ids); + bool is_neox, torch::Tensor& position_ids, + int64_t forced_token_heads_per_warp); void apply_repetition_penalties_(torch::Tensor& logits, const torch::Tensor& prompt_mask, @@ -308,4 +309,16 @@ int64_t qr_max_size(); #ifndef USE_ROCM void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b); -#endif \ No newline at end of file +#endif + +#ifndef USE_ROCM +torch::Tensor minimax_allreduce_rms(torch::Tensor const& input, + torch::Tensor const& norm_weight, + torch::Tensor workspace, int64_t const rank, + int64_t const nranks, double const eps); +std::tuple minimax_allreduce_rms_qk( + torch::Tensor qkv, torch::Tensor const& norm_weight_q, + torch::Tensor const& norm_weight_k, torch::Tensor workspace, + int64_t const q_size, int64_t const kv_size, int64_t const rank, + int64_t const nranks, double const eps); +#endif diff --git a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh index 81f5cb83f3e..7ae644d81d4 100644 --- a/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -639,7 +639,9 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { // function with template. #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ - if (KV_DTYPE == "auto") { \ + vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ + vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ + if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ if (SRC_DTYPE == at::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ } else if (SRC_DTYPE == at::ScalarType::Half) { \ @@ -649,21 +651,18 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { } else { \ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ - } else { \ - if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ + } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh index 421e8092474..3b7e25dc56b 100644 --- a/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh +++ b/csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -543,7 +543,9 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { // function with template. #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ - if (KV_DTYPE == "auto") { \ + vllm::Fp8KVCacheDataType KV_CACHE_DTYPE = \ + vllm::get_fp8_kv_cache_data_type(KV_DTYPE); \ + if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { \ if (SRC_DTYPE == at::ScalarType::Float) { \ FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ } else if (SRC_DTYPE == at::ScalarType::Half) { \ @@ -553,43 +555,28 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { } else { \ TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ - } else { \ - if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ - } else if (KV_DTYPE == "fp8_e5m2") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ - } else if (KV_DTYPE == "fp8_ds_mla") { \ - if (SRC_DTYPE == at::ScalarType::Float) { \ - FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::Half) { \ - FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ - FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ - } else { \ - TORCH_CHECK(false, \ - "Unsupported input type of kv cache: ", SRC_DTYPE); \ - } \ + } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ } else { \ - TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ } \ + } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else { \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ } } // namespace fp8 diff --git a/csrc/sampler.cu b/csrc/sampler.cu index 2e76873c8f1..c0cc03a08ad 100644 --- a/csrc/sampler.cu +++ b/csrc/sampler.cu @@ -564,8 +564,9 @@ template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( const float* logits, const int* seqLens, int* outIndices, int stride0, - int stride1, const int topK, int next_n, float* outLogits = nullptr, - const int numBlocksToMerge = 0, const int* indices = nullptr) { + int stride1, const int topK, int next_n, int seqLensIs2D = 0, + float* outLogits = nullptr, const int numBlocksToMerge = 0, + const int* indices = nullptr) { // The number of bins in the histogram. static constexpr int kNumBins = 2048; @@ -574,8 +575,16 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( // The range of logits within the row. int rowStart = 0; - int seq_len = seqLens[rowIdx / next_n]; - int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1); + int batch_idx = rowIdx / next_n; + int next_n_idx = rowIdx % next_n; + // seqLensIs2D=0: 1D seqLens — all rows in a batch share the same seq_len; + // kernel computes per-row effective length via offset. + // seqLensIs2D=1: 2D seqLens — each logit row has its own pre-computed + // effective length (flat index rowIdx = b*next_n + j maps + // directly to seqLens[b, j] in C-contiguous layout). + int seq_len = seqLensIs2D ? seqLens[rowIdx] : seqLens[batch_idx]; + int rowEnd = + seqLensIs2D ? max(0, seq_len) : max(0, seq_len - next_n + next_n_idx + 1); // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) { @@ -653,6 +662,11 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); const auto numColumns = logits.size(1); + // True if seqLens is 2D (B, next_n): each logit row has its own pre-computed + // effective seq_len. False if seqLens is 1D (B,): all rows in a batch share + // the same seq_len and the kernel computes the per-row offset itself. + int seqLensIs2D = seqLens.dim() == 2 ? 1 : 0; + if (numColumns < kSortingAlgorithmThreshold) { // Use insertion sort vllm::topKPerRowDecode @@ -660,7 +674,7 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, logits.data_ptr(), seqLens.data_ptr(), indices.data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), - static_cast(next_n)); + static_cast(next_n), seqLensIs2D); } else if (numColumns < kSplitWorkThreshold) { // From this threshold, use radix sort instead vllm::topKPerRowDecode @@ -668,7 +682,7 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, logits.data_ptr(), seqLens.data_ptr(), indices.data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), - static_cast(next_n)); + static_cast(next_n), seqLensIs2D); } else { // Long sequences are run in two steps constexpr auto multipleBlocksPerRowConfig = 10; @@ -686,15 +700,16 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, logits.data_ptr(), seqLens.data_ptr(), outIndicesAux.data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), - static_cast(next_n), outLogitsAux.data_ptr()); + static_cast(next_n), seqLensIs2D, + outLogitsAux.data_ptr()); constexpr int kNumThreadsPerBlockMerge = 1024; vllm::topKPerRowDecode <<>>( outLogitsAux.data_ptr(), seqLens.data_ptr(), indices.data_ptr(), multipleBlocksPerRowConfig * topK, 1, - static_cast(topK), static_cast(next_n), nullptr, - multipleBlocksPerRowConfig, outIndicesAux.data_ptr()); + static_cast(topK), static_cast(next_n), seqLensIs2D, + nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr()); } } diff --git a/csrc/topk.cu b/csrc/topk.cu index 402b64b027a..f48e7cbc4fc 100644 --- a/csrc/topk.cu +++ b/csrc/topk.cu @@ -21,13 +21,15 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths, TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32"); TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32"); TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); - TORCH_CHECK(lengths.dim() == 1, "lengths must be 1D"); + TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, + "lengths must be 1D or 2D"); + TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); TORCH_CHECK(output.dim() == 2, "output must be 2D"); const int64_t num_rows = logits.size(0); const int64_t stride = logits.size(1); - TORCH_CHECK(lengths.size(0) == num_rows, "lengths size mismatch"); + TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, "output size mismatch"); namespace P = vllm::persistent; diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 78156cfced9..48062c3f47b 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -173,7 +173,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " "int num_heads_k, int num_heads_v, int head_dim, float eps, " "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, " - "bool is_neox, Tensor position_ids) -> ()"); + "bool is_neox, Tensor position_ids, " + "int forced_token_heads_per_warp=-1) -> ()"); ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope); // Apply repetition penalties to logits in-place @@ -496,6 +497,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? b_qzeros, " "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt " "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor"); + + ops.def( + "minimax_allreduce_rms(" + "Tensor input," + "Tensor norm_weight," + "Tensor workspace," + "int rank," + "int nranks," + "float eps) -> Tensor"); + ops.impl("minimax_allreduce_rms", torch::kCUDA, &minimax_allreduce_rms); + ops.def( + "minimax_allreduce_rms_qk(" + "Tensor qkv," + "Tensor norm_weight_q," + "Tensor norm_weight_k," + "Tensor workspace," + "int q_size," + "int kv_size," + "int rank," + "int nranks," + "float eps) -> (Tensor, Tensor)"); + ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk); + // conditionally compiled so impl in source file #endif } diff --git a/docker/Dockerfile b/docker/Dockerfile index 6bbd34f9543..12942b5c807 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -204,7 +204,7 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL ARG PYTORCH_NIGHTLY # Install build dependencies -COPY requirements/build.txt requirements/build.txt +COPY requirements/build/cuda.txt requirements/build/cuda.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -219,13 +219,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing build requirements without torch..." \ && python3 use_existing_torch.py --prefix \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ && echo "Installing torch nightly..." \ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ echo "Installing build requirements..." \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi @@ -315,7 +315,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \ #################### CSRC BUILD IMAGE #################### #################### EXTENSIONS BUILD IMAGE #################### -# Build DeepGEMM, DeepEP - runs in PARALLEL with csrc-build +# Build DeepEP - runs in PARALLEL with csrc-build # This stage is independent and doesn't affect csrc cache FROM base AS extensions-build ARG CUDA_VERSION @@ -327,21 +327,6 @@ ENV UV_LINK_MODE=copy WORKDIR /workspace -# Build DeepGEMM wheel -# Default moved here from tools/install_deepgemm.sh for centralized version management -ARG DEEPGEMM_GIT_REF=477618cd51baffca09c4b0b87e97c03fe827ef03 -COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh -RUN --mount=type=cache,target=/root/.cache/uv \ - mkdir -p /tmp/deepgemm/dist && \ - VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh \ - --cuda-version "${CUDA_VERSION}" \ - ${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} \ - --wheel-dir /tmp/deepgemm/dist || \ - echo "DeepGEMM build skipped (CUDA version requirement not met)" - -# Ensure the wheel dir exists so COPY won't fail when DeepGEMM is skipped -RUN mkdir -p /tmp/deepgemm/dist && touch /tmp/deepgemm/dist/.deepgemm_skipped - # Build DeepEP wheels COPY tools/ep_kernels/install_python_libraries.sh /tmp/install_python_libraries.sh # Defaults moved here from tools/ep_kernels/install_python_libraries.sh for centralized version management @@ -370,7 +355,7 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL ARG PYTORCH_NIGHTLY # Install build dependencies -COPY requirements/build.txt requirements/build.txt +COPY requirements/build/cuda.txt requirements/build/cuda.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -385,13 +370,13 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing build requirements without torch..." \ && python3 use_existing_torch.py --prefix \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ && echo "Installing torch nightly..." \ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | grep -i "^torch=" | xargs) --pre \ --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ else \ echo "Installing build requirements..." \ - && uv pip install --python /opt/venv/bin/python3 -r requirements/build.txt \ + && uv pip install --python /opt/venv/bin/python3 -r requirements/build/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ fi @@ -426,7 +411,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 # Copy extension wheels from extensions-build stage for later use -COPY --from=extensions-build /tmp/deepgemm/dist /tmp/deepgemm/dist COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist # Check the size of the wheel if RUN_WHEEL_CHECK is true @@ -466,8 +450,8 @@ ARG PYTORCH_NIGHTLY # Install development dependencies COPY requirements/lint.txt requirements/lint.txt -COPY requirements/test.in requirements/test.in -COPY requirements/test.txt requirements/test.txt +COPY requirements/test/cuda.in requirements/test/cuda.in +COPY requirements/test/cuda.txt requirements/test/cuda.txt COPY requirements/dev.txt requirements/dev.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -475,8 +459,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing dev requirements plus torch nightly..." \ && python3 use_existing_torch.py --prefix \ - && cat torch_lib_versions.txt >> requirements/test.in \ - && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \ + && cat torch_lib_versions.txt >> requirements/test/cuda.in \ + && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ && uv pip install --python /opt/venv/bin/python3 $(cat torch_lib_versions.txt | xargs) --pre \ -r requirements/dev.txt \ @@ -554,7 +538,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \ cuda-nvrtc-${CUDA_VERSION_DASH} \ cuda-cuobjdump-${CUDA_VERSION_DASH} \ libcurand-dev-${CUDA_VERSION_DASH} \ - libcublas-${CUDA_VERSION_DASH} && \ + libcublas-${CUDA_VERSION_DASH} \ + # Required by fastsafetensors (fixes #20384) + libnuma-dev && \ # Fixes nccl_allocator requiring nccl.h at runtime # https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22 # NCCL packages don't use the cuda-MAJOR-MINOR naming convention, @@ -693,15 +679,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ . /etc/environment && \ uv pip list -# Install deepgemm wheel that has been built in the `build` stage -RUN --mount=type=cache,target=/root/.cache/uv \ - --mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \ - sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \ - uv pip install --system /tmp/deepgemm/dist/*.whl; \ - else \ - echo "No DeepGEMM wheels to install; skipping."; \ - fi' - # Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH @@ -750,8 +727,8 @@ ARG PYTORCH_NIGHTLY # Install development dependencies (for testing) COPY requirements/lint.txt requirements/lint.txt -COPY requirements/test.in requirements/test.in -COPY requirements/test.txt requirements/test.txt +COPY requirements/test/cuda.in requirements/test/cuda.in +COPY requirements/test/cuda.txt requirements/test/cuda.txt COPY requirements/dev.txt requirements/dev.txt COPY use_existing_torch.py use_existing_torch.py COPY --from=base /workspace/torch_lib_versions.txt torch_lib_versions.txt @@ -761,8 +738,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \ if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \ echo "Installing dev requirements plus torch nightly..." \ && python3 use_existing_torch.py --prefix \ - && cat torch_lib_versions.txt >> requirements/test.in \ - && uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match \ + && cat torch_lib_versions.txt >> requirements/test/cuda.in \ + && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \ && uv pip install --system $(cat torch_lib_versions.txt | xargs) --pre \ -r requirements/dev.txt \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index afcb388beb2..0600f7da82f 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -107,10 +107,10 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \ fi # Copy build requirements -COPY requirements/cpu-build.txt requirements/build.txt +COPY requirements/build/cpu.txt requirements/build/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/build.txt + uv pip install -r requirements/build/cpu.txt COPY . . @@ -127,26 +127,28 @@ FROM base AS vllm-test-deps WORKDIR /vllm-workspace # Copy test requirements -COPY requirements/test.in requirements/cpu-test.in +COPY requirements/test/cuda.in requirements/test/cpu.in RUN \ - sed -i '/mamba_ssm/d' requirements/cpu-test.in && \ + sed -i '/mamba_ssm/d' requirements/test/cpu.in && \ remove_packages_not_supported_on_aarch64() { \ case "$(uname -m)" in \ aarch64|arm64) \ - sed -i '/decord/d' requirements/cpu-test.in; \ - sed -i '/terratorch/d' requirements/cpu-test.in; \ + sed -i '/decord/d' requirements/test/cpu.in; \ + sed -i '/terratorch/d' requirements/test/cpu.in; \ ;; \ esac; \ }; \ remove_packages_not_supported_on_aarch64 && \ - sed -i 's/^torch==.*/torch==2.11.0/g' requirements/cpu-test.in && \ - sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \ - sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \ - uv pip compile requirements/cpu-test.in -o requirements/cpu-test.txt --index-strategy unsafe-best-match --torch-backend cpu + sed -i 's/^torch==.*/torch==2.11.0/g' requirements/test/cpu.in && \ + sed -i 's/torchaudio.*/torchaudio/g' requirements/test/cpu.in && \ + sed -i 's/torchvision.*/torchvision/g' requirements/test/cpu.in && \ + # Related issue: https://github.com/vllm-project/vllm/pull/38800#issuecomment-4228314305 + sed -i 's/^sentence-transformers.*/sentence-transformers==5.3.0/g' requirements/test/cpu.in && \ + uv pip compile requirements/test/cpu.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install -r requirements/cpu-test.txt + uv pip install -r requirements/test/cpu.txt ######################### DEV IMAGE ######################### FROM vllm-build AS vllm-dev @@ -168,7 +170,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=.git,target=.git \ VLLM_TARGET_DEVICE=cpu python3 setup.py develop -COPY --from=vllm-test-deps /vllm-workspace/requirements/cpu-test.txt requirements/test.txt +COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt RUN --mount=type=cache,target=/root/.cache/uv \ uv pip install -r requirements/dev.txt && \ diff --git a/docker/Dockerfile.nightly_torch b/docker/Dockerfile.nightly_torch index 045a09a42dc..39e1cc18759 100644 --- a/docker/Dockerfile.nightly_torch +++ b/docker/Dockerfile.nightly_torch @@ -107,7 +107,7 @@ COPY . . RUN python3 use_existing_torch.py RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/build.txt + uv pip install --system -r requirements/build/cuda.txt ARG GIT_REPO_CHECK=0 RUN --mount=type=bind,source=.git,target=.git \ @@ -261,7 +261,7 @@ FROM vllm-base as test COPY tests/ tests/ # install build and runtime dependencies without stable torch version -COPY requirements/nightly_torch_test.txt requirements/nightly_torch_test.txt +COPY requirements/test/nightly-torch.txt requirements/test/nightly-torch.txt # This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out # Reference: https://github.com/astral-sh/uv/pull/1694 @@ -277,7 +277,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ ENV HF_HUB_ENABLE_HF_TRANSFER 1 RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --system -r requirements/nightly_torch_test.txt + uv pip install --system -r requirements/test/nightly-torch.txt # Logging to confirm the torch versions RUN pip freeze | grep -E 'torch|vllm|flashinfer' diff --git a/docker/Dockerfile.ppc64le b/docker/Dockerfile.ppc64le index 07b64a509a4..845d900c39c 100644 --- a/docker/Dockerfile.ppc64le +++ b/docker/Dockerfile.ppc64le @@ -251,7 +251,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ make -C /numactl install && \ # sentencepiece.pc is in some pkgconfig inside uv cache export PKG_CONFIG_PATH=$(find / -type d -name "pkgconfig" 2>/dev/null | tr '\n' ':') && \ - nanobind_DIR=$(uv pip show nanobind | grep Location | sed 's/^Location: //;s/$/\/nanobind\/cmake/') && uv pip install -r /src/requirements/common.txt -r /src/requirements/cpu.txt -r /src/requirements/build.txt --no-build-isolation && \ + nanobind_DIR=$(uv pip show nanobind | grep Location | sed 's/^Location: //;s/$/\/nanobind\/cmake/') && uv pip install -r /src/requirements/common.txt -r /src/requirements/cpu.txt -r /src/requirements/build/cuda.txt --no-build-isolation && \ cd /src/ && \ uv build --wheel --out-dir /vllmwheel/ --no-build-isolation && \ uv pip install /vllmwheel/*.whl diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index fad79af9cff..801847d4999 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -19,7 +19,8 @@ ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} # Install some basic utilities RUN apt-get update -q -y && apt-get install -q -y \ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ - apt-transport-https ca-certificates wget curl + apt-transport-https ca-certificates wget curl \ + libnuma-dev RUN python3 -m pip install --upgrade pip # Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) ARG USE_SCCACHE @@ -328,14 +329,14 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ --mount=type=cache,target=/root/.cache/uv \ cd /install \ && uv pip install --system -r requirements/rocm.txt \ - && uv pip install --system -r requirements/rocm-test.txt \ + && uv pip install --system -r requirements/test/rocm.txt \ && pip uninstall -y vllm \ && uv pip install --system *.whl -# Verify that PyTorch is the ROCm build, not CUDA -RUN python3 -c "import torch; assert torch.version.hip is not None, \ - f'Expected ROCm PyTorch but got CUDA (torch.version.cuda={torch.version.cuda}, torch.version.hip={torch.version.hip})'; \ - print(f'Verified: PyTorch {torch.__version__} with ROCm (HIP {torch.version.hip})')" +# Persist the built wheel in the image so python_only_compile_rocm.sh can +# reinstall it after removing compilers. The bind-mounted /install contents +# above are not available once that RUN step completes. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ # Install RIXL wheel RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index 51ce663a1b2..1ab2d7229bc 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -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.12" +ARG AITER_BRANCH="v0.1.10.post3" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="2d02c6a9" ARG MORI_REPO="https://github.com/ROCm/mori.git" diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x index e90f2fdfc4c..6b5d965be76 100644 --- a/docker/Dockerfile.s390x +++ b/docker/Dockerfile.s390x @@ -262,7 +262,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ $OPENCV_WHL_FILE \ $OUTLINES_CORE_WHL_FILE \ --index-strategy unsafe-best-match \ - -r requirements/cpu-build.txt \ + -r requirements/build/cpu.txt \ -r requirements/cpu.txt diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index b8ec60a8232..555c1f14420 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -76,20 +76,14 @@ ENV UV_LINK_MODE="copy" RUN --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,src=requirements/common.txt,target=/workspace/vllm/requirements/common.txt \ --mount=type=bind,src=requirements/xpu.txt,target=/workspace/vllm/requirements/xpu.txt \ - --mount=type=bind,src=requirements/xpu-test.in,target=/workspace/vllm/requirements/xpu-test.in \ + --mount=type=bind,src=requirements/test/xpu.txt,target=/workspace/vllm/requirements/test/xpu.txt \ uv pip install --upgrade pip && \ uv pip install -r requirements/xpu.txt && \ - uv pip compile /workspace/vllm/requirements/xpu-test.in \ - -o /workspace/vllm/requirements/xpu-test.txt \ - -c /workspace/vllm/requirements/xpu.txt \ - --index-strategy unsafe-best-match \ - --extra-index-url ${PIP_EXTRA_INDEX_URL} \ - --python-version ${PYTHON_VERSION} && \ uv pip install grpcio-tools protobuf nanobind && \ source /opt/intel/oneapi/setvars.sh --force && \ source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \ - uv pip install --no-build-isolation -r /workspace/vllm/requirements/xpu-test.txt + uv pip install --no-build-isolation -r /workspace/vllm/requirements/test/xpu.txt diff --git a/docker/versions.json b/docker/versions.json index 71caab20b0c..625e363f64c 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -52,9 +52,6 @@ "vllm_target_device": { "default": "cuda" }, - "DEEPGEMM_GIT_REF": { - "default": "477618cd51baffca09c4b0b87e97c03fe827ef03" - }, "DEEPEP_COMMIT_HASH": { "default": "73b6ea4" }, diff --git a/docs/README.md b/docs/README.md index 4b480c463ab..2bdc3e1a3e4 100644 --- a/docs/README.md +++ b/docs/README.md @@ -25,7 +25,7 @@ hide: vLLM is a fast and easy-to-use library for LLM inference and serving. -Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has evolved into a community-driven project with contributions from both academia and industry. +Originally developed in the [Sky Computing Lab](https://sky.cs.berkeley.edu) at UC Berkeley, vLLM has grown into one of the most active open-source AI projects built and maintained by a diverse community of many dozens of academic institutions and companies from over 2000 contributors. Where to get started with vLLM depends on the type of user. If you are looking to: @@ -42,23 +42,37 @@ vLLM is fast with: - State-of-the-art serving throughput - Efficient management of attention key and value memory with [**PagedAttention**](https://blog.vllm.ai/2023/06/20/vllm.html) -- Continuous batching of incoming requests -- Fast model execution with CUDA/HIP graph -- Quantization: [GPTQ](https://arxiv.org/abs/2210.17323), [AWQ](https://arxiv.org/abs/2306.00978), INT4, INT8, and FP8 -- Optimized CUDA kernels, including integration with FlashAttention and FlashInfer. -- Speculative decoding -- Chunked prefill +- Continuous batching of incoming requests, chunked prefill, prefix caching +- Fast and flexible model execution with piecewise and full CUDA/HIP graphs +- Quantization: FP8, MXFP8/MXFP4, NVFP4, INT8, INT4, GPTQ/AWQ, GGUF, compressed-tensors, ModelOpt, TorchAO, and [more](https://docs.vllm.ai/en/latest/features/quantization/index.html) +- Optimized attention kernels including FlashAttention, FlashInfer, TRTLLM-GEN, FlashMLA, and Triton +- Optimized GEMM/MoE kernels for various precisions using CUTLASS, TRTLLM-GEN, CuTeDSL +- Speculative decoding including n-gram, suffix, EAGLE, DFlash +- Automatic kernel generation and graph-level transformations using torch.compile +- Disaggregated prefill, decode, and encode vLLM is flexible and easy to use with: -- Seamless integration with popular HuggingFace models +- Seamless integration with popular Hugging Face models - High-throughput serving with various decoding algorithms, including *parallel sampling*, *beam search*, and more -- Tensor, pipeline, data and expert parallelism support for distributed inference +- Tensor, pipeline, data, expert, and context parallelism for distributed inference - Streaming outputs -- OpenAI-compatible API server -- Support for NVIDIA GPUs, AMD CPUs and GPUs, Intel CPUs and GPUs, PowerPC CPUs, Arm CPUs, and TPU. Additionally, support for diverse hardware plugins such as Intel Gaudi, IBM Spyre and Huawei Ascend. -- Prefix caching support -- Multi-LoRA support +- Generation of structured outputs using xgrammar or guidance +- Tool calling and reasoning parsers +- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support +- Efficient multi-LoRA support for dense and MoE layers +- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more. + +vLLM seamlessly supports 200+ model architectures on HuggingFace, including: + +- Decoder-only LLMs (e.g., Llama, Qwen, Gemma) +- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS) +- Hybrid attention and state-space models (e.g., Mamba, Qwen3.5) +- Multi-modal models (e.g., LLaVA, Qwen-VL, Pixtral) +- Embedding and retrieval models (e.g., E5-Mistral, GTE, ColBERT) +- Reward and classification models (e.g., Qwen-Math) + +Find the full list of supported models [here](./models/supported_models.md). For more information, check out the following: diff --git a/docs/assets/contributing/dockerfile-stages-dependency.png b/docs/assets/contributing/dockerfile-stages-dependency.png index 5ea354f34a0..a6b0bf7b8fd 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 f78ae8a9536..a7e50c907d5 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -37,6 +37,7 @@ th { | HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` | | HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` | | Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` | +| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` | | Custom | ✅ | ✅ | Local file: `data.jsonl` | | Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` | @@ -239,6 +240,69 @@ vllm bench serve \ --spec-bench-category "summarization" ``` +#### SPEED-Bench Benchmark with Speculative Decoding + +[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k). + +!!! note + This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets. + +First, download the dataset to a folder, using this one liner: + +```bash +curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 - +``` + +The command supports also the following arguments: + +- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets. +- `--output_dir`: download to a specified folder. By default, it will download to the current directory. + +Start a server with speculative decoding: + +```bash +vllm serve meta-llama/Llama-3.3-70B-Instruct \ + --speculative-config $'{"method": "eagle3", + "num_speculative_tokens": 3, + "model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}' +``` + +Run all categories in the Qualitative split: + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --dataset-path "/data/speed_bench" \ + --num-prompts -1 +``` + +Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`. + +Run only a specific category like "multilingual": + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --dataset-path "/data/speed_bench" \ + --num-prompts -1 + --speed-bench-category "multilingual" +``` + +Run all categories in the Throughput split (2k ISL): + +```bash +vllm bench serve \ + --model meta-llama/Llama-3.3-70B-Instruct \ + --dataset-name speed_bench \ + --speed-bench-dataset-subset throughput_2k + --dataset-path "/data/speed_bench/" \ + --num-prompts -1 +``` + +Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card. + #### Other HuggingFaceDataset Examples ```bash diff --git a/docs/contributing/README.md b/docs/contributing/README.md index 24e7d1c5be0..a538a408e98 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -49,10 +49,10 @@ If you are developing vLLM's Python and CUDA/C++ code, install Pytorch first: uv pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu129 ``` -Then install the necessary build dependencies from `requirements/build.txt`, skipping `torch` as it was installed in the previous step: +Then install the necessary build dependencies from `requirements/build/cuda.txt`, skipping `torch` as it was installed in the previous step: ```bash -grep -v '^torch==' requirements/build.txt | uv pip install -r - +grep -v '^torch==' requirements/build/cuda.txt | uv pip install -r - ``` Finally install vLLM using: diff --git a/docs/contributing/incremental_build.md b/docs/contributing/incremental_build.md index cc01a60ce1e..6be35af6131 100644 --- a/docs/contributing/incremental_build.md +++ b/docs/contributing/incremental_build.md @@ -16,10 +16,10 @@ Before setting up the incremental build: 2. **CUDA Toolkit:** Verify that the NVIDIA CUDA Toolkit is correctly installed and `nvcc` is accessible in your `PATH`. CMake relies on `nvcc` to compile CUDA code. You can typically find `nvcc` in `$CUDA_HOME/bin/nvcc` or by running `which nvcc`. If you encounter issues, refer to the [official CUDA Toolkit installation guides](https://developer.nvidia.com/cuda-toolkit-archive) and vLLM's main [GPU installation documentation](../getting_started/installation/gpu.md#troubleshooting) for troubleshooting. The `CMAKE_CUDA_COMPILER` variable in your `CMakeUserPresets.json` should also point to your `nvcc` binary. -3. **Build Tools:** It is highly recommended to install `ccache` for fast rebuilds by caching compilation results (e.g., `sudo apt install ccache` or `conda install ccache`). Also, ensure the core build dependencies like `cmake` and `ninja` are installed. These are installable through `requirements/build.txt` or your system's package manager. +3. **Build Tools:** It is highly recommended to install `ccache` for fast rebuilds by caching compilation results (e.g., `sudo apt install ccache` or `conda install ccache`). Also, ensure the core build dependencies like `cmake` and `ninja` are installed. These are installable through `requirements/build/cuda.txt` or your system's package manager. ```console - uv pip install -r requirements/build.txt --torch-backend=auto + uv pip install -r requirements/build/cuda.txt --torch-backend=auto ``` ## Setting up the CMake Build Environment diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 242cc6b3b1e..65e93657e78 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -178,6 +178,7 @@ Priority is **1 = highest** (tried first). | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A | | `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any | +| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. > diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 5515c91a8b6..9c15f02858c 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -28,6 +28,7 @@ Multiple CUDA Graphs are pre-captured at different **token budget** levels (e.g. class BudgetGraphMetadata: token_budget: int max_batch_size: int + max_frames_per_batch: int graph: torch.cuda.CUDAGraph input_buffer: torch.Tensor # e.g. pixel_values metadata_buffers: dict[str, torch.Tensor] # e.g. embeddings, seq metadata @@ -51,6 +52,15 @@ For each graph replay: When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`. +### Video inference support (experimental) + +Following (ViT full CUDA graph support for image inference), extends the encoder CUDA graph framework to support video inference for Qwen3-VL. Previously, the CUDA graph capture/replay path only handled image inputs (`pixel_values` + `image_grid_thw`). Video inputs use different keys (`pixel_values_videos` + `video_grid_thw`) and require larger `cu_seqlens` buffers because each video item contributes multiple frames (`T` attention sequences). This PR generalizes the protocol and manager to handle both modalities through a single shared graph manager. + +!!! note + Video CUDA graphs are automatically disabled when EVS (Efficient Video Sampling) pruning is enabled, since EVS makes the token count data-dependent and incompatible with CUDA graph capture. + + Currently, we only support image-only or video-only inputs when enabling CUDA graph, mixed inputs (image + video) are not supported yet (we will work on it in the near future). Thus, it's recommended to turn off the image modality by `--limit-mm-per-prompt '{"image": 0}'` for video-only inputs. + ## Model integration via `SupportsEncoderCudaGraph` Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGraph][vllm.model_executor.models.interfaces.SupportsEncoderCudaGraph] protocol. This protocol encapsulates all model-specific logic so that the manager remains model-agnostic. The protocol defines the following methods: @@ -65,12 +75,17 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra * `prepare_encoder_cudagraph_replay_buffers(...)` — computes new buffer values from actual batch inputs before replay. * `encoder_cudagraph_forward(...)` — forward pass using precomputed buffers (called during capture and replay). * `encoder_eager_forward(...)` — fallback eager forward when no graph fits. - -Currently supported: **Qwen3-VL** (see `vllm/model_executor/models/qwen3_vl.py`). +* `get_input_modality(...)` - return the modality of the inputs. !!! note The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager. +**Supported models:** + +| Architecture | Models | CG for Image | CG for Video | +| ------------ | ------ | ------------ | ------------ | +| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | + !!! note Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs. @@ -80,10 +95,13 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs: * `cudagraph_mm_encoder` (`bool`, default `False`) — enable CUDA Graph capture for multimodal encoder. When enabled, captures the full encoder forward as a CUDA Graph for each token budget level. * `encoder_cudagraph_token_budgets` (`list[int]`, default `[]`) — token budget levels for capture. If empty (default), auto-inferred from model architecture as power-of-2 levels. User-provided values override auto-inference. -* `encoder_cudagraph_max_images_per_batch` (`int`, default `0`) — maximum number of images per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`. +* `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`. +* `encoder_cudagraph_max_frames_per_batch` (`int`, default `0`) — maximum number of video frames per batch during capture. If 0 (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * 2` (to be optimized). ## Usage guide +### Image inference + Enable encoder CUDA Graphs via `compilation_config`: ```bash @@ -95,7 +113,7 @@ With explicit budgets: ```bash vllm serve Qwen/Qwen3-VL-32B \ - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}' + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8}' ``` Python example: @@ -107,7 +125,7 @@ compilation_config = { "cudagraph_mm_encoder": True, # Optional: override auto-inferred budgets # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], - # "encoder_cudagraph_max_images_per_batch": 8, + # "encoder_cudagraph_max_vision_items_per_batch": 8, } model = vllm.LLM( @@ -118,6 +136,44 @@ model = vllm.LLM( The manager tracks hit/miss statistics and logs them periodically. A "hit" means an image was processed via CUDA Graph replay; a "miss" means eager fallback (image exceeded all budgets). +### Video inference + +Enable encoder CUDA Graphs via `compilation_config`: + +```bash +vllm serve Qwen/Qwen3-VL-32B \ + --limit-mm-per-prompt '{"image": 0}' \ + --compilation-config '{"cudagraph_mm_encoder": true}' +``` + +With explicit budgets: + +```bash +vllm serve Qwen/Qwen3-VL-32B \ + --limit-mm-per-prompt '{"image": 0}' \ + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_vision_items_per_batch": 8, "encoder_cudagraph_max_frames_per_batch": 64}' +``` + +Python example: + +```python +import vllm + +compilation_config = { + "cudagraph_mm_encoder": True, + # Optional: override auto-inferred budgets + # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], + # "encoder_cudagraph_max_vision_items_per_batch": 8, + # "encoder_cudagraph_max_frames_per_batch": 64, +} + +model = vllm.LLM( + model="Qwen/Qwen3-VL-32B", + limit_mm_per_prompt='{"image": 0}', + compilation_config=compilation_config, +) +``` + ## About the Performance The following benchmarks were run on Blackwell GPUs (GB200) using `vllm bench mm-processor`. See [#35963](https://github.com/vllm-project/vllm/pull/35963) for full details. @@ -140,7 +196,7 @@ vllm bench mm-processor \ --num-prompts 3000 --num-warmups 300 \ --max-model-len 32768 --seed 42 \ --mm-encoder-attn-backend FLASH_ATTN \ - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}' + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}' ``` ### Multi-GPU (4x GB200, TP=4, DP=4) @@ -165,5 +221,8 @@ vllm bench mm-processor \ --max-model-len 8192 --seed 42 \ --mm-encoder-attn-backend FLASHINFER \ --tensor-parallel-size 4 --mm-encoder-tp-mode data \ - --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}' + --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_vision_items_per_batch": 8}' ``` + +!!! note + Find more details about benchmarks on GPUs (A100) for video inference at [#38061](https://github.com/vllm-project/vllm/pull/38061). diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index c9dc1292d09..231bca3646f 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -59,7 +59,7 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes. - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod] - [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod] - [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod] -- [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod] +- [`GptOssMxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.GptOssMxfp4MoEMethod] - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod] ## Fused Experts Kernels @@ -86,7 +86,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k | cutlass_fp4 | standard,
batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] | | cutlass_fp8 | standard,
batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],
[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] | | flashinfer | standard | nvfp4,
fp8 | T | 5 | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] | -| gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],
[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] | +| gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],
[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] | | marlin | standard,
batched | 3 / N/A | 3 / N/A | silu,
swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],
[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],
[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] | | trtllm | standard | mxfp4,
nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],
[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],
[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],
[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] | | rocm aiter moe | standard | mxfp4,
fp8 | G(32),G(128),A,T | silu, gelu,
swigluoai | Y | N | `rocm_aiter_fused_experts`,
`AiterExperts` | diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md index d9542746419..af6da94aa62 100644 --- a/docs/features/disagg_encoder.md +++ b/docs/features/disagg_encoder.md @@ -72,4 +72,4 @@ For the PD disaggregation part, the Prefill instance receives cache exactly the `docs/features/disagg_prefill.md` shows the brief idea about the disaggregated prefill (v0) -We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D; +We create the example setup with the **NixlConnector** from `vllm/distributed/kv_transfer/kv_connector/v1/nixl/` and referred to the `tests/v1/kv_connector/nixl_integration/toy_proxy_server.py` to facilitate the kv transfer between P and D; diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index 9d41185f1ae..5c99eb20bc7 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -249,7 +249,7 @@ Token counting starts from `reasoning_start_str`. Once the reasoning token count To use this feature: - `--reasoning-parser` enables reasoning extraction. -- `--reasoning-config` defines the reasoning boundary tokens (e.g., `reasoning_start_str`, `reasoning_end_str`). +- `--reasoning-config` defines the reasoning boundary tokens (e.g., `reasoning_start_str`, `reasoning_end_str`). If not set, vLLM will attempt to automatically initialize these tokens from the reasoning parser. - `thinking_token_budget` (a sampling parameter) sets the per-request reasoning token limit. If `thinking_token_budget` is not specified, no explicit reasoning limit is applied beyond normal generation constraints such as `max_tokens`. diff --git a/docs/getting_started/installation/cpu.arm.inc.md b/docs/getting_started/installation/cpu.arm.inc.md index b266e96db55..8d9cc65f2be 100644 --- a/docs/getting_started/installation/cpu.arm.inc.md +++ b/docs/getting_started/installation/cpu.arm.inc.md @@ -96,14 +96,14 @@ cd vllm_source Third, install required dependencies: ```bash -uv pip install -r requirements/cpu-build.txt --torch-backend cpu +uv pip install -r requirements/build/cpu.txt --torch-backend cpu uv pip install -r requirements/cpu.txt --torch-backend cpu ``` ??? console "pip" ```bash pip install --upgrade pip - pip install -v -r requirements/cpu-build.txt --extra-index-url https://download.pytorch.org/whl/cpu + pip install -v -r requirements/build/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu pip install -v -r requirements/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu ``` diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index eeb20b8bf06..3078c1537c0 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -46,10 +46,10 @@ Execute the following commands to build and install vLLM from source. Please build the following dependencies, `torchvision`, `pyarrow` from source before building vLLM. ```bash - sed -i '/^torch/d' requirements/build.txt # remove torch from requirements/build.txt since we use nightly builds + sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds uv pip install -v \ --torch-backend auto \ - -r requirements/build.txt \ + -r requirements/build/cuda.txt \ -r requirements/cpu.txt \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ uv pip install dist/*.whl @@ -57,10 +57,10 @@ Execute the following commands to build and install vLLM from source. ??? console "pip" ```bash - sed -i '/^torch/d' requirements/build.txt # remove torch from requirements/build.txt since we use nightly builds + sed -i '/^torch/d' requirements/build/cuda.txt # remove torch from requirements/build/cuda.txt since we use nightly builds pip install -v \ --extra-index-url https://download.pytorch.org/whl/nightly/cpu \ - -r requirements/build.txt \ + -r requirements/build/cuda.txt \ -r requirements/cpu.txt \ VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \ pip install dist/*.whl diff --git a/docs/getting_started/installation/cpu.x86.inc.md b/docs/getting_started/installation/cpu.x86.inc.md index 8b855e919f4..ad051d22dc8 100644 --- a/docs/getting_started/installation/cpu.x86.inc.md +++ b/docs/getting_started/installation/cpu.x86.inc.md @@ -88,14 +88,14 @@ cd vllm_source Install the required dependencies: ```bash -uv pip install -r requirements/cpu-build.txt --torch-backend cpu +uv pip install -r requirements/build/cpu.txt --torch-backend cpu uv pip install -r requirements/cpu.txt --torch-backend cpu ``` ??? console "pip" ```bash pip install --upgrade pip - pip install -v -r requirements/cpu-build.txt --extra-index-url https://download.pytorch.org/whl/cpu + pip install -v -r requirements/build/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu pip install -v -r requirements/cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu ``` diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index e46fecc45cd..db181c60e77 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -1,12 +1,12 @@ --8<-- [start:installation] -vLLM contains pre-compiled C++ and CUDA (12.8) binaries. +vLLM contains pre-compiled C++ and CUDA (12.9) binaries. --8<-- [end:installation] --8<-- [start:requirements] -- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.) +- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, B200, etc.) --8<-- [end:requirements] --8<-- [start:set-up-using-python] @@ -31,7 +31,7 @@ uv pip install vllm --torch-backend=auto pip install vllm --extra-index-url https://download.pytorch.org/whl/cu129 ``` -We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu128`), set `--torch-backend=cu128` (or `UV_TORCH_BACKEND=cu128`). If this doesn't work, try running `uv self update` to update `uv` first. +We recommend leveraging `uv` to [automatically select the appropriate PyTorch index at runtime](https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection) by inspecting the installed CUDA driver version via `--torch-backend=auto` (or `UV_TORCH_BACKEND=auto`). To select a specific backend (e.g., `cu130`), set `--torch-backend=cu130` (or `UV_TORCH_BACKEND=cu130`). If this doesn't work, try running `uv self update` to update `uv` first. !!! note NVIDIA Blackwell GPUs (B200, GB200) require a minimum of CUDA 12.8, so make sure you are installing PyTorch wheels with at least that version. PyTorch itself offers a [dedicated interface](https://pytorch.org/get-started/locally/) to determine the appropriate pip command to run for a given target configuration. @@ -93,7 +93,7 @@ If you only need to change Python code, you can build and install vLLM without c ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -VLLM_USE_PRECOMPILED=1 uv pip install --editable . +VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto ``` This command will do the following: @@ -107,10 +107,10 @@ This command will do the following: 1. If you change C++ or kernel code, you cannot use Python-only build; otherwise you will see an import error about library not found or undefined symbol. 2. If you rebase your dev branch, it is recommended to uninstall vllm and re-run the above command to make sure your libraries are up to date. -In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the main branch was just merged and the wheel is being built. In this case, you can wait for around an hour to try again, or manually assign the previous commit in the installation using the `VLLM_PRECOMPILED_WHEEL_LOCATION` environment variable. +In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the `main` branch was just merged and its precompiled wheel is not available yet. You can wait around an hour and retry, or set `VLLM_PRECOMPILED_WHEEL_COMMIT=nightly` to automatically select the most recent already-built commit on `main`. ```bash -export VLLM_PRECOMPILED_WHEEL_COMMIT=$(git rev-parse HEAD~1) # or earlier commit on main +export VLLM_PRECOMPILED_WHEEL_COMMIT=nightly export VLLM_USE_PRECOMPILED=1 uv pip install --editable . ``` @@ -134,7 +134,7 @@ If you want to modify C++ or CUDA code, you'll need to build vLLM from source. T ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -uv pip install -e . +uv pip install -e . --torch-backend=auto ``` !!! tip @@ -162,7 +162,7 @@ To build vLLM using an existing PyTorch installation: git clone https://github.com/vllm-project/vllm.git cd vllm python use_existing_torch.py -uv pip install -r requirements/build.txt +uv pip install -r requirements/build/cuda.txt uv pip install --no-build-isolation -e . ``` @@ -185,7 +185,7 @@ To achieve this, you can set the environment variable VLLM_CUTLASS_SRC_DIR to po ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -VLLM_CUTLASS_SRC_DIR=/path/to/cutlass uv pip install -e . +VLLM_CUTLASS_SRC_DIR=/path/to/cutlass uv pip install -e . --torch-backend=auto ``` ##### Troubleshooting diff --git a/docs/governance/committers.md b/docs/governance/committers.md index df874418f1c..e6f8f317ebb 100644 --- a/docs/governance/committers.md +++ b/docs/governance/committers.md @@ -56,6 +56,7 @@ Sorted alphabetically by GitHub handle: - [@zhuohan123](https://github.com/zhuohan123): Project lead, RL integration, numerics - [@zou3519](https://github.com/zou3519): Compilation - [@BoyuanFeng](https://github.com/BoyuanFeng): Compilation, CUDAGraph +- [@xuechendi](https://github.com/xuechendi): Intel CPU/XPU integration, KV connector ### Emeritus Committers @@ -175,7 +176,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r - Plugin Interface: @youkaichao, @Yikun - NVIDIA GPU: @pavanimajety - AMD GPU: @gshtras, @tjtanaa -- Intel CPU/GPU: @jikunshang, @bigPYJ1151 +- Intel CPU/GPU: @jikunshang, @bigPYJ1151, @xuechendi - Google TPU: @yaochengji ### Ecosystem Projects diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 49a845e6c25..a798b8e3453 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -46,7 +46,7 @@ mock_if_no_torch( # Mock any version checks by reading from compiled CI requirements -with open(ROOT_DIR / "requirements/test.txt") as f: +with open(ROOT_DIR / "requirements/test/cuda.txt") as f: VERSIONS = dict(line.strip().split("==") for line in f if "==" in line) importlib.metadata.version = lambda name: VERSIONS.get(name) or "0.0.0" diff --git a/docs/mkdocs/hooks/generate_metrics.py b/docs/mkdocs/hooks/generate_metrics.py index 4565861c4f7..97282aaee7d 100644 --- a/docs/mkdocs/hooks/generate_metrics.py +++ b/docs/mkdocs/hooks/generate_metrics.py @@ -19,7 +19,7 @@ METRIC_SOURCE_FILES = [ "output": "spec_decode.inc.md", }, { - "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py", + "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py", "output": "nixl_connector.inc.md", }, {"path": "vllm/v1/metrics/perf.py", "output": "perf.inc.md"}, diff --git a/docs/models/pooling_models/classify.md b/docs/models/pooling_models/classify.md index 69a6fe75d37..c857e4e061e 100644 --- a/docs/models/pooling_models/classify.md +++ b/docs/models/pooling_models/classify.md @@ -267,12 +267,39 @@ You can modify the `problem_type` via problem_type in the Hugging Face config. T Implement alignment with transformers [ForSequenceClassificationLoss](https://github.com/huggingface/transformers/blob/57bb6db6ee4cfaccc45b8d474dfad5a17811ca60/src/transformers/loss/loss_utils.py#L92). -### Logit bias +### Affine Score Calibration -You can modify the `logit_bias` (aka `sigmoid_normalize`) through the logit_bias parameter in `vllm.config.PoolerConfig`. +Affine Score Calibration, also known as [Platt Scaling](https://en.wikipedia.org/wiki/Platt_scaling) (Platt, 1999), is the most widely used method for calibrating classifier outputs into well-calibrated probabilities. + +The calibration follows the transformation: + +`activation((logit - logit_mean) / logit_sigma)` + +| Parameter | Default | Description | +| --------- | ------- | ----------- | +| `logit_mean` | `None` | Mean subtracted from logits (centers scores) | +| `logit_sigma` | `None` | Standard deviation used to scale logits after mean subtraction | + +The computation order is as follows: + +```python +logits -= logit_mean # subtract mean (center scores) +logits /= logit_sigma # divide by sigma (scale) +logits = activation(logits) # e.g. sigmoid +``` + +Example configuration: + +```bash +--pooler-config '{"use_activation": true, "logit_mean": 4.5, "logit_sigma": 1.0}' +``` ## Removed Features ### Remove softmax from PoolingParams We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function. + +### Remove `logit_bias` and `logit_scale` + +`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21. diff --git a/docs/models/pooling_models/token_embed.md b/docs/models/pooling_models/token_embed.md index 3396f4eac2d..b0e094267db 100644 --- a/docs/models/pooling_models/token_embed.md +++ b/docs/models/pooling_models/token_embed.md @@ -71,6 +71,14 @@ Models of any architecture can be converted into embedding models using `--conve If your model is not in the above list, we will try to automatically convert the model using [as_embedding_model][vllm.model_executor.models.adapters.as_embedding_model]. +### Special models + +| Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) | +| ------------ | ------ | ----------------- | -------------------- | ------------------------- | +| `JinaForRanking` | Qwen3-based | `jinaai/jina-reranker-v3` | | | + +jina-reranker-v3 is a listwise document reranker model with a novel `last but not late interaction` architecture. More information can be found at: [examples/pooling/token_embed/jina_reranker_v3_offline.py](../../../examples/pooling/token_embed/jina_reranker_v3_offline.py) + --8<-- [end:supported-token-embed-models] ## Offline Inference diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b0b6c080ac4..ef1f5901ed0 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -550,6 +550,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I+ | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ | | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + IE+ | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I+/ V+ | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ | +| `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + IE+ | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ | | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ | | `Gemma3ForConditionalGeneration` | Gemma 3 | T + IE+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | @@ -660,11 +661,12 @@ Speech2Text models trained specifically for Automatic Speech Recognition. | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `CohereAsrForConditionalGeneration` | Cohere-Transcribe | `CohereLabs/cohere-transcribe-03-2026` | | | | `FireRedASR2ForConditionalGeneration` | FireRedASR2 | `allendou/FireRedASR2-LLM-vllm`, etc. | | | +| `FireRedLIDForConditionalGeneration` | FireRedLID | `PatchyTisa/FireRedLID-vllm`, etc. | | | | `FunASRForConditionalGeneration` | FunASR | `allendou/Fun-ASR-Nano-2512-vllm`, etc. | | | | `Gemma3nForConditionalGeneration` | Gemma3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | | `GlmAsrForConditionalGeneration` | GLM-ASR | `zai-org/GLM-ASR-Nano-2512` | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | `ibm-granite/granite-4.0-1b-speech`, `ibm-granite/granite-speech-3.3-2b`, etc. | ✅︎ | ✅︎ | -| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | | ✅︎ | +| `Qwen3ASRForConditionalGeneration` | Qwen3-ASR | `Qwen/Qwen3-ASR-1.7B`, etc. | ✅︎ | ✅︎ | | `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, etc. | | ✅︎ | | `VoxtralForConditionalGeneration` | Voxtral (Mistral format) | `mistralai/Voxtral-Mini-3B-2507`, `mistralai/Voxtral-Small-24B-2507`, etc. | ✅︎ | ✅︎ | | `WhisperForConditionalGeneration` | Whisper | `openai/whisper-small`, `openai/whisper-large-v3-turbo`, etc. | | | diff --git a/examples/offline_inference/audio_language.py b/examples/offline_inference/audio_language.py index 690aada03ab..c480f1b4145 100755 --- a/examples/offline_inference/audio_language.py +++ b/examples/offline_inference/audio_language.py @@ -537,9 +537,30 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData: ) +# FireRedLID +def run_fireredlid(question: str, audio_count: int) -> ModelRequestData: + assert audio_count == 1, "FireRedLID only supports single audio input per prompt" + model_name = "PatchyTisa/FireRedLID-vllm" + + prompt = "" + + engine_args = EngineArgs( + model=model_name, + max_model_len=8, + max_num_seqs=5, + limit_mm_per_prompt={"audio": audio_count}, + ) + + return ModelRequestData( + engine_args=engine_args, + prompt=prompt, + ) + + model_example_map = { "audioflamingo3": run_audioflamingo3, "cohere_asr": run_cohere_asr, + "fireredlid": run_fireredlid, "funaudiochat": run_funaudiochat, "gemma3n": run_gemma3n, "glmasr": run_glmasr, diff --git a/examples/offline_inference/encoder_decoder_multimodal.py b/examples/offline_inference/encoder_decoder_multimodal.py index 2f72b7d0670..4fc74e9555f 100644 --- a/examples/offline_inference/encoder_decoder_multimodal.py +++ b/examples/offline_inference/encoder_decoder_multimodal.py @@ -55,7 +55,91 @@ def run_whisper(): ) +def run_fireredasr2(): + """ + FireRedASR2 – Automatic Speech Recognition model. + + This model uses a Conformer encoder + Qwen2 LLM decoder architecture + for speech-to-text transcription. Audio is passed via the implicit + prompt format with the ``<|AUDIO|>`` placeholder token. + """ + engine_args = EngineArgs( + model="allendou/FireRedASR2-LLM-vllm", + max_model_len=448, + max_num_seqs=16, + limit_mm_per_prompt={"audio": 1}, + ) + + prompt_str = ( + "<|im_start|>user\n<|AUDIO|>请转写音频为文字<|im_end|>\n<|im_start|>assistant\n" + ) + + prompts = [ + { # Implicit prompt with audio + "prompt": prompt_str, + "multi_modal_data": { + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, + }, + }, + { # Another audio sample + "prompt": prompt_str, + "multi_modal_data": { + "audio": AudioAsset("winning_call").audio_and_sample_rate, + }, + }, + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + +def run_fireredlid(): + """ + FireRedLID – Language Identification model. + + This encoder-decoder model identifies the spoken language of an audio + clip. It outputs at most 2 tokens representing the detected language + (e.g. "en", "zh mandarin"). + """ + engine_args = EngineArgs( + model="PatchyTisa/FireRedLID-vllm", + max_model_len=8, + max_num_seqs=16, + limit_mm_per_prompt={"audio": 1}, + ) + + prompts = [ + { # Test explicit encoder/decoder prompt + "encoder_prompt": { + "prompt": "", + "multi_modal_data": { + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, + }, + }, + "decoder_prompt": "", + }, + { # Another audio sample + "encoder_prompt": { + "prompt": "", + "multi_modal_data": { + "audio": AudioAsset("winning_call").audio_and_sample_rate, + }, + }, + "decoder_prompt": "", + }, + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + model_example_map = { + "fireredasr2": run_fireredasr2, + "fireredlid": run_fireredlid, "whisper": run_whisper, } diff --git a/examples/offline_inference/vision_language.py b/examples/offline_inference/vision_language.py index d62f25c2285..2c0bd52c0e3 100755 --- a/examples/offline_inference/vision_language.py +++ b/examples/offline_inference/vision_language.py @@ -421,6 +421,43 @@ def run_ernie45_vl(questions: list[str], modality: str) -> ModelRequestData: ) +# EXAONE-4.5 +def run_exaone4_5(questions: list[str], modality: str) -> ModelRequestData: + model_name = "LGAI-EXAONE/EXAONE-4.5-33B" + + engine_args = EngineArgs( + model=model_name, + max_model_len=4096, + max_num_seqs=5, + mm_processor_kwargs={ + "min_pixels": 28 * 28, + "max_pixels": 1280 * 28 * 28, + "fps": 1, + }, + limit_mm_per_prompt={modality: 1}, + ) + + if modality == "image": + placeholder = "<|image_pad|>" + elif modality == "video": + placeholder = "<|video_pad|>" + + prompts = [ + ( + "<|system|>\nYou are a helpful assistant.<|endofturn|>\n" + f"<|user|>\n{placeholder}" + f"{question}<|endofturn|>\n" + "<|assistant|>\n" + ) + for question in questions + ] + + return ModelRequestData( + engine_args=engine_args, + prompts=prompts, + ) + + # Fuyu def run_fuyu(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2199,6 +2236,7 @@ model_example_map = { "dots_ocr": run_dots_ocr, "eagle2_5": run_eagle2_5, "ernie45_vl": run_ernie45_vl, + "exaone4_5": run_exaone4_5, "fuyu": run_fuyu, "gemma3": run_gemma3, "gemma3n": run_gemma3n, diff --git a/examples/offline_inference/vision_language_multi_image.py b/examples/offline_inference/vision_language_multi_image.py index 1963ffff791..1a3b7fd954c 100755 --- a/examples/offline_inference/vision_language_multi_image.py +++ b/examples/offline_inference/vision_language_multi_image.py @@ -241,6 +241,41 @@ def load_deepseek_ocr(question: str, image_urls: list[str]) -> ModelRequestData: ) +# exaone4_5 +def load_exaone4_5(question: str, image_urls: list[str]) -> ModelRequestData: + model_name = "LGAI-EXAONE/EXAONE-4.5-33B" + + engine_args = EngineArgs( + model=model_name, + max_model_len=8192, + 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_gemma3(question: str, image_urls: list[str]) -> ModelRequestData: model_name = "google/gemma-3-4b-it" @@ -1450,6 +1485,7 @@ model_example_map = { "command_a_vision": load_command_a_vision, "deepseek_vl_v2": load_deepseek_vl2, "deepseek_ocr": load_deepseek_ocr, + "exaone4_5": load_exaone4_5, "gemma3": load_gemma3, "h2ovl_chat": load_h2ovl, "hunyuan_vl": load_hunyuan_vl, diff --git a/examples/online_serving/kv_events_subscriber.py b/examples/online_serving/kv_events_subscriber.py index 499ab1f3946..0512297fcf4 100644 --- a/examples/online_serving/kv_events_subscriber.py +++ b/examples/online_serving/kv_events_subscriber.py @@ -43,10 +43,13 @@ class BlockStored(KVCacheEvent): prompt embeddings data, etc. for that specific block. """ + group_idx: int | None = None + class BlockRemoved(KVCacheEvent): block_hashes: list[ExternalBlockHash] medium: str | None + group_idx: int | None = None class AllBlocksCleared(KVCacheEvent): diff --git a/examples/online_serving/openai_lid_client.py b/examples/online_serving/openai_lid_client.py new file mode 100644 index 00000000000..0ce0fbc9225 --- /dev/null +++ b/examples/online_serving/openai_lid_client.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Language Identification (LID) demo using the FireRedLID model on vLLM. + +FireRedLID is an audio encoder-decoder model that identifies the spoken +language of an audio clip. Unlike ASR models that output full transcriptions, +FireRedLID outputs at most 2 tokens representing the detected language +(e.g. "en", "zh mandarin"). + +Start the vLLM server: + + vllm serve PatchyTisa/FireRedLID-vllm + +Then run this script: + + # Use the built-in sample audio + python examples/online_serving/openai_lid_client.py + + # Use your own audio file(s) + python examples/online_serving/openai_lid_client.py \ + --audio_paths audio_en.wav audio_zh.wav audio_fr.wav + + # Batch-identify multiple files in one run + python examples/online_serving/openai_lid_client.py \ + --audio_paths /path/to/dir/*.wav + +Requirements: +- vLLM with audio support +- openai Python SDK +- kaldi_native_fbank (pulled in by the model) +""" + +import argparse +import json +import os + +from openai import OpenAI + +from vllm.assets.audio import AudioAsset + +# ────────────────────────────────────────────────────────────────────── +# Helpers +# ────────────────────────────────────────────────────────────────────── + + +def identify_language( + audio_path: str, + client: OpenAI, + model: str, +) -> str: + """ + Send a single audio file to the vLLM transcription endpoint and return + the detected language tag. + + FireRedLID re-uses the OpenAI-compatible ``/v1/audio/transcriptions`` + endpoint. The "transcription" it returns is actually the language label + (e.g. ``"en"`` or ``"zh mandarin"``). + """ + with open(audio_path, "rb") as f: + result = client.audio.transcriptions.create( + file=f, + model=model, + response_format="json", + temperature=0.0, + ) + return result.text.strip() + + +def identify_language_raw( + audio_path: str, + model: str, + api_base: str, +) -> str: + """ + Same as :func:`identify_language` but uses raw HTTP so that the demo + works without the ``openai`` SDK (useful for quick debugging). + """ + import requests + + url = f"{api_base}/audio/transcriptions" + with open(audio_path, "rb") as f: + files = {"file": (os.path.basename(audio_path), f)} + data = { + "model": model, + "response_format": "json", + } + resp = requests.post(url, files=files, data=data) + resp.raise_for_status() + return resp.json()["text"].strip() + + +def identify_language_streaming( + audio_path: str, + model: str, + api_base: str, +) -> str: + """ + Streaming variant – demonstrates the streaming transcription endpoint. + For a 1-2 token output the stream finishes almost instantly, but this + shows that the API path works end-to-end. + """ + import requests + + url = f"{api_base}/audio/transcriptions" + with open(audio_path, "rb") as f: + files = {"file": (os.path.basename(audio_path), f)} + data = { + "stream": "true", + "model": model, + "response_format": "json", + } + response = requests.post(url, files=files, data=data, stream=True) + response.raise_for_status() + + tokens: list[str] = [] + for chunk in response.iter_lines( + chunk_size=8192, decode_unicode=False, delimiter=b"\n" + ): + if not chunk: + continue + payload = json.loads(chunk[len("data: ") :].decode("utf-8")) + choice = payload["choices"][0] + delta = choice.get("delta", {}).get("content", "") + if delta: + tokens.append(delta) + if choice.get("finish_reason") is not None: + break + + return "".join(tokens).strip() + + +# ────────────────────────────────────────────────────────────────────── +# Main +# ────────────────────────────────────────────────────────────────────── + + +def main(args: argparse.Namespace) -> None: + api_base = args.api_base.rstrip("/") + client = OpenAI(api_key="EMPTY", base_url=api_base) + model = client.models.list().data[0].id + print(f"Model : {model}") + print(f"Server: {api_base}\n") + + # Resolve audio paths ------------------------------------------------ + if args.audio_paths: + audio_paths = args.audio_paths + else: + # Fall back to the built-in vLLM sample audios (both are English). + audio_paths = [ + str(AudioAsset("mary_had_lamb").get_local_path()), + str(AudioAsset("winning_call").get_local_path()), + ] + + # Run LID for each file ---------------------------------------------- + print(f"{'Audio File':<50} {'Language (sync)':<20} {'Language (stream)'}") + print("-" * 90) + + for path in audio_paths: + basename = os.path.basename(path) + + # 1) Synchronous via OpenAI SDK + lang_sync = identify_language(path, client, model) + + # 2) Streaming via raw HTTP + lang_stream = identify_language_streaming(path, model, api_base) + + print(f"{basename:<50} {lang_sync:<20} {lang_stream}") + + print() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="FireRedLID – Language Identification demo via vLLM", + ) + parser.add_argument( + "--audio_paths", + nargs="+", + default=None, + help=( + "One or more audio files to identify. " + "If omitted, uses vLLM's built-in sample audios." + ), + ) + parser.add_argument( + "--api_base", + type=str, + default="http://localhost:8000/v1", + help="vLLM API base URL (default: http://localhost:8000/v1)", + ) + args = parser.parse_args() + main(args) diff --git a/examples/pooling/token_embed/jina_reranker_v3_offline.py b/examples/pooling/token_embed/jina_reranker_v3_offline.py new file mode 100644 index 00000000000..c250eccc62a --- /dev/null +++ b/examples/pooling/token_embed/jina_reranker_v3_offline.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +import torch.nn.functional as F + +from vllm import LLM + +query = "What are the health benefits of green tea?" +documents = [ + "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.", + "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.", + "Studies show that drinking green tea regularly can improve brain function and boost metabolism.", + "Basketball is one of the most popular sports in the United States.", + "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。", + "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.", +] + + +def main(): + # Initialize model + llm = LLM( + model="jinaai/jina-reranker-v3", + runner="pooling", + ) + + # Generate scores. + outputs = llm.score(query, documents) + + # Print the outputs. + print("\nGenerated Outputs:\n" + "-" * 60) + for document, output in zip(documents, outputs): + score = output.outputs.score + print(f"Pair: {[query, document]!r} \nScore: {score}") + print("-" * 60) + + # Generate embeddings. + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + outputs = llm.encode(documents + [query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + # Print the outputs. + print("\nGenerated Outputs:\n" + "-" * 60) + for document, score in zip(documents, scores): + print(f"Pair: {[query, document]!r} \nScore: {score}") + print("-" * 60) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 2758c3e0ac1..f55dd9308bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -# Should be mirrored in requirements/build.txt +# Should be mirrored in requirements/build/cuda.txt requires = [ "cmake>=3.26.1", "ninja", @@ -120,7 +120,8 @@ python = "./.venv" [tool.typos.files] # these files may be written in non english words extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*", - "benchmarks/sonnet.txt", "tests/lora/data/*", "examples/pooling/token_embed/*", "build/*", + "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*", + "examples/pooling/token_embed/*", "tests/models/language/pooling/*", "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py", "docs/governance/process.md", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"] ignore-hidden = false @@ -169,6 +170,9 @@ eles = "eles" datas = "datas" ser = "ser" ure = "ure" +# Walsh-Hadamard Transform +wht = "wht" +WHT = "WHT" [tool.uv] no-build-isolation-package = ["torch"] diff --git a/requirements/cpu-build.txt b/requirements/build/cpu.txt similarity index 100% rename from requirements/cpu-build.txt rename to requirements/build/cpu.txt diff --git a/requirements/build.txt b/requirements/build/cuda.txt similarity index 100% rename from requirements/build.txt rename to requirements/build/cuda.txt diff --git a/requirements/rocm-build.txt b/requirements/build/rocm.txt similarity index 94% rename from requirements/rocm-build.txt rename to requirements/build/rocm.txt index b71a847c6a7..9e2ea57e1fd 100644 --- a/requirements/rocm-build.txt +++ b/requirements/build/rocm.txt @@ -1,5 +1,5 @@ # Common dependencies --r common.txt +-r ../common.txt --extra-index-url https://download.pytorch.org/whl/rocm7.1 diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 75831c39e2c..c6b8a82eac6 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -15,6 +15,9 @@ flashinfer-cubin==0.6.7 # breaking changes in 1.19.0 nvidia-cudnn-frontend>=1.13.0,<1.19.0 +# Required for faster safetensors model loading +fastsafetensors >= 0.2.2 + # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) nvidia-cutlass-dsl>=4.4.2 quack-kernels>=0.3.3 diff --git a/requirements/dev.txt b/requirements/dev.txt index e75821eb4a8..fe0b9eaaf96 100644 --- a/requirements/dev.txt +++ b/requirements/dev.txt @@ -1,5 +1,5 @@ -r lint.txt --r test.txt +-r test/cuda.txt # Avoid adding requirements directly to this file. # Instead, modify the two files referenced above. diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index c0b24d99ee1..700213feda1 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -1,3 +1,5 @@ lmcache >= 0.3.9 nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill +nixl-cu12 >= 0.7.1, < 0.10.0 +nixl-cu13 >= 0.7.1, < 0.10.0 mooncake-transfer-engine >= 0.3.8 diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 6639e71a4b9..011d0e53fd1 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -20,4 +20,6 @@ conch-triton-kernels==1.2.1 timm>=1.0.17 # amd-quark: required for Quark quantization on ROCm # To be consistent with test_quark.py -amd-quark>=0.8.99 \ No newline at end of file +amd-quark>=0.8.99 +# Required for faster safetensors model loading +fastsafetensors >= 0.2.2 \ No newline at end of file diff --git a/requirements/test.in b/requirements/test/cuda.in similarity index 100% rename from requirements/test.in rename to requirements/test/cuda.in diff --git a/requirements/test.txt b/requirements/test/cuda.txt similarity index 89% rename from requirements/test.txt rename to requirements/test/cuda.txt index 5675a2a829a..548ca9310ff 100644 --- a/requirements/test.txt +++ b/requirements/test/cuda.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile requirements/test.in -c requirements/common.txt -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12 +# uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12 absl-py==2.1.0 # via # rouge-score @@ -29,7 +29,7 @@ albucore==0.0.16 # via terratorch albumentations==1.4.6 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # terratorch alembic==1.16.4 # via optuna @@ -46,7 +46,7 @@ anyio==4.6.2.post1 # httpx # starlette arctic-inference==0.1.1 - # via -r requirements/test.in + # via -r requirements/test/cuda.in argcomplete==3.5.1 # via datamodel-code-generator arrow==1.3.0 @@ -64,7 +64,7 @@ attrs==24.2.0 audioread==3.0.1 # via librosa av==16.1.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in azure-core==1.38.2 # via # azure-identity @@ -75,16 +75,16 @@ azure-storage-blob==12.28.0 # via runai-model-streamer-azure backoff==2.2.1 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # schemathesis bitsandbytes==0.49.2 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # lightning black==24.10.0 # via datamodel-code-generator blobfile==3.0.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in bm25s==0.2.13 # via mteb boto3==1.35.57 @@ -98,7 +98,7 @@ botocore==1.35.57 bounded-pool-executor==0.0.3 # via pqdm buildkite-test-collector==0.1.9 - # via -r requirements/test.in + # via -r requirements/test/cuda.in cachetools==5.5.2 # via google-auth certifi==2024.8.30 @@ -176,21 +176,21 @@ cupy-cuda12x==13.6.0 cycler==0.12.1 # via matplotlib datamodel-code-generator==0.26.3 - # via -r requirements/test.in + # via -r requirements/test/cuda.in dataproperty==1.0.1 # via # pytablewriter # tabledata datasets==3.3.0 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # evaluate # lm-eval # mteb decorator==5.1.1 # via librosa decord==0.6.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in diffusers==0.36.0 # via terratorch dill==0.3.8 @@ -211,7 +211,7 @@ docstring-parser==0.17.0 # via jsonargparse einops==0.8.1 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # encodec # terratorch # torchgeo @@ -236,7 +236,9 @@ fastparquet==2024.11.0 fastrlock==0.8.2 # via cupy-cuda12x fastsafetensors==0.2.2 - # via -r requirements/test.in + # via + # -c requirements/cuda.txt + # -r requirements/test/cuda.in filelock==3.16.1 # via # -c requirements/common.txt @@ -273,7 +275,7 @@ fsspec==2024.12.0 ftfy==6.3.1 # via open-clip-torch genai-perf==0.0.16 - # via -r requirements/test.in + # via -r requirements/test/cuda.in genson==1.3.0 # via datamodel-code-generator geopandas==1.0.1 @@ -306,19 +308,19 @@ google-resumable-media==2.7.2 googleapis-common-protos==1.70.0 # via google-api-core gpt-oss==0.0.8 - # via -r requirements/test.in + # 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.in + # -r requirements/test/cuda.in # grpcio-reflection # ray # tensorboard grpcio-reflection==1.78.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in h11==0.14.0 # via # httpcore @@ -341,7 +343,7 @@ httpcore==1.0.6 # via httpx httpx==0.27.2 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # diffusers # perceptron # schemathesis @@ -386,7 +388,7 @@ idna==3.10 # requests # yarl imagehash==4.3.2 - # via -r requirements/test.in + # via -r requirements/test/cuda.in imageio==2.37.0 # via scikit-image importlib-metadata==8.7.0 @@ -400,7 +402,7 @@ inflect==5.6.2 iniconfig==2.0.0 # via pytest instanttensor==0.1.5 - # via -r requirements/test.in + # via -r requirements/test/cuda.in isodate==0.7.2 # via azure-storage-blob isoduration==20.11.0 @@ -414,7 +416,7 @@ jinja2==3.1.6 # lm-eval # torch jiwer==3.0.5 - # via -r requirements/test.in + # via -r requirements/test/cuda.in jmespath==1.0.1 # via # boto3 @@ -445,7 +447,7 @@ jsonschema-specifications==2024.10.1 junit-xml==1.9 # via schemathesis kaldi-native-fbank==1.22.3 - # via -r requirements/test.in + # via -r requirements/test/cuda.in kaleido==0.2.1 # via genai-perf kiwisolver==1.4.7 @@ -461,7 +463,7 @@ lazy-loader==0.4 libnacl==2.1.0 # via tensorizer librosa==0.10.2.post1 - # via -r requirements/test.in + # via -r requirements/test/cuda.in lightly==1.5.22 # via # terratorch @@ -480,7 +482,7 @@ lightning-utilities==0.14.3 llvmlite==0.44.0 # via numba lm-eval==0.4.11 - # via -r requirements/test.in + # via -r requirements/test/cuda.in lxml==5.3.0 # via # blobfile @@ -499,7 +501,7 @@ markupsafe==3.0.1 # werkzeug matplotlib==3.9.2 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # lightning # pycocotools # torchgeo @@ -513,7 +515,7 @@ mdurl==0.1.2 mistral-common==1.11.0 # via # -c requirements/common.txt - # -r requirements/test.in + # -r requirements/test/cuda.in more-itertools==10.5.0 # via lm-eval mpmath==1.3.0 @@ -529,7 +531,7 @@ msgpack==1.1.0 # librosa # ray mteb==2.8.3 - # via -r requirements/test.in + # via -r requirements/test/cuda.in multidict==6.1.0 # via # aiohttp @@ -547,15 +549,16 @@ networkx==3.2.1 nltk==3.9.1 # via rouge-score num2words==0.5.14 - # via -r requirements/test.in + # via -r requirements/test/cuda.in numba==0.61.2 # via - # -r requirements/test.in + # -c requirements/cuda.txt + # -r requirements/test/cuda.in # librosa # resampy numpy==2.2.6 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # accelerate # albucore # albumentations @@ -661,7 +664,7 @@ omegaconf==2.3.0 # hydra-core # lightning open-clip-torch==2.32.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in openai-harmony==0.0.4 # via # -c requirements/common.txt @@ -673,12 +676,12 @@ opencensus-context==0.1.3 opencv-python-headless==4.13.0.90 # via # -c requirements/common.txt - # -r requirements/test.in + # -r requirements/test/cuda.in # albucore # albumentations # mistral-common openpyxl==3.1.5 - # via -r requirements/test.in + # via -r requirements/test/cuda.in opentelemetry-api==1.35.0 # via # -c requirements/common.txt @@ -754,9 +757,9 @@ pathvalidate==3.2.1 patsy==1.0.1 # via statsmodels peft==0.16.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in perceptron==0.1.4 - # via -r requirements/test.in + # via -r requirements/test/cuda.in perf-analyzer==0.1.0 # via genai-perf pillow==10.4.0 @@ -782,7 +785,7 @@ platformdirs==4.3.6 # wandb plotly==5.24.1 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # genai-perf pluggy==1.5.0 # via @@ -795,7 +798,7 @@ pooch==1.8.2 portalocker==2.10.1 # via sacrebleu pqdm==0.2.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in prometheus-client==0.22.0 # via # -c requirements/common.txt @@ -852,7 +855,7 @@ pycryptodomex==3.22.0 pydantic==2.12.0 # via # -c requirements/common.txt - # -r requirements/test.in + # -r requirements/test/cuda.in # albumentations # datamodel-code-generator # fastapi @@ -891,7 +894,7 @@ pytablewriter==1.2.0 # via lm-eval pytest==8.3.5 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # buildkite-test-collector # genai-perf # pytest-asyncio @@ -904,21 +907,21 @@ pytest==8.3.5 # pytest-timeout # schemathesis pytest-asyncio==0.24.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in pytest-cov==6.3.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in pytest-forked==1.6.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in pytest-mock==3.14.0 # via genai-perf pytest-rerunfailures==14.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in pytest-shard==0.1.2 - # via -r requirements/test.in + # via -r requirements/test/cuda.in pytest-subtests==0.14.1 # via schemathesis pytest-timeout==2.3.1 - # via -r requirements/test.in + # via -r requirements/test/cuda.in python-box==7.3.2 # via terratorch python-dateutil==2.9.0.post0 @@ -972,7 +975,7 @@ rasterio==1.4.3 # terratorch # torchgeo ray==2.48.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in redis==5.2.0 # via tensorizer referencing==0.35.1 @@ -1015,7 +1018,7 @@ requests==2.32.3 # transformers # wandb resampy==0.4.3 - # via -r requirements/test.in + # via -r requirements/test/cuda.in responses==0.25.3 # via genai-perf rfc3339-validator==0.1.4 @@ -1043,7 +1046,7 @@ rsa==4.9.1 rtree==1.4.0 # via torchgeo runai-model-streamer==0.15.7 - # via -r requirements/test.in + # via -r requirements/test/cuda.in runai-model-streamer-azure==0.15.7 # via runai-model-streamer runai-model-streamer-gcs==0.15.7 @@ -1064,7 +1067,7 @@ safetensors==0.4.5 # timm # transformers schemathesis==3.39.15 - # via -r requirements/test.in + # via -r requirements/test/cuda.in scikit-image==0.25.2 # via # albumentations @@ -1091,12 +1094,12 @@ scipy==1.13.1 # vocos segmentation-models-pytorch==0.5.0 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # terratorch # torchgeo sentence-transformers==5.2.0 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # mteb sentry-sdk==2.52.0 # via wandb @@ -1136,7 +1139,7 @@ sortedcontainers==2.4.0 # via hypothesis soundfile==0.12.1 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # genai-perf # librosa # mistral-common @@ -1172,7 +1175,7 @@ tabulate==0.9.0 tacoreader==0.5.6 # via terratorch tblib==3.1.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in tcolorpy==0.1.6 # via pytablewriter tenacity==9.1.2 @@ -1187,13 +1190,13 @@ tensorboard-data-server==0.7.2 tensorboardx==2.6.4 # via lightning tensorizer==2.10.1 - # via -r requirements/test.in + # via -r requirements/test/cuda.in termcolor==3.1.0 # via # gpt-oss # terratorch terratorch==1.2.2 - # via -r requirements/test.in + # via -r requirements/test/cuda.in threadpoolctl==3.5.0 # via scikit-learn tifffile==2025.3.30 @@ -1208,7 +1211,7 @@ tiktoken==0.12.0 # mistral-common timm==1.0.17 # via - # -r requirements/test.in + # -r requirements/test/cuda.in # open-clip-torch # segmentation-models-pytorch # terratorch @@ -1216,7 +1219,7 @@ timm==1.0.17 tokenizers==0.22.0 # via # -c requirements/common.txt - # -r requirements/test.in + # -r requirements/test/cuda.in # transformers tomli==2.2.1 # via schemathesis @@ -1224,7 +1227,8 @@ tomli-w==1.2.0 # via schemathesis torch==2.11.0+cu130 # via - # -r requirements/test.in + # -c requirements/cuda.txt + # -r requirements/test/cuda.in # accelerate # bitsandbytes # encodec @@ -1249,7 +1253,8 @@ torch==2.11.0+cu130 # vocos torchaudio==2.11.0+cu130 # via - # -r requirements/test.in + # -c requirements/cuda.txt + # -r requirements/test/cuda.in # encodec # vocos torchgeo==0.7.0 @@ -1262,7 +1267,8 @@ torchmetrics==1.7.4 # torchgeo torchvision==0.26.0+cu130 # via - # -r requirements/test.in + # -c requirements/cuda.txt + # -r requirements/test/cuda.in # lightly # open-clip-torch # segmentation-models-pytorch @@ -1292,17 +1298,17 @@ tqdm==4.67.3 transformers==4.57.5 # via # -c requirements/common.txt - # -r requirements/test.in + # -r requirements/test/cuda.in # genai-perf # peft # sentence-transformers # transformers-stream-generator transformers-stream-generator==0.0.5 - # via -r requirements/test.in + # via -r requirements/test/cuda.in triton==3.6.0 # via torch tritonclient==2.64.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in typepy==1.3.2 # via # dataproperty @@ -1371,11 +1377,11 @@ urllib3==2.2.3 uvicorn==0.35.0 # via gpt-oss vector-quantize-pytorch==1.21.2 - # via -r requirements/test.in + # via -r requirements/test/cuda.in virtualenv==20.31.2 # via ray vocos==0.1.0 - # via -r requirements/test.in + # via -r requirements/test/cuda.in wandb==0.24.2 # via terratorch wcwidth==0.2.13 diff --git a/requirements/nightly_torch_test.txt b/requirements/test/nightly-torch.txt similarity index 100% rename from requirements/nightly_torch_test.txt rename to requirements/test/nightly-torch.txt diff --git a/requirements/rocm-test.in b/requirements/test/rocm.in similarity index 97% rename from requirements/rocm-test.in rename to requirements/test/rocm.in index 23c3a0f91e0..b5a9451b36f 100644 --- a/requirements/rocm-test.in +++ b/requirements/test/rocm.in @@ -1,4 +1,4 @@ --r common.txt +-r ../common.txt # testing pytest @@ -78,7 +78,7 @@ datasets>=3.3.0,<=3.6.0 openpyxl # required for perf comparison excel report plotly # required for perf comparison html report -# ROCm-specific extras (not in CUDA test.in) +# ROCm-specific extras (not in CUDA cuda.in) rapidfuzz torchgeo==0.7.0 multiprocess==0.70.16 diff --git a/requirements/rocm-test.txt b/requirements/test/rocm.txt similarity index 81% rename from requirements/rocm-test.txt rename to requirements/test/rocm.txt index a441bfef04d..e1efae912ee 100644 --- a/requirements/rocm-test.txt +++ b/requirements/test/rocm.txt @@ -1,5 +1,5 @@ # This file was autogenerated by uv via the following command: -# uv pip compile requirements/rocm-test.in -o requirements/rocm-test.txt --index-strategy unsafe-best-match -c requirements/rocm.txt --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx +# uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13 absl-py==2.4.0 # via # rouge-score @@ -15,7 +15,7 @@ aiohappyeyeballs==2.6.1 aiohttp==3.13.3 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # aiohttp-cors # fsspec # gpt-oss @@ -29,7 +29,7 @@ albucore==0.1.2 # via terratorch albumentations==1.4.6 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # terratorch alembic==1.18.4 # via optuna @@ -42,7 +42,7 @@ annotated-types==0.7.0 anthropic==0.89.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt antlr4-python3-runtime==4.9.3 # via # hydra-core @@ -57,7 +57,7 @@ anyio==4.13.0 # starlette # watchfiles arctic-inference==0.1.1 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in argcomplete==3.6.3 # via datamodel-code-generator arrow==1.4.0 @@ -76,7 +76,7 @@ attrs==26.1.0 audioread==3.0.1 # via librosa av==16.1.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in azure-core==1.39.0 # via # azure-identity @@ -87,18 +87,18 @@ azure-storage-blob==12.28.0 # via runai-model-streamer-azure backoff==2.2.1 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # schemathesis bitsandbytes==0.49.2 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # lightning black==26.3.1 # via datamodel-code-generator blake3==1.0.8 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt blobfile==3.0.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in bm25s==0.2.13 # via mteb boto3==1.42.74 @@ -112,11 +112,11 @@ botocore==1.42.74 bounded-pool-executor==0.0.3 # via pqdm buildkite-test-collector==0.1.9 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in cachetools==7.0.5 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt cbor2==5.9.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt certifi==2026.2.25 # via # fiona @@ -162,7 +162,7 @@ cligj==0.7.2 # fiona # rasterio cloudpickle==3.1.2 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt colorama==0.4.6 # via # perceptron @@ -175,7 +175,7 @@ colorlog==6.10.1 compressed-tensors==0.14.0.1 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt contourpy==1.3.3 # via matplotlib coverage==7.13.5 @@ -192,25 +192,25 @@ cryptography==46.0.0 cycler==0.12.1 # via matplotlib datamodel-code-generator==0.55.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in dataproperty==1.1.0 # via # pytablewriter # tabledata datasets==3.6.0 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # evaluate # lm-eval # mteb decorator==5.2.1 # via librosa decord==0.6.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in depyf==0.20.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt diffusers==0.37.0 # via terratorch dill==0.3.8 @@ -223,7 +223,7 @@ dill==0.3.8 diskcache==5.6.3 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt distlib==0.4.0 # via virtualenv distro==1.9.0 @@ -242,8 +242,8 @@ docstring-parser==0.17.0 # jsonargparse einops==0.8.2 # via - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in # encodec # terratorch # torchgeo @@ -264,7 +264,7 @@ evaluate==0.4.6 fastapi==0.135.2 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss # model-hosting-container-standards fastapi-cli==0.0.24 @@ -276,11 +276,13 @@ fastar==0.9.0 fastparquet==2026.3.0 # via genai-perf fastsafetensors==0.2.2 - # via -r requirements/rocm-test.in + # via + # -c requirements/rocm.txt + # -r requirements/test/rocm.in filelock==3.25.2 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # blobfile # datasets # diffusers @@ -315,7 +317,7 @@ fsspec==2025.3.0 ftfy==6.3.1 # via open-clip-torch genai-perf==0.0.16 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in genson==1.3.0 # via datamodel-code-generator geopandas==1.1.3 @@ -323,7 +325,7 @@ geopandas==1.1.3 gguf==0.18.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt gitdb==4.0.12 # via gitpython gitpython==3.1.46 @@ -355,7 +357,7 @@ googleapis-common-protos==1.73.0 # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http gpt-oss==0.0.8 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in graphql-core==3.2.8 # via hypothesis-graphql greenlet==3.3.2 @@ -363,7 +365,7 @@ greenlet==3.3.2 grpcio==1.78.0 # via # -c requirements/rocm.txt - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # grpcio-reflection # opentelemetry-exporter-otlp-proto-grpc # ray @@ -371,7 +373,7 @@ grpcio==1.78.0 grpcio-reflection==1.78.0 # via # -c requirements/rocm.txt - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in h11==0.16.0 # via # httpcore @@ -396,7 +398,7 @@ httptools==0.7.1 # via uvicorn httpx==0.27.2 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # anthropic # diffusers # fastapi @@ -410,7 +412,7 @@ httpx-sse==0.4.3 # via mcp huggingface-hub==0.36.2 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # accelerate # datasets # diffusers @@ -450,9 +452,9 @@ idna==3.11 # requests # yarl ijson==3.5.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt imagehash==4.3.2 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in imageio==2.37.3 # via scikit-image importlib-metadata==8.7.1 @@ -466,7 +468,7 @@ inflect==7.5.0 iniconfig==2.3.0 # via pytest instanttensor==0.1.6 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in interegular==0.3.3 # via lm-format-enforcer isodate==0.7.2 @@ -487,7 +489,7 @@ jiter==0.13.0 # anthropic # openai jiwer==4.0.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in jmespath==1.1.0 # via # boto3 @@ -520,7 +522,7 @@ jsonschema-specifications==2025.9.1 junit-xml==1.9 # via schemathesis kaldi-native-fbank==1.22.3 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in kaleido==1.0.0 # via genai-perf kiwisolver==1.5.0 @@ -532,7 +534,7 @@ kornia-rs==0.1.10 lark==1.2.2 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt lazy-loader==0.4 # via # librosa @@ -540,7 +542,7 @@ lazy-loader==0.4 libnacl==2.1.0 # via tensorizer librosa==0.10.2.post1 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in lightly==1.5.22 # via # terratorch @@ -559,15 +561,15 @@ lightning-utilities==0.15.3 llguidance==1.3.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt llvmlite==0.44.0 # via numba lm-eval==0.4.11 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in lm-format-enforcer==0.11.3 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt logistro==2.0.1 # via # choreographer @@ -592,7 +594,7 @@ markupsafe==3.0.3 # werkzeug matplotlib==3.10.8 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # lightning # torchgeo mbstrdecoder==1.1.4 @@ -601,18 +603,18 @@ mbstrdecoder==1.1.4 # pytablewriter # typepy mcp==1.27.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt mdurl==0.1.2 # via markdown-it-py mistral-common==1.11.0 # via # -c requirements/common.txt - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in model-hosting-container-standards==0.1.14 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt more-itertools==10.8.0 # via # inflect @@ -630,16 +632,16 @@ msgpack==1.1.2 # librosa # ray msgspec==0.20.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt mteb==2.11.5 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in multidict==6.7.1 # via # aiohttp # yarl multiprocess==0.70.16 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # datasets # evaluate mypy-extensions==1.1.0 @@ -651,23 +653,23 @@ networkx==3.6.1 # scikit-image # torch ninja==1.13.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt nltk==3.9.3 # via rouge-score num2words==0.5.14 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in numba==0.61.2 # via # -c requirements/rocm.txt - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # librosa # resampy numkong==7.1.1 # via albucore numpy==2.2.6 # via - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in # accelerate # albucore # albumentations @@ -739,15 +741,15 @@ omegaconf==2.3.0 # hydra-core # lightning open-clip-torch==2.32.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in openai==2.30.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt openai-harmony==0.0.8 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss opencensus==0.11.4 # via ray @@ -756,16 +758,16 @@ opencensus-context==0.1.3 opencv-python-headless==4.13.0.92 # via # -c requirements/common.txt - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in # albumentations # mistral-common openpyxl==3.1.5 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in opentelemetry-api==1.40.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # opentelemetry-exporter-prometheus @@ -774,7 +776,7 @@ opentelemetry-api==1.40.0 opentelemetry-exporter-otlp==1.40.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt opentelemetry-exporter-otlp-proto-common==1.40.0 # via # opentelemetry-exporter-otlp-proto-grpc @@ -794,7 +796,7 @@ opentelemetry-proto==1.40.0 opentelemetry-sdk==1.40.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # opentelemetry-exporter-otlp-proto-grpc # opentelemetry-exporter-otlp-proto-http # opentelemetry-exporter-prometheus @@ -807,7 +809,7 @@ opentelemetry-semantic-conventions==0.61b0 opentelemetry-semantic-conventions-ai==0.5.1 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt optuna==3.6.1 # via genai-perf orjson==3.11.7 @@ -817,7 +819,7 @@ orjson==3.11.7 outlines-core==0.2.11 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt packaging==26.0 # via # -c requirements/rocm.txt @@ -868,7 +870,7 @@ pandas==3.0.1 # torchgeo # xarray partial-json-parser==0.2.1.1.post7 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt pathspec==1.0.4 # via black pathvalidate==3.3.1 @@ -876,14 +878,14 @@ pathvalidate==3.3.1 patsy==1.0.2 # via statsmodels peft==0.18.1 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in perceptron==0.1.4 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in perf-analyzer==0.1.0 # via genai-perf pillow==12.1.1 # via - # -r requirements/common.txt + # -r requirements/test/../common.txt # diffusers # genai-perf # imagehash @@ -906,7 +908,7 @@ platformdirs==4.3.6 # wandb plotly==6.6.0 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # genai-perf pluggy==1.6.0 # via @@ -921,18 +923,18 @@ pooch==1.8.2 portalocker==3.2.0 # via sacrebleu pqdm==0.2.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in prometheus-client==0.24.1 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # opentelemetry-exporter-prometheus # prometheus-fastapi-instrumentator # ray prometheus-fastapi-instrumentator==7.1.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt propcache==0.4.1 # via # aiohttp @@ -942,7 +944,7 @@ proto-plus==1.27.1 protobuf==6.33.6 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # google-api-core # googleapis-common-protos # grpcio-reflection @@ -955,14 +957,14 @@ protobuf==6.33.6 # wandb psutil==7.2.2 # via - # -r requirements/common.txt + # -r requirements/test/../common.txt # accelerate # peft # tensorizer py==1.11.0 # via pytest-forked py-cpuinfo==9.0.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt py-spy==0.4.1 # via ray pyarrow==23.0.1 @@ -976,7 +978,7 @@ pyasn1==0.6.3 pyasn1-modules==0.4.2 # via google-auth pybase64==1.4.3 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt pycocotools==2.0.11 # via terratorch pycountry==26.2.16 @@ -988,8 +990,8 @@ pycryptodomex==3.23.0 pydantic==2.12.5 # via # -c requirements/common.txt - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in # albumentations # anthropic # compressed-tensors @@ -1045,7 +1047,7 @@ pytablewriter==1.2.1 # via lm-eval pytest==8.3.5 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # buildkite-test-collector # genai-perf # pytest-asyncio @@ -1058,21 +1060,21 @@ pytest==8.3.5 # pytest-timeout # schemathesis pytest-asyncio==0.24.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in pytest-cov==6.3.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in pytest-forked==1.6.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in pytest-mock==3.15.1 # via genai-perf pytest-rerunfailures==14.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in pytest-shard==0.1.2 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in pytest-subtests==0.14.2 # via schemathesis pytest-timeout==2.3.1 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in python-box==7.4.1 # via terratorch python-dateutil==2.9.0.post0 @@ -1090,7 +1092,7 @@ python-dotenv==1.2.2 # pydantic-settings # uvicorn python-json-logger==4.1.0 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt python-multipart==0.0.22 # via # fastapi @@ -1111,7 +1113,7 @@ pywavelets==1.9.0 # via imagehash pyyaml==6.0.3 # via - # -r requirements/common.txt + # -r requirements/test/../common.txt # accelerate # albumentations # datamodel-code-generator @@ -1137,10 +1139,10 @@ pyyaml==6.0.3 pyzmq==27.1.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt rapidfuzz==3.12.1 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # jiwer rasterio==1.5.0 # via @@ -1148,7 +1150,7 @@ rasterio==1.5.0 # terratorch # torchgeo ray==2.54.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in redis==7.3.0 # via tensorizer referencing==0.37.0 @@ -1157,7 +1159,7 @@ referencing==0.37.0 # jsonschema-specifications regex==2026.2.28 # via - # -r requirements/common.txt + # -r requirements/test/../common.txt # diffusers # nltk # open-clip-torch @@ -1167,7 +1169,7 @@ regex==2026.2.28 requests==2.32.5 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # azure-core # buildkite-test-collector # datasets @@ -1195,7 +1197,7 @@ requests==2.32.5 # transformers # wandb resampy==0.4.3 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in responses==0.26.0 # via genai-perf rfc3339-validator==0.1.4 @@ -1230,7 +1232,7 @@ rtree==1.4.1 runai-model-streamer==0.15.7 # via # -c requirements/rocm.txt - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in runai-model-streamer-azure==0.15.7 # via runai-model-streamer runai-model-streamer-gcs==0.15.7 @@ -1251,7 +1253,7 @@ safetensors==0.7.0 # timm # transformers schemathesis==3.39.15 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in scikit-image==0.26.0 # via # albumentations @@ -1279,26 +1281,26 @@ scipy==1.17.1 # vocos segmentation-models-pytorch==0.5.0 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # terratorch # torchgeo sentence-transformers==5.3.0 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # mteb sentencepiece==0.2.1 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt sentry-sdk==2.55.0 # via # fastapi-cloud-cli # wandb setproctitle==1.3.7 - # via -r requirements/common.txt + # via -r requirements/test/../common.txt setuptools==79.0.1 # via # -c requirements/common.txt # -c requirements/rocm.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # model-hosting-container-standards # pytablewriter # tensorboard @@ -1316,7 +1318,7 @@ simplejson==3.20.2 six==1.17.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # junit-xml # lightly # opencensus @@ -1336,7 +1338,7 @@ sortedcontainers==2.4.0 # via hypothesis soundfile==0.13.1 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # genai-perf # librosa # mistral-common @@ -1382,7 +1384,7 @@ tabulate==0.10.0 tacoreader==0.5.6 # via terratorch tblib==3.1.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in tcolorpy==0.1.7 # via pytablewriter tenacity==9.1.4 @@ -1398,13 +1400,13 @@ tensorboardx==2.6.4 tensorizer==2.10.1 # via # -c requirements/rocm.txt - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in termcolor==3.3.0 # via # gpt-oss # terratorch terratorch==1.2.2 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in threadpoolctl==3.6.0 # via scikit-learn tifffile==2026.3.3 @@ -1414,14 +1416,14 @@ tifffile==2026.3.3 tiktoken==0.12.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # gpt-oss # lm-eval # mistral-common timm==1.0.17 # via # -c requirements/rocm.txt - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # open-clip-torch # segmentation-models-pytorch # terratorch @@ -1429,8 +1431,8 @@ timm==1.0.17 tokenizers==0.22.0 # via # -c requirements/common.txt - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in # transformers tomli==2.4.0 # via schemathesis @@ -1438,7 +1440,7 @@ tomli-w==1.2.0 # via schemathesis torchgeo==0.7.0 # via - # -r requirements/rocm-test.in + # -r requirements/test/rocm.in # terratorch torchmetrics==1.9.0 # via @@ -1448,7 +1450,7 @@ torchmetrics==1.9.0 # torchgeo tqdm==4.67.3 # via - # -r requirements/common.txt + # -r requirements/test/../common.txt # datasets # evaluate # gguf @@ -1472,8 +1474,8 @@ tqdm==4.67.3 transformers==4.57.5 # via # -c requirements/common.txt - # -r requirements/common.txt - # -r requirements/rocm-test.in + # -r requirements/test/../common.txt + # -r requirements/test/rocm.in # compressed-tensors # genai-perf # peft @@ -1481,9 +1483,9 @@ transformers==4.57.5 # transformers-stream-generator # xgrammar transformers-stream-generator==0.0.5 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in tritonclient==2.66.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in typeguard==4.5.1 # via inflect typepy==1.3.4 @@ -1502,7 +1504,7 @@ typeshed-client==2.9.0 typing-extensions==4.15.0 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt # aiosignal # albumentations # alembic @@ -1575,16 +1577,16 @@ uvicorn==0.42.0 uvloop==0.22.1 # via uvicorn vector-quantize-pytorch==1.28.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in virtualenv==21.2.0 # via ray vocos==0.1.0 - # via -r requirements/rocm-test.in + # via -r requirements/test/rocm.in wandb==0.25.1 # via terratorch watchfiles==1.1.1 # via - # -r requirements/common.txt + # -r requirements/test/../common.txt # uvicorn wcwidth==0.6.0 # via ftfy @@ -1605,7 +1607,7 @@ xarray==2026.2.0 xgrammar==0.1.33 # via # -c requirements/common.txt - # -r requirements/common.txt + # -r requirements/test/../common.txt xxhash==3.6.0 # via # datasets @@ -1632,14 +1634,14 @@ zstandard==0.25.0 # nvidia-cuda-cupti # nvidia-cuda-nvrtc # nvidia-cuda-runtime -# nvidia-cudnn-cu13 # nvidia-cufft # nvidia-cufile # nvidia-curand # nvidia-cusolver # nvidia-cusparse +# nvidia-nvjitlink +# nvidia-nvtx +# nvidia-cudnn-cu13 # nvidia-cusparselt-cu13 # nvidia-nccl-cu13 -# nvidia-nvjitlink # nvidia-nvshmem-cu13 -# nvidia-nvtx diff --git a/requirements/xpu-test.in b/requirements/test/xpu.in similarity index 100% rename from requirements/xpu-test.in rename to requirements/test/xpu.in diff --git a/requirements/xpu-test.txt b/requirements/test/xpu.txt similarity index 89% rename from requirements/xpu-test.txt rename to requirements/test/xpu.txt index 2a83fd90f27..51810592c46 100644 --- a/requirements/xpu-test.txt +++ b/requirements/test/xpu.txt @@ -1,11 +1,11 @@ # This file was autogenerated by uv via the following command: -# uv pip compile requirements/xpu-test.in -o requirements/xpu-test.txt -c requirements/xpu.txt --python-version 3.12 --index-strategy unsafe-best-match +# uv pip compile requirements/test/xpu.in -c requirements/xpu.txt -o requirements/test/xpu.txt --index-strategy unsafe-best-match --torch-backend xpu --python-platform x86_64-manylinux_2_39 --python-version 3.12 absl-py==2.4.0 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # rouge-score accelerate==1.13.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in aiohappyeyeballs==2.6.1 # via aiohttp aiohttp==3.13.4 @@ -17,7 +17,7 @@ aiohttp==3.13.4 aiosignal==1.4.0 # via aiohttp albumentations==1.4.6 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in annotated-doc==0.0.4 # via fastapi annotated-types==0.7.0 @@ -27,7 +27,7 @@ anyio==4.13.0 # httpx # starlette arctic-inference==0.1.1 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in attrs==26.1.0 # via # aiohttp @@ -36,13 +36,13 @@ attrs==26.1.0 # referencing audioread==3.0.1 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # librosa blobfile==3.0.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in bm25s==0.2.13 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # mteb bounded-pool-executor==0.0.3 # via pqdm @@ -124,7 +124,7 @@ fsspec==2026.2.0 # huggingface-hub # torch gpt-oss==0.0.8 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in graphql-core==3.2.8 # via hypothesis-graphql h11==0.16.0 @@ -134,7 +134,7 @@ h11==0.16.0 harfile==0.4.0 # via schemathesis hf-transfer==0.1.9 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in hf-xet==1.4.2 # via huggingface-hub html2text==2025.4.15 @@ -218,7 +218,7 @@ jinja2==3.1.6 # lm-eval # torch jiwer==4.0.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in joblib==1.5.3 # via # librosa @@ -242,11 +242,11 @@ lazy-loader==0.5 # librosa # scikit-image librosa==0.10.2.post1 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in llvmlite==0.44.0 # via numba lm-eval==0.4.11 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in lxml==6.0.2 # via # blobfile @@ -265,10 +265,10 @@ mbstrdecoder==1.1.4 # typepy mdurl==0.1.2 # via markdown-it-py -mistral-common==1.10.0 +mistral-common==1.11.0 # via # -c requirements/common.txt - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in mkl==2025.3.0 # via # onemkl-sycl-blas @@ -278,7 +278,7 @@ mkl==2025.3.0 # onemkl-sycl-sparse # torch modelscope==1.35.3 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in more-itertools==10.8.0 # via lm-eval mpmath==1.3.0 @@ -286,7 +286,7 @@ mpmath==1.3.0 msgpack==1.1.2 # via librosa mteb==2.12.7 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in multidict==6.7.1 # via # aiohttp @@ -302,7 +302,7 @@ networkx==3.6.1 nltk==3.9.4 # via rouge-score num2words==0.5.14 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in numba==0.61.2 # via # -c requirements/xpu.txt @@ -405,12 +405,12 @@ polars-runtime-32==1.39.3 # via polars pooch==1.8.2 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # librosa portalocker==3.2.0 # via sacrebleu pqdm==0.2.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in propcache==0.4.1 # via # aiohttp @@ -451,13 +451,13 @@ pyrate-limiter==4.1.0 # via schemathesis pystemmer==3.0.0 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # mteb pytablewriter==1.2.1 # via lm-eval pytest==9.0.2 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # pytest-asyncio # pytest-cov # pytest-forked @@ -466,17 +466,17 @@ pytest==9.0.2 # pytest-timeout # schemathesis pytest-asyncio==1.3.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in pytest-cov==6.3.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in pytest-forked==1.6.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in pytest-rerunfailures==14.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in pytest-shard==0.1.2 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in pytest-timeout==2.3.1 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in python-dateutil==2.9.0.post0 # via # pandas @@ -496,7 +496,7 @@ pyyaml==6.0.3 # transformers rapidfuzz==3.12.1 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # jiwer referencing==0.37.0 # via @@ -543,7 +543,7 @@ safetensors==0.7.0 # timm # transformers schemathesis==4.14.2 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in scikit-image==0.26.0 # via albumentations scikit-learn==1.8.0 @@ -582,12 +582,12 @@ sortedcontainers==2.4.0 # via hypothesis soundfile==0.13.1 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # librosa # mistral-common soxr==0.5.0.post1 # via - # -r requirements/xpu-test.in + # -r requirements/test/xpu.in # librosa # mistral-common sqlitedict==2.1.0 @@ -612,7 +612,7 @@ tbb==2022.3.0 # mkl # torch tblib==3.1.0 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in tcmlib==1.4.1 # via # tbb @@ -638,7 +638,7 @@ tiktoken==0.12.0 # lm-eval # mistral-common timm==1.0.17 - # via -r requirements/xpu-test.in + # via -r requirements/test/xpu.in tokenizers==0.22.2 # via # -c requirements/common.txt diff --git a/requirements/xpu.txt b/requirements/xpu.txt index 3be85dcb5f4..26ba38f3efa 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -11,7 +11,7 @@ jinja2>=3.1.6 datasets # for benchmark scripts numba == 0.61.2 # Required for N-gram speculative decoding --extra-index-url=https://download.pytorch.org/whl/xpu -torch==2.11.0+xpu +torch==2.10.0+xpu torchaudio torchvision diff --git a/setup.py b/setup.py index 06c2f247448..b0cca73bb91 100644 --- a/setup.py +++ b/setup.py @@ -379,6 +379,20 @@ class cmake_build_ext(build_ext): dirs_exist_ok=True, ) + if _is_cuda(): + # copy vendored deep_gemm package from build_lib to source tree + # for editable installs + deep_gemm_build = os.path.join( + self.build_lib, "vllm", "third_party", "deep_gemm" + ) + if os.path.exists(deep_gemm_build): + print(f"Copying {deep_gemm_build} to vllm/third_party/deep_gemm") + shutil.copytree( + deep_gemm_build, + "vllm/third_party/deep_gemm", + dirs_exist_ok=True, + ) + class precompiled_build_ext(build_ext): """Disables extension building when using precompiled binaries.""" @@ -679,17 +693,29 @@ class precompiled_wheel_utils: flash_attn_regex = re.compile( r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" ) + # __init__.py and flash_attn_interface.py are source-controlled + # in vllm and should not be overwritten (matches cmake exclusions) + flash_attn_files_to_skip = { + "vllm/vllm_flash_attn/__init__.py", + "vllm/vllm_flash_attn/flash_attn_interface.py", + } triton_kernels_regex = re.compile( r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" ) flashmla_regex = re.compile( r"vllm/third_party/flashmla/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py" ) + # DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.) + deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*") file_members = list( filter(lambda x: x.filename in files_to_copy, wheel.filelist) ) file_members += list( - filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist) + filter( + lambda x: flash_attn_regex.match(x.filename) + and x.filename not in flash_attn_files_to_skip, + wheel.filelist, + ) ) file_members += list( filter( @@ -699,6 +725,9 @@ class precompiled_wheel_utils: file_members += list( filter(lambda x: flashmla_regex.match(x.filename), wheel.filelist) ) + file_members += list( + filter(lambda x: deep_gemm_regex.match(x.filename), wheel.filelist) + ) for file in file_members: print(f"[extract] {file.filename}") @@ -987,6 +1016,12 @@ if _is_cuda(): ext_modules.append( CMakeExtension(name="vllm._flashmla_extension_C", optional=True) ) + if envs.VLLM_USE_PRECOMPILED or ( + CUDA_HOME and get_nvcc_cuda_version() >= Version("12.3") + ): + # DeepGEMM requires CUDA 12.3+ (SM90/SM100) + # Optional since it won't build on unsupported architectures + ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True)) if _is_cpu(): import platform @@ -1014,6 +1049,10 @@ package_data = { "entrypoints/serve/instrumentator/static/*.js", "entrypoints/serve/instrumentator/static/*.css", "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp", + # DeepGEMM JIT include headers (vendored via cmake) + "third_party/deep_gemm/include/**/*.cuh", + "third_party/deep_gemm/include/**/*.h", + "third_party/deep_gemm/include/**/*.hpp", ] } diff --git a/tests/basic_correctness/test_cumem.py b/tests/basic_correctness/test_cumem.py index b1a16cfcaba..8d8f87f0a3c 100644 --- a/tests/basic_correctness/test_cumem.py +++ b/tests/basic_correctness/test_cumem.py @@ -13,6 +13,8 @@ from vllm.utils.mem_constants import GiB_bytes from ..utils import create_new_process_for_each_test, requires_fp8 +DEVICE_TYPE = current_platform.device_type + @create_new_process_for_each_test("fork" if not current_platform.is_rocm() else "spawn") def test_python_error(): @@ -26,13 +28,13 @@ def test_python_error(): tensors = [] with allocator.use_memory_pool(): # allocate 70% of the total memory - x = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda") + x = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE) tensors.append(x) # release the memory allocator.sleep() # allocate more memory than the total memory - y = torch.empty(alloc_bytes, dtype=torch.uint8, device="cuda") + y = torch.empty(alloc_bytes, dtype=torch.uint8, device=DEVICE_TYPE) tensors.append(y) with pytest.raises(RuntimeError): # when the allocator is woken up, it should raise an error @@ -44,17 +46,17 @@ def test_python_error(): def test_basic_cumem(): # some tensors from default memory pool shape = (1024, 1024) - x = torch.empty(shape, device="cuda") + x = torch.empty(shape, device=DEVICE_TYPE) x.zero_() # some tensors from custom memory pool allocator = CuMemAllocator.get_instance() with allocator.use_memory_pool(): # custom memory pool - y = torch.empty(shape, device="cuda") + y = torch.empty(shape, device=DEVICE_TYPE) y.zero_() y += 1 - z = torch.empty(shape, device="cuda") + z = torch.empty(shape, device=DEVICE_TYPE) z.zero_() z += 2 @@ -77,16 +79,16 @@ def test_basic_cumem(): def test_cumem_with_cudagraph(): allocator = CuMemAllocator.get_instance() with allocator.use_memory_pool(): - weight = torch.eye(1024, device="cuda") + weight = torch.eye(1024, device=DEVICE_TYPE) with allocator.use_memory_pool(tag="discard"): - cache = torch.empty(1024, 1024, device="cuda") + cache = torch.empty(1024, 1024, device=DEVICE_TYPE) def model(x): out = x @ weight cache[: out.size(0)].copy_(out) return out + 1 - x = torch.empty(128, 1024, device="cuda") + x = torch.empty(128, 1024, device=DEVICE_TYPE) # warmup model(x) diff --git a/tests/benchmarks/test_sampling_params.py b/tests/benchmarks/test_sampling_params.py new file mode 100644 index 00000000000..3bc34a84b37 --- /dev/null +++ b/tests/benchmarks/test_sampling_params.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest + +from vllm.benchmarks.datasets.utils import get_sampling_params +from vllm.tokenizers import TokenizerLike + + +class _FakeTokenizer(TokenizerLike): + """Minimal tokenizer implementing the TokenizerLike protocol + for testing get_sampling_params.""" + + def __init__(self, vocab_size: int = 1000, num_special_tokens: int = 0) -> None: + self._vocab_size = vocab_size + self._num_special_tokens = num_special_tokens + + # -- Properties required by TokenizerLike -- + + @classmethod + def from_pretrained(cls, path_or_repo_id, *a, **kw): # type: ignore[override] + return cls() + + @property + def vocab_size(self) -> int: + return self._vocab_size + + @property + def all_special_tokens(self) -> list[str]: + return [] + + @property + def all_special_ids(self) -> list[int]: + return [] + + @property + def bos_token_id(self) -> int: + return 0 + + @property + def eos_token_id(self) -> int: + return 1 + + @property + def pad_token_id(self) -> int: + return 2 + + @property + def is_fast(self) -> bool: + return False + + @property + def max_token_id(self) -> int: + return self._vocab_size - 1 + + @property + def max_chars_per_token(self) -> int: + return 4 + + @property + def truncation_side(self) -> str: + return "right" + + def num_special_tokens_to_add(self) -> int: + return self._num_special_tokens + + def __call__(self, text, text_pair=None, **kw): # type: ignore[override] + raise NotImplementedError + + def get_vocab(self) -> dict[str, int]: + return {} + + def get_added_vocab(self) -> dict[str, int]: + return {} + + def encode(self, text, **kw) -> list[int]: # type: ignore[override] + raise NotImplementedError + + def apply_chat_template(self, messages, **kw): # type: ignore[override] + raise NotImplementedError + + def convert_tokens_to_ids(self, tokens): # type: ignore[override] + raise NotImplementedError + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + raise NotImplementedError + + def decode(self, ids, skip_special_tokens: bool = False) -> str: # type: ignore[override] + raise NotImplementedError + + def convert_ids_to_tokens( # type: ignore[override] + self, ids, skip_special_tokens: bool = False + ) -> list[str]: + raise NotImplementedError + + +class TestGetSamplingParams: + """Tests for ``get_sampling_params`` in ``vllm.benchmarks.datasets.shared``.""" + + # -- helpers -- + + @staticmethod + def _tok(vocab_size: int = 1000, num_special: int = 0) -> _FakeTokenizer: + return _FakeTokenizer(vocab_size=vocab_size, num_special_tokens=num_special) + + # -- return shape / dtype -- + + def test_returns_three_arrays(self): + rng = np.random.default_rng(0) + result = get_sampling_params(rng, 5, 0.0, 100, 50, self._tok()) + assert len(result) == 3 + for arr in result: + assert isinstance(arr, np.ndarray) + + @pytest.mark.parametrize("n", [1, 10, 100]) + def test_output_length_matches_num_requests(self, n: int): + rng = np.random.default_rng(42) + input_lens, output_lens, offsets = get_sampling_params( + rng, n, 0.0, 64, 32, self._tok() + ) + assert input_lens.shape == (n,) + assert output_lens.shape == (n,) + assert offsets.shape == (n,) + + # -- fixed lengths (range_ratio = 0) -- + + def test_zero_range_ratio_gives_constant_lengths(self): + rng = np.random.default_rng(7) + input_lens, output_lens, _ = get_sampling_params( + rng, 20, 0.0, 128, 64, self._tok() + ) + assert np.all(input_lens == 128) + assert np.all(output_lens == 64) + + def test_special_tokens_subtracted_from_input_only(self): + rng = np.random.default_rng(7) + input_lens, output_lens, _ = get_sampling_params( + rng, 10, 0.0, 100, 50, self._tok(num_special=4) + ) + # real_input_len = 100 - 4 = 96, range_ratio 0 → all 96 + assert np.all(input_lens == 96) + # special tokens are not subtracted from output length + assert np.all(output_lens == 50) + + # -- range ratios -- + + def test_input_range_bounds(self): + rng = np.random.default_rng(0) + ratio = 0.5 + base = 200 + input_lens, _, _ = get_sampling_params( + rng, 500, {"input": ratio, "output": 0.0}, base, 50, self._tok() + ) + lo = int(np.floor(base * (1 - ratio))) + hi = int(np.ceil(base * (1 + ratio))) + assert np.all(input_lens >= lo) + assert np.all(input_lens <= hi) + + def test_output_range_bounds(self): + rng = np.random.default_rng(0) + ratio = 0.3 + base = 100 + _, output_lens, _ = get_sampling_params( + rng, 500, {"input": 0.0, "output": ratio}, 50, base, self._tok() + ) + lo = max(1, int(np.floor(base * (1 - ratio)))) + hi = int(np.ceil(base * (1 + ratio))) + assert np.all(output_lens >= lo) + assert np.all(output_lens <= hi) + + def test_output_low_clamped_to_one(self): + """Even with a high ratio that would push output_low to 0, + the function clamps it to 1.""" + rng = np.random.default_rng(0) + # output_len=1, ratio=0.99 → floor(1*0.01)=0, should clamp to 1 + _, output_lens, _ = get_sampling_params( + rng, 50, {"input": 0.0, "output": 0.99}, 100, 1, self._tok() + ) + assert np.all(output_lens >= 1) + + # -- offsets bounded by vocab_size -- + + @pytest.mark.parametrize("vocab", [100, 32000, 128256]) + def test_offsets_within_vocab(self, vocab: int): + rng = np.random.default_rng(0) + _, _, offsets = get_sampling_params( + rng, 200, 0.0, 64, 32, self._tok(vocab_size=vocab) + ) + assert np.all(offsets >= 0) + assert np.all(offsets < vocab) + + # -- reproducibility -- + + def test_same_seed_same_results(self): + tok = self._tok() + rr = {"input": 0.3, "output": 0.2} + a = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok) + b = get_sampling_params(np.random.default_rng(42), 50, rr, 256, 64, tok) + for arr_a, arr_b in zip(a, b): + np.testing.assert_array_equal(arr_a, arr_b) + + def test_different_seed_different_results(self): + tok = self._tok() + rr = {"input": 0.3, "output": 0.2} + a = get_sampling_params(np.random.default_rng(0), 50, rr, 256, 64, tok) + b = get_sampling_params(np.random.default_rng(1), 50, rr, 256, 64, tok) + # Extremely unlikely all three arrays match with different seeds + assert not all(np.array_equal(arr_a, arr_b) for arr_a, arr_b in zip(a, b)) + + # -- validation / error paths -- + + @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5]) + def test_invalid_input_range_ratio(self, bad_ratio: float): + rng = np.random.default_rng(0) + with pytest.raises(ValueError, match="input_range_ratio"): + get_sampling_params( + rng, 10, {"input": bad_ratio, "output": 0.0}, 100, 50, self._tok() + ) + + @pytest.mark.parametrize("bad_ratio", [-0.1, 1.0, 1.5]) + def test_invalid_output_range_ratio(self, bad_ratio: float): + rng = np.random.default_rng(0) + with pytest.raises(ValueError, match="output_range_ratio"): + get_sampling_params( + rng, 10, {"input": 0.0, "output": bad_ratio}, 100, 50, self._tok() + ) + + def test_invalid_dict_missing_keys(self): + rng = np.random.default_rng(0) + with pytest.raises(ValueError, match="input.*output"): + get_sampling_params(rng, 10, {"input": 0.1}, 100, 50, self._tok()) + + def test_input_len_zero_with_special_tokens(self): + """input_len < num_special_tokens → real_input_len = 0, which is fine + (range [0, 0]).""" + rng = np.random.default_rng(0) + input_lens, _, _ = get_sampling_params( + rng, 5, 0.0, 5, 50, self._tok(num_special=10) + ) + # real_input_len = max(0, 5 - 10) = 0 + assert np.all(input_lens == 0) + + # -- edge cases -- + + def test_single_request(self): + rng = np.random.default_rng(0) + i, o, off = get_sampling_params(rng, 1, 0.0, 100, 50, self._tok()) + assert i.shape == (1,) + assert o.shape == (1,) + assert off.shape == (1,) + + def test_large_num_requests(self): + rng = np.random.default_rng(0) + i, o, off = get_sampling_params(rng, 10_000, 0.5, 512, 128, self._tok()) + assert i.shape == (10_000,) + assert o.shape == (10_000,) + assert off.shape == (10_000,) diff --git a/tests/benchmarks/test_txt_slices_dataset.py b/tests/benchmarks/test_txt_slices_dataset.py new file mode 100644 index 00000000000..7821e9a925a --- /dev/null +++ b/tests/benchmarks/test_txt_slices_dataset.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from pathlib import Path + +import pytest +from transformers import AutoTokenizer, PreTrainedTokenizerBase + +from vllm.benchmarks.datasets import CustomDataset +from vllm.benchmarks.datasets.create_txt_slices_dataset import create_txt_slices_jsonl + + +@pytest.fixture(scope="session") +def hf_tokenizer() -> PreTrainedTokenizerBase: + # Use a small, commonly available tokenizer + return AutoTokenizer.from_pretrained("gpt2") + + +text_content = """ +Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor +incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud +exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. +Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat +nulla pariatur. Excepteur sint occaecat cupidatat non proident, +sunt in culpa qui officia deserunt mollit anim id est laborum. +""" + + +@pytest.mark.benchmark +def test_create_txt_slices_jsonl( + hf_tokenizer: PreTrainedTokenizerBase, tmp_path: Path +) -> None: + """Test that create_txt_slices_jsonl produces valid JSONL for CustomDataset.""" + txt_path = tmp_path / "input.txt" + jsonl_path = tmp_path / "input.txt.jsonl" + + txt_path.write_text(text_content) + + create_txt_slices_jsonl( + input_path=str(txt_path), + output_path=str(jsonl_path), + tokenizer_name="gpt2", + num_prompts=10, + input_len=10, + output_len=10, + ) + + # Verify the JSONL file is valid and has the expected structure + records = [json.loads(line) for line in jsonl_path.read_text().splitlines()] + + assert len(records) == 10 + for record in records: + assert "prompt" in record + assert "output_tokens" in record + assert isinstance(record["prompt"], str) + assert record["output_tokens"] == 10 + + # Verify the JSONL file can be loaded by CustomDataset + dataset = CustomDataset(dataset_path=str(jsonl_path)) + samples = dataset.sample( + tokenizer=hf_tokenizer, + num_requests=10, + output_len=10, + skip_chat_template=True, + ) + + assert len(samples) == 10 + assert all(sample.expected_output_len == 10 for sample in samples) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 7edceee9811..4fbf958d867 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -31,6 +31,7 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type FP8_DTYPE = current_platform.fp8_dtype() prompts = [ @@ -299,7 +300,7 @@ def async_tp_pass_on_test_model( ): set_random_seed(0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -324,7 +325,7 @@ def async_tp_pass_on_test_model( fuse_gemm_comms=True, ), ) - vllm_config.device_config = DeviceConfig(device=torch.device("cuda")) + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) # this is a fake model name to construct the model config # in the vllm_config, it's not really used. diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index a50d3ca8e3e..e2c461e6692 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -37,6 +37,8 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type + class TestAllReduceRMSNormModel(torch.nn.Module): def __init__( @@ -268,7 +270,7 @@ def all_reduce_fusion_pass_on_test_model( ): set_random_seed(0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -300,7 +302,7 @@ def all_reduce_fusion_pass_on_test_model( vllm_config.compilation_config.pass_config = PassConfig( fuse_allreduce_rms=True, eliminate_noops=True ) - vllm_config.device_config = DeviceConfig(device=torch.device("cuda")) + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) vllm_config.parallel_config.rank = local_rank # Setup rank for debug path # this is a fake model name to construct the model config diff --git a/tests/compile/passes/distributed/test_sequence_parallelism.py b/tests/compile/passes/distributed/test_sequence_parallelism.py index 667ef4e04fb..7b240acead5 100644 --- a/tests/compile/passes/distributed/test_sequence_parallelism.py +++ b/tests/compile/passes/distributed/test_sequence_parallelism.py @@ -35,6 +35,8 @@ from vllm.platforms import current_platform from vllm.utils.system_utils import update_environment_variables from vllm.utils.torch_utils import set_random_seed +DEVICE_TYPE = current_platform.device_type + pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA") FP8_DTYPE = current_platform.fp8_dtype() @@ -228,7 +230,7 @@ def sequence_parallelism_pass_on_test_model( ): set_random_seed(0) - device = torch.device(f"cuda:{local_rank}") + device = torch.device(f"{DEVICE_TYPE}:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) torch.set_default_dtype(dtype) @@ -258,7 +260,7 @@ def sequence_parallelism_pass_on_test_model( eliminate_noops=True, ), ) # NoOp needed for fusion - device_config = DeviceConfig(device=torch.device("cuda")) + device_config = DeviceConfig(device=torch.device(DEVICE_TYPE)) # this is a fake model name to construct the model config # in the vllm_config, it's not really used. diff --git a/tests/compile/passes/test_fusion_attn.py b/tests/compile/passes/test_fusion_attn.py index 2bbf0bda626..b776f6af98a 100644 --- a/tests/compile/passes/test_fusion_attn.py +++ b/tests/compile/passes/test_fusion_attn.py @@ -39,8 +39,9 @@ from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.registry import AttentionBackendEnum -from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode +DEVICE_TYPE = current_platform.device_type FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 @@ -53,7 +54,6 @@ class AttentionQuantPatternModel(torch.nn.Module): num_qo_heads: int, num_kv_heads: int, head_size: int, - kv_cache_dtype: torch.dtype, device: torch.device, vllm_config: VllmConfig, block_size: int, @@ -63,7 +63,6 @@ class AttentionQuantPatternModel(torch.nn.Module): self.num_qo_heads = num_qo_heads self.num_kv_heads = num_kv_heads self.head_size = head_size - self.kv_cache_dtype = kv_cache_dtype self.device = device self.vllm_config = vllm_config self.dtype = vllm_config.model_config.dtype @@ -81,13 +80,14 @@ class AttentionQuantPatternModel(torch.nn.Module): self.block_size = block_size - # Initialize attn MetadataBuilder + # Initialize attn MetadataBuilder (match Attention.get_kv_cache_spec) self.builder = self.attn.attn_backend.get_builder_cls()( kv_cache_spec=AttentionSpec( block_size=self.block_size, num_kv_heads=self.num_kv_heads, head_size=self.head_size, - dtype=self.kv_cache_dtype, + dtype=self.attn.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype), ), layer_names=[self.attn.layer_name], vllm_config=self.vllm_config, @@ -126,7 +126,7 @@ class AttentionQuantPatternModel(torch.nn.Module): # Create dummy KV cache raw_tensor = torch.zeros( 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size, - dtype=self.kv_cache_dtype, + dtype=self.attn.kv_cache_torch_dtype, device=self.device, ) raw_tensor = raw_tensor.view(kv_cache_shape) @@ -301,7 +301,7 @@ def test_attention_quant_pattern( custom_ops_list = custom_ops.split(",") if custom_ops else [] - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") torch.set_default_dtype(dtype) torch.manual_seed(42) @@ -348,7 +348,6 @@ def test_attention_quant_pattern( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_size=head_size, - kv_cache_dtype=FP8_DTYPE, device=device, vllm_config=vllm_config_unfused, block_size=block_size, @@ -376,7 +375,6 @@ def test_attention_quant_pattern( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_size=head_size, - kv_cache_dtype=FP8_DTYPE, device=device, vllm_config=vllm_config, w=model_unfused.w, diff --git a/tests/compile/passes/test_mla_attn_quant_fusion.py b/tests/compile/passes/test_mla_attn_quant_fusion.py index ce1fa642ad5..a5875a6b396 100644 --- a/tests/compile/passes/test_mla_attn_quant_fusion.py +++ b/tests/compile/passes/test_mla_attn_quant_fusion.py @@ -45,6 +45,7 @@ from vllm.v1.kv_cache_interface import MLAAttentionSpec FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 +DEVICE_TYPE = current_platform.device_type class MLAAttentionQuantPatternModel(torch.nn.Module): @@ -356,7 +357,7 @@ def test_mla_attention_quant_pattern( custom_ops_list = custom_ops.split(",") if custom_ops else [] - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") torch.set_default_dtype(dtype) torch.manual_seed(42) diff --git a/tests/compile/passes/test_noop_elimination.py b/tests/compile/passes/test_noop_elimination.py index 412e8056f9c..c31acfaf723 100644 --- a/tests/compile/passes/test_noop_elimination.py +++ b/tests/compile/passes/test_noop_elimination.py @@ -8,6 +8,9 @@ import vllm from tests.compile.backend import TestBackend from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @@ -17,7 +20,7 @@ from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConf ) @pytest.mark.parametrize("hidden_size", [64, 4096]) def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size): - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) torch.set_default_dtype(dtype) torch.manual_seed(1) @@ -88,7 +91,7 @@ def test_non_noop_slice_preserved(): Regression test for a bug where end=-1 was treated like an inferred dimension (reshape semantics) leading to incorrect elimination. """ - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) x = torch.randn(16, 16) class SliceModel(torch.nn.Module): diff --git a/tests/compile/passes/test_rope_kvcache_fusion.py b/tests/compile/passes/test_rope_kvcache_fusion.py index eea21c9179b..bab70c12a89 100644 --- a/tests/compile/passes/test_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rope_kvcache_fusion.py @@ -28,6 +28,7 @@ from vllm.forward_context import get_forward_context, set_forward_context from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding from vllm.platforms import current_platform +from vllm.utils.torch_utils import _encode_layer_name from vllm.v1.attention.backend import ( AttentionBackend, CommonAttentionMetadata, @@ -170,7 +171,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module): k = k.view(-1, self.num_kv_heads, self.head_size) v = v.view(-1, self.num_kv_heads, self.head_size) kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update( - k, v, self.layer_name + k, v, _encode_layer_name(self.layer_name) ) return q, k, v, kv_cache_dummy_dep diff --git a/tests/compile/passes/test_scatter_split_replace.py b/tests/compile/passes/test_scatter_split_replace.py index 65996089640..e85fd9f9efc 100644 --- a/tests/compile/passes/test_scatter_split_replace.py +++ b/tests/compile/passes/test_scatter_split_replace.py @@ -13,6 +13,9 @@ from vllm.compilation.passes.utility.scatter_split_replace import ( from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass from vllm.config import CompilationConfig, CompilationMode, VllmConfig from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type class ScatterSplitReplacementModel(nn.Module): @@ -61,7 +64,7 @@ class ScatterSplitReplacementModel(nn.Module): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_scatter_split_replace(dtype): - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) torch.set_default_dtype(dtype) torch.manual_seed(0) diff --git a/tests/compile/passes/test_split_coalescing.py b/tests/compile/passes/test_split_coalescing.py index a217a4af9f2..ab7a0be1a21 100644 --- a/tests/compile/passes/test_split_coalescing.py +++ b/tests/compile/passes/test_split_coalescing.py @@ -8,6 +8,9 @@ import vllm from tests.compile.backend import TestBackend from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type class SplitCoalescingModel(torch.nn.Module): @@ -28,7 +31,7 @@ class SplitCoalescingModel(torch.nn.Module): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_split_coalescing(dtype): - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) torch.set_default_dtype(dtype) torch.manual_seed(0) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index c0f0dcca8d7..12518b4cfa2 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -31,6 +31,8 @@ from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher # This import automatically registers `torch.ops.silly.attention` from . import silly_attention # noqa: F401 +DEVICE_TYPE = current_platform.device_type + def test_version(): # Test the version comparison logic using the private function @@ -456,7 +458,7 @@ def test_cached_compilation_config(default_vllm_config): from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape dtype = torch.bfloat16 - device = torch.device("cuda:0") + device = torch.device(f"{DEVICE_TYPE}:0") batch_size, num_qo_heads, head_size = 8, 16, 128 # access and cache default compilation config @@ -478,7 +480,7 @@ def test_cached_compilation_config(default_vllm_config): query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR) query_quant = torch.compile(query_quant) - _q_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda") + _q_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE_TYPE) query = torch.randn( batch_size, num_qo_heads * head_size, dtype=dtype, device=device ) diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index bbd62237c5e..1775b2c9deb 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards( torch.randn(1, 10).cuda(), is_01_specialization=True, ) + + +@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10") +def test_piecewise_backend_empty_sym_shape_indices(): + """Test that PiecewiseBackend handles empty sym_shape_indices correctly. + + When all inputs have static shapes (no torch.SymInt), sym_shape_indices + will be empty. The fix in PiecewiseBackend.__call__ handles this case + by using the first compiled range_entry. + """ + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() + + # Use small max_model_len and max_num_batched_tokens to encourage + # static shape compilation with empty sym_shape_indices + llm = LLM( + model="Qwen/Qwen3-0.6B", + max_model_len=512, + max_num_batched_tokens=1, + compilation_config={ + "mode": CompilationMode.VLLM_COMPILE, + "dynamic_shapes_config": { + "type": DynamicShapesType.BACKED.value, + }, + }, + ) + + sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10) + + # Generate with static shape inputs + output = llm.generate("Hello, my name is", sampling_params=sampling_params) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output" + + # Generate again to verify compilation works with empty sym_shape_indices + output = llm.generate("The capital of France is", sampling_params=sampling_params) + result = output[0].outputs[0].text + assert len(result) > 0, "Should generate non-empty output on second run" + + del llm + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py index 0b490e97f3f..4cb199b5897 100644 --- a/tests/compile/test_graph_partition.py +++ b/tests/compile/test_graph_partition.py @@ -9,12 +9,19 @@ import torch._dynamo import torch.fx as fx from torch.fx.experimental.proxy_tensor import make_fx -from vllm.compilation.backends import _is_empty_allocation_node, split_graph +from vllm.compilation.backends import ( + _decompose_size_nodes, + _is_empty_allocation_node, + split_graph, +) from vllm.compilation.passes.fx_utils import find_op_nodes +from vllm.platforms import current_platform # This import automatically registers `torch.ops.silly.attention` from . import silly_attention # noqa: F401 +DEVICE_TYPE = current_platform.device_type + def test_getitem_moved_to_producer_subgraph(): """ @@ -147,7 +154,7 @@ def test_consecutive_ops_in_split(): final_result = torch.sigmoid(attn_inout) return final_result - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) # Create the traced FX graph for the model x = torch.randn(8, 4) @@ -325,7 +332,7 @@ def test_builtin_empty_only_partition_is_merged(): "Expected two builtin empty_like nodes in merged non-splitting subgraph" ) - x = torch.randn(2, 3, device="cuda") + x = torch.randn(2, 3, device=DEVICE_TYPE) output_original = gm(x) output_split = split_gm(x) assert torch.allclose(output_original, output_split), "Output mismatch after split" @@ -622,3 +629,73 @@ def test_sym_size_metadata_propagated(): else: example_inputs.append(int(ev)) standalone_compile(submod, example_inputs, dynamic_shapes="from_example_inputs") + + +def test_decompose_size_with_getitem_user(): + """ + Regression test: _decompose_size_nodes must handle getitem users of size() + correctly. + + When a graph contains x.shape[i], it can appear as: + + %size = call_method[target="size"](args = (%x,)) + %getitem = call_function[target=operator.getitem](args = (%size, 1)) + + The old code spliced *all* per-dim values into every user's args + unconditionally, turning the 2-arg getitem into a malformed 3-arg node: + + %getitem(args = (%sym_size_int, 5120, 1)) # TypeError at runtime + + The fix detects getitem users and replaces them with dims[idx] directly. + """ + # Build a graph manually to guarantee the size() + getitem pattern. + # + # Graph: + # %x = placeholder + # %size = x.size() + # %dim1 = getitem(%size, 1) <-- the getitem branch we're testing + # %relu = relu(%x) + # %view = view(%relu, -1, %dim1) + # return %view + graph = fx.Graph() + x = graph.placeholder("x") + size_node = graph.call_method("size", args=(x,)) + getitem_node = graph.call_function(operator.getitem, args=(size_node, 1)) + relu_node = graph.call_function(torch.ops.aten.relu.default, args=(x,)) + view_node = graph.call_function( + torch.ops.aten.view.default, args=(relu_node, [-1, getitem_node]) + ) + graph.output(view_node) + + # Attach example_value metadata so _decompose_size_nodes can inspect dims. + # dim 0 is dynamic (SymInt), dim 1 is static (8). + from torch._dynamo.source import LocalSource + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + shape_env = ShapeEnv() + src = LocalSource("batch_size") + sym_batch = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4) + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_x = torch.empty_strided((sym_batch, 8), (8, 1)) + x.meta["example_value"] = fake_x + + gm = fx.GraphModule(torch.nn.Module(), graph) + + # Run decomposition — this would produce a 3-arg getitem without the fix + _decompose_size_nodes(gm) + + # Verify no size() nodes remain + remaining_size_nodes = list(gm.graph.find_nodes(op="call_method", target="size")) + assert len(remaining_size_nodes) == 0, ( + f"size() nodes should be fully decomposed, found {len(remaining_size_nodes)}" + ) + + # Verify no malformed getitem nodes (3+ args) + for node in gm.graph.nodes: + if node.op == "call_function" and node.target is operator.getitem: + assert len(node.args) == 2, ( + f"getitem node '{node.name}' has {len(node.args)} args " + f"(expected 2): {node.args}" + ) diff --git a/tests/compile/test_rotary_embedding_compile.py b/tests/compile/test_rotary_embedding_compile.py index 76f5382534e..69a4cc05084 100644 --- a/tests/compile/test_rotary_embedding_compile.py +++ b/tests/compile/test_rotary_embedding_compile.py @@ -16,6 +16,8 @@ from vllm.config.compilation import CompilationMode, CUDAGraphMode from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + @support_torch_compile class RotaryEmbeddingCompileModule(torch.nn.Module): @@ -45,7 +47,7 @@ def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch): monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1") monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0") - device = "cuda" + device = DEVICE_TYPE positions = torch.arange(16, device=device) query = torch.randn(16, 32, device=device, dtype=torch.bfloat16) key = torch.randn(16, 32, device=device, dtype=torch.bfloat16) diff --git a/tests/compile/test_structured_logging.py b/tests/compile/test_structured_logging.py index 7813b7429b1..10b0ed139cf 100644 --- a/tests/compile/test_structured_logging.py +++ b/tests/compile/test_structured_logging.py @@ -17,8 +17,10 @@ from vllm.config.compilation import ( ) from vllm.config.scheduler import SchedulerConfig from vllm.forward_context import set_forward_context +from vllm.platforms import current_platform MLP_SIZE = 64 +DEVICE_TYPE = current_platform.device_type @support_torch_compile @@ -71,7 +73,7 @@ class TraceStructuredCapture: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache): """Test that all expected vLLM artifacts are logged during compilation.""" - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) capture = TraceStructuredCapture() diff --git a/tests/distributed/test_kv_cache_events.py b/tests/distributed/test_kv_cache_events.py new file mode 100644 index 00000000000..57d1c9b546a --- /dev/null +++ b/tests/distributed/test_kv_cache_events.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.distributed.kv_events import BlockRemoved, BlockStored + +# Minimal ExternalBlockHash for testing (bytes are a valid ExternalBlockHash). +_FAKE_HASH: bytes = b"\xab" * 32 + + +def _make_block_stored(group_idx: int | None = None) -> BlockStored: + return BlockStored( + block_hashes=[_FAKE_HASH], + parent_block_hash=None, + token_ids=[1, 2, 3, 4], + block_size=4, + lora_id=None, + medium="GPU", + lora_name=None, + group_idx=group_idx, + ) + + +def _make_block_removed(group_idx: int | None = None) -> BlockRemoved: + return BlockRemoved( + block_hashes=[_FAKE_HASH], + medium="GPU", + group_idx=group_idx, + ) + + +def test_block_stored_default_group_idx_is_none(): + """group_idx defaults to None when not provided.""" + event = _make_block_stored() + assert event.group_idx is None + + +def test_block_removed_default_group_idx_is_none(): + """group_idx defaults to None when not provided.""" + event = _make_block_removed() + assert event.group_idx is None + + +@pytest.mark.parametrize("group_idx", [1, 2, 3]) +def test_block_stored_hash_differs_by_group_idx(group_idx: int): + """BlockStored events that differ only in group_idx must hash differently.""" + other_group_idx = group_idx + 1 + event_a = _make_block_stored(group_idx=group_idx) + event_b = _make_block_stored(group_idx=other_group_idx) + assert hash(event_a) != hash(event_b) + + +def test_block_stored_hash_same_for_equal_group_idx(): + """Two BlockStored events with identical fields produce the same hash.""" + event_a = _make_block_stored(group_idx=1) + event_b = _make_block_stored(group_idx=1) + assert hash(event_a) == hash(event_b) + + +@pytest.mark.parametrize("group_idx", [1, 2, 3]) +def test_block_removed_hash_differs_by_group_idx(group_idx: int): + """BlockRemoved events that differ only in group_idx must hash differently.""" + other_group_idx = group_idx + 1 + event_a = _make_block_removed(group_idx=group_idx) + event_b = _make_block_removed(group_idx=other_group_idx) + assert hash(event_a) != hash(event_b) + + +def test_block_removed_hash_same_for_equal_group_idx(): + """Two BlockRemoved events with identical fields produce the same hash.""" + event_a = _make_block_removed(group_idx=1) + event_b = _make_block_removed(group_idx=1) + assert hash(event_a) == hash(event_b) diff --git a/tests/entrypoints/openai/chat_completion/test_chat_error.py b/tests/entrypoints/openai/chat_completion/test_chat_error.py index 46070e4810b..f1fb7c7518b 100644 --- a/tests/entrypoints/openai/chat_completion/test_chat_error.py +++ b/tests/entrypoints/openai/chat_completion/test_chat_error.py @@ -87,7 +87,6 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, @@ -123,7 +122,6 @@ async def test_chat_error_non_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -173,7 +171,6 @@ async def test_chat_error_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index cb356e0e198..39d59d28f85 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -567,7 +567,6 @@ def _build_serving_render( return OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=model_registry, request_logger=None, chat_template=CHAT_TEMPLATE, @@ -599,7 +598,6 @@ def _build_serving_chat(engine: AsyncLLM) -> OpenAIServingChat: class MockEngine: model_config: MockModelConfig = field(default_factory=MockModelConfig) input_processor: MagicMock = field(default_factory=MagicMock) - io_processor: MagicMock = field(default_factory=MagicMock) renderer: MagicMock = field(default_factory=MagicMock) @@ -632,7 +630,6 @@ async def test_serving_chat_returns_correct_model_name(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -662,7 +659,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -693,7 +689,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Initialize the serving chat @@ -737,7 +732,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -779,7 +773,6 @@ async def test_serving_chat_should_set_correct_max_tokens(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Initialize the serving chat @@ -823,7 +816,6 @@ async def test_serving_chat_mistral_token_ids_prompt_is_validated(): mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_tokenizer = MagicMock(spec=MistralTokenizer) mock_renderer = MistralRenderer( @@ -863,7 +855,6 @@ async def test_serving_chat_mistral_token_ids_prompt_too_long_is_rejected(): mock_engine.errored = False mock_engine.model_config = MockModelConfig(skip_tokenizer_init=True) mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_tokenizer = MagicMock(spec=MistralTokenizer) mock_renderer = MistralRenderer( @@ -906,7 +897,6 @@ async def test_serving_chat_could_load_correct_generation_config(): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Initialize the serving chat @@ -952,7 +942,6 @@ async def test_serving_chat_did_set_correct_cache_salt(model_type): mock_engine.errored = False mock_engine.model_config = mock_model_config mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_chat = _build_serving_chat(mock_engine) @@ -1003,7 +992,6 @@ async def test_serving_chat_data_parallel_rank_extraction(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) # Mock the generate method to return an async generator @@ -1095,7 +1083,6 @@ class TestServingChatWithHarmony: mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) return mock_engine @@ -1732,7 +1719,6 @@ async def test_tool_choice_validation_without_parser(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) models = OpenAIServingModels( @@ -1802,7 +1788,6 @@ async def test_streaming_n_gt1_independent_tool_parsers(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) models = OpenAIServingModels( diff --git a/tests/v1/entrypoints/openai/test_thinking_token_budget.py b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py similarity index 72% rename from tests/v1/entrypoints/openai/test_thinking_token_budget.py rename to tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py index b3f6d53ab1d..d2db50082a5 100644 --- a/tests/v1/entrypoints/openai/test_thinking_token_budget.py +++ b/tests/entrypoints/openai/chat_completion/test_thinking_token_budget.py @@ -24,6 +24,24 @@ def server(): "--max-model-len", "2048", "--enforce-eager", + "--gpu-memory-utilization", + "0.4", + "--no-async-scheduling", + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest.fixture(scope="module") +def server_with_auto_reasoning_config(): + args = [ + "--reasoning-parser", + "qwen3", + "--max-model-len", + "2048", + "--enforce-eager", + "--gpu-memory-utilization", + "0.4", "--no-async-scheduling", ] with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: @@ -31,12 +49,18 @@ def server(): @pytest_asyncio.fixture -async def client(server): - async with server.get_async_client() as async_client: +async def client(request, server, server_with_auto_reasoning_config): + server_map = { + "default": server, + "auto_config": server_with_auto_reasoning_config, + } + target_server = server_map[request.param] + async with target_server.get_async_client() as async_client: yield async_client @pytest.mark.asyncio +@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True) async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): """Test that mixed requests (some with thinking_token_budget, some without) complete successfully without errors.""" @@ -61,6 +85,7 @@ async def test_thinking_token_budget_mixed_requests(client: openai.AsyncOpenAI): @pytest.mark.asyncio +@pytest.mark.parametrize("client", ["default", "auto_config"], indirect=True) async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI): """Test that thinking_token_budget limits the number of reasoning tokens. @@ -82,6 +107,6 @@ async def test_thinking_token_budget_limits_reasoning(client: openai.AsyncOpenAI reasoning_token_count += 1 assert reasoning_token_count == THINK_BUDGET, ( - f"reasoning tokens ({reasoning_token_count}) != " + f"reasoning tokens ({reasoning_token_count}) exceeded " f"thinking_token_budget ({THINK_BUDGET})" ) diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 46eb02e3c59..3349f4126bc 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -79,7 +79,6 @@ def _build_serving_completion(engine: AsyncLLM) -> OpenAIServingCompletion: serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, @@ -107,7 +106,6 @@ async def test_completion_error_non_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_completion = _build_serving_completion(mock_engine) @@ -157,7 +155,6 @@ async def test_completion_error_stream(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) serving_completion = _build_serving_completion(mock_engine) diff --git a/tests/entrypoints/openai/completion/test_lora_resolvers.py b/tests/entrypoints/openai/completion/test_lora_resolvers.py index 8d5283de5cf..6a0bec92516 100644 --- a/tests/entrypoints/openai/completion/test_lora_resolvers.py +++ b/tests/entrypoints/openai/completion/test_lora_resolvers.py @@ -137,7 +137,6 @@ def mock_serving_setup(): mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() mock_engine.renderer = _build_renderer(mock_engine.model_config) models = OpenAIServingModels( @@ -148,7 +147,6 @@ def mock_serving_setup(): serving_render = OpenAIServingRender( model_config=mock_engine.model_config, renderer=mock_engine.renderer, - io_processor=mock_engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, diff --git a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py b/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py index a260027af0f..632c4bcc90a 100644 --- a/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py +++ b/tests/entrypoints/openai/generative_scoring/test_generative_scoring.py @@ -77,7 +77,6 @@ def _create_mock_engine(): mock_engine.errored = False mock_engine.model_config = MockModelConfig() mock_engine.input_processor = MagicMock() - mock_engine.io_processor = MagicMock() # renderer is accessed by OpenAIServing.__init__ and serving.py mock_renderer = MagicMock() diff --git a/tests/entrypoints/openai/responses/test_function_call.py b/tests/entrypoints/openai/responses/test_function_call.py index bacb084c7eb..f198e1cfa70 100644 --- a/tests/entrypoints/openai/responses/test_function_call.py +++ b/tests/entrypoints/openai/responses/test_function_call.py @@ -249,40 +249,74 @@ async def test_function_calling_with_streaming_expected_arguments( "additionalProperties": False, }, "strict": True, - } + }, + { + "type": "function", + "name": "get_time", + "description": "Get current local time for provided location.", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + }, + "required": ["location"], + "additionalProperties": False, + }, + "strict": True, + }, ] stream_response = await client.responses.create( model=model_name, - input="Can you tell me what the current weather is in Berlin?", + input=( + "Use tools only. Call get_weather for Berlin and get_time for Tokyo. " + "Do not answer directly." + ), tools=tools, stream=True, ) - tool_call_item = None - completed_event = None + tool_call_items = {} + arguments_done_events = {} + completed_events = {} async for event in stream_response: if ( event.type == "response.output_item.added" and event.item.type == "function_call" ): - tool_call_item = event.item - elif event.type == "response.function_call_arguments.delta" and tool_call_item: + tool_call_items[event.output_index] = event.item + elif event.type == "response.function_call_arguments.delta": + tool_call_item = tool_call_items[event.output_index] tool_call_item.arguments += event.delta + elif event.type == "response.function_call_arguments.done": + arguments_done_events[event.output_index] = event elif ( event.type == "response.output_item.done" and event.item.type == "function_call" ): - completed_event = event - assert tool_call_item is not None - assert tool_call_item.type == "function_call" - assert tool_call_item.name == "get_weather" - assert completed_event is not None - assert tool_call_item.arguments == completed_event.item.arguments - assert tool_call_item.name == completed_event.item.name - args = json.loads(tool_call_item.arguments) - assert "location" in args - assert args["location"] is not None + completed_events[event.output_index] = event + assert len(tool_call_items) >= 2 + assert len(arguments_done_events) >= 2 + assert len(completed_events) >= 2 + + tool_calls_by_name = { + event.item.name: ( + tool_call_items[output_index], + arguments_done_events[output_index], + event.item, + ) + for output_index, event in completed_events.items() + } + assert {"get_weather", "get_time"}.issubset(tool_calls_by_name) + for added_item, arguments_done_event, completed_item in tool_calls_by_name.values(): + assert added_item.type == "function_call" + assert added_item.arguments == arguments_done_event.arguments + assert added_item.arguments == completed_item.arguments + assert added_item.name == arguments_done_event.name + assert added_item.name == completed_item.name + args = json.loads(added_item.arguments) + assert "location" in args + assert args["location"] is not None @pytest.mark.asyncio diff --git a/tests/entrypoints/openai/responses/test_harmony.py b/tests/entrypoints/openai/responses/test_harmony.py index 74f3360df45..88dd2d38457 100644 --- a/tests/entrypoints/openai/responses/test_harmony.py +++ b/tests/entrypoints/openai/responses/test_harmony.py @@ -999,17 +999,21 @@ async def test_mcp_tool_multi_turn(client: OpenAI, model_name: str, server): (msg.get("recipient") or "").startswith("python") for msg in response1.output_messages ) + parsed_output_messages = [ + Message.from_dict(msg) for msg in response1.output_messages + ] tool_response_found = any( - msg.get("author", {}).get("role") == "tool" - and (msg.get("author", {}).get("name") or "").startswith("python") - for msg in response1.output_messages + (msg.author.role == "tool" and (msg.author.name or "").startswith("python")) + for msg in parsed_output_messages ) assert tool_call_found, "MCP tool call not found in output_messages" assert tool_response_found, "MCP tool response not found in output_messages" # No developer messages expected for elevated tools developer_msgs = [ - msg for msg in response1.input_messages if msg["author"]["role"] == "developer" + msg + for msg in (Message.from_dict(raw) for raw in response1.input_messages) + if msg.author.role == "developer" ] assert len(developer_msgs) == 0, "No developer message expected for elevated tools" @@ -1119,12 +1123,10 @@ async def test_function_call_with_previous_input_messages( num_system = 0 num_developer = 0 num_tool = 0 - for msg_dict in response_2.input_messages: - # input_messages use {"author": {"role": "..."}} format, - # not the top-level {"role": "..."} that Message.from_dict - # expects. - author = msg_dict.get("author", {}) - role = author.get("role") if isinstance(author, dict) else None + for message in ( + Message.from_dict(msg_dict) for msg_dict in response_2.input_messages + ): + role = message.author.role if role == "system": num_system += 1 elif role == "developer": @@ -1183,12 +1185,8 @@ async def test_system_prompt_override_no_duplication(client: OpenAI, model_name: assert response.output_text is not None num_system = 0 - for msg in response.input_messages: - # input_messages use {"author": {"role": "system"}} format, - # not the top-level {"role": "system"} that Message.from_dict expects. - author = msg.get("author", {}) - role = author.get("role") if isinstance(author, dict) else None - if role == "system": + for message in (Message.from_dict(msg) for msg in response.input_messages): + if message.author.role == "system": num_system += 1 assert num_system == 1, f"Expected 1 system message, got {num_system}" diff --git a/tests/entrypoints/openai/responses/test_mcp_tools.py b/tests/entrypoints/openai/responses/test_mcp_tools.py index 763e2b20855..330d4b9e4bc 100644 --- a/tests/entrypoints/openai/responses/test_mcp_tools.py +++ b/tests/entrypoints/openai/responses/test_mcp_tools.py @@ -7,7 +7,7 @@ from __future__ import annotations import pytest import pytest_asyncio from openai import OpenAI -from openai_harmony import ToolDescription, ToolNamespaceConfig +from openai_harmony import Message, ToolDescription, ToolNamespaceConfig from tests.utils import RemoteOpenAIServer from vllm.entrypoints.mcp.tool_server import MCPToolServer @@ -173,10 +173,10 @@ class TestMCPEnabled: if recipient and recipient.startswith("python"): tool_call_found = True assert message.get("channel") == "commentary" - author = message.get("author", {}) - if author.get("role") == "tool" and (author.get("name") or "").startswith( - "python" - ): + parsed_message = Message.from_dict(message) + if parsed_message.author.role == "tool" and ( + parsed_message.author.name or "" + ).startswith("python"): tool_response_found = True assert message.get("channel") == "commentary" @@ -188,7 +188,7 @@ class TestMCPEnabled: assert tool_response_found, "No Python tool response found" for message in response.input_messages: - assert message.get("author", {}).get("role") != "developer" + assert Message.from_dict(message).author.role != "developer" @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) diff --git a/tests/entrypoints/openai/responses/test_responses_utils.py b/tests/entrypoints/openai/responses/test_responses_utils.py index 3a4476984d3..9dab4186d40 100644 --- a/tests/entrypoints/openai/responses/test_responses_utils.py +++ b/tests/entrypoints/openai/responses/test_responses_utils.py @@ -22,6 +22,7 @@ from vllm.entrypoints.openai.responses.utils import ( _construct_single_message_from_response_item, _maybe_combine_reasoning_and_tool_call, construct_chat_messages_with_tool_call, + construct_input_messages, convert_tool_responses_to_completions_format, should_continue_final_message, ) @@ -738,3 +739,71 @@ class TestMaybeCombineReasoningAndToolCall: result = _maybe_combine_reasoning_and_tool_call(item, messages) assert result is None + + +class TestConstructInputMessagesInstructionsLeak: + """Regression tests for #37697: instructions from a prior response + should NOT leak through previous_response_id.""" + + def test_old_instructions_stripped_from_prev_msg(self): + """System message in prev_msg must be dropped so the new request's + instructions are the only system message in the conversation.""" + prev = [ + {"role": "system", "content": "old instructions"}, + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + ] + msgs = construct_input_messages( + request_instructions="new instructions", + request_input="What is 3+3?", + prev_msg=prev, + ) + system_msgs = [m for m in msgs if m.get("role") == "system"] + assert len(system_msgs) == 1 + assert system_msgs[0]["content"] == "new instructions" + + def test_no_instructions_in_new_request(self): + """If the new request has no instructions, old ones should still + be stripped -- they must not carry over.""" + prev = [ + {"role": "system", "content": "old instructions"}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + msgs = construct_input_messages( + request_instructions=None, + request_input="What is 3+3?", + prev_msg=prev, + ) + system_msgs = [m for m in msgs if m.get("role") == "system"] + assert len(system_msgs) == 0 + + def test_non_system_messages_preserved(self): + """User/assistant messages from prev_msg must remain intact.""" + prev = [ + {"role": "system", "content": "old instructions"}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"}, + ] + msgs = construct_input_messages( + request_instructions="new instructions", + request_input="Follow up", + prev_msg=prev, + ) + roles = [m["role"] for m in msgs] + assert roles == ["system", "user", "assistant", "user"] + assert msgs[0]["content"] == "new instructions" + assert msgs[1]["content"] == "Hi" + assert msgs[2]["content"] == "Hello" + assert msgs[3]["content"] == "Follow up" + + def test_no_prev_msg(self): + """Baseline: when there's no prev_msg, instructions work normally.""" + msgs = construct_input_messages( + request_instructions="be helpful", + request_input="hello", + prev_msg=None, + ) + assert len(msgs) == 2 + assert msgs[0] == {"role": "system", "content": "be helpful"} + assert msgs[1] == {"role": "user", "content": "hello"} diff --git a/tests/entrypoints/openai/responses/test_serving_responses.py b/tests/entrypoints/openai/responses/test_serving_responses.py index 39429cb9bf9..90c3183939a 100644 --- a/tests/entrypoints/openai/responses/test_serving_responses.py +++ b/tests/entrypoints/openai/responses/test_serving_responses.py @@ -11,8 +11,12 @@ from openai.types.responses import ( ResponseReasoningItem, ResponseReasoningTextDeltaEvent, ResponseReasoningTextDoneEvent, + ResponseTextConfig, ResponseTextDeltaEvent, ) +from openai.types.responses.response_format_text_json_schema_config import ( + ResponseFormatTextJSONSchemaConfig, +) from openai.types.responses.tool import ( CodeInterpreterContainerCodeInterpreterToolAuto, LocalShell, @@ -23,12 +27,20 @@ from openai.types.responses.tool import ( import vllm.envs as envs from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.engine.protocol import ( + DeltaFunctionCall, DeltaMessage, + DeltaToolCall, ErrorResponse, RequestResponseMetadata, ) from vllm.entrypoints.openai.responses.context import ConversationContext, SimpleContext -from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.openai.responses.protocol import ( + ResponseCreatedEvent, + ResponseRawMessageAndToken, + ResponsesRequest, + ResponsesResponse, + serialize_message, +) from vllm.entrypoints.openai.responses.serving import ( OpenAIServingResponses, _extract_allowed_tools_from_mcp_requests, @@ -73,6 +85,16 @@ class MockConversationContext(ConversationContext): pass +def test_serialize_message_pydantic_model_returns_dict() -> None: + msg = ResponseRawMessageAndToken(message="hello", tokens=[1, 2, 3]) + + serialized = serialize_message(msg) + + assert isinstance(serialized, dict) + assert serialized["type"] == "raw_message_tokens" + assert serialized["message"] == "hello" + + @pytest.fixture def mock_serving_responses(): """Create a mock OpenAIServingResponses instance""" @@ -132,6 +154,56 @@ def test_extract_tool_types(monkeypatch: pytest.MonkeyPatch) -> None: } +@pytest.mark.skip_global_cleanup +def test_response_created_event_uses_public_json_schema_alias() -> None: + schema = { + "type": "object", + "properties": { + "event_name": {"type": "string"}, + "date": {"type": "string"}, + "participants": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["event_name", "date", "participants"], + "additionalProperties": False, + } + text = ResponseTextConfig() + text.format = ResponseFormatTextJSONSchemaConfig( + type="json_schema", + name="calendar_event", + schema=schema, + description="A calendar event.", + strict=True, + ) + request = ResponsesRequest( + model="test-model", + input="Alice and Bob are going to a science fair on Friday.", + text=text, + ) + sampling_params = request.to_sampling_params(default_max_tokens=64) + initial_response = ResponsesResponse.from_request( + request=request, + sampling_params=sampling_params, + model_name="test-model", + created_time=0, + output=[], + status="in_progress", + usage=None, + ).model_dump(mode="json", by_alias=True) + + fmt = initial_response["text"]["format"] + assert fmt["schema"] == schema + assert "schema_" not in fmt + + event = ResponseCreatedEvent( + type="response.created", + sequence_number=0, + response=initial_response, + ) + assert event.response.text is not None + assert event.response.text.format is not None + assert event.response.text.format.model_dump(by_alias=True)["schema"] == schema + + class TestInitializeToolSessions: """Test class for _initialize_tool_sessions method""" @@ -148,7 +220,6 @@ class TestInitializeToolSessions: engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() models = MagicMock() @@ -237,7 +308,6 @@ class TestValidateGeneratorInput: engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() models = MagicMock() @@ -299,7 +369,6 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch): model_config.get_diff_sampling_param.return_value = {} engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() tokenizer = FakeTokenizer() @@ -602,7 +671,6 @@ def _make_serving_instance_with_reasoning(): model_config.get_diff_sampling_param.return_value = {} engine_client.model_config = model_config engine_client.input_processor = MagicMock() - engine_client.io_processor = MagicMock() engine_client.renderer = MagicMock() models = MagicMock() @@ -862,3 +930,197 @@ class TestStreamingReasoningToContentTransition: ] assert len(item_done_events) == 1 assert isinstance(item_done_events[0].item, ResponseReasoningItem) + + +class TestAutoToolStreaming: + @staticmethod + async def _collect_events(delta_sequence: list[DeltaMessage]): + serving = _make_serving_instance_with_reasoning() + _mock_parser_with_reasoning(serving, delta_sequence) + + contexts = [ + _make_simple_context_with_output("chunk", [i]) + for i in range(len(delta_sequence)) + ] + + async def result_generator(): + for ctx in contexts: + yield ctx + + request = ResponsesRequest( + input="hi", + tools=[ + { + "type": "function", + "name": "get_weather", + "description": "Get weather.", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + "additionalProperties": False, + }, + } + ], + tool_choice="auto", + stream=True, + ) + sampling_params = SamplingParams(max_tokens=64) + metadata = RequestResponseMetadata(request_id="req") + _identity_increment._counter = 0 # type: ignore + + events = [] + async for event in serving._process_simple_streaming_events( + request=request, + sampling_params=sampling_params, + result_generator=result_generator(), + context=SimpleContext(), + model_name="test-model", + tokenizer=MagicMock(), + request_metadata=metadata, + created_time=0, + _increment_sequence_number_and_return=_identity_increment, + ): + events.append(event) + return events + + @pytest.mark.skip_global_cleanup + @pytest.mark.asyncio + async def test_auto_multi_tool_streaming_opens_one_item_per_tool(self, monkeypatch): + monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) + + delta_sequence = [ + DeltaMessage( + tool_calls=[ + DeltaToolCall( + id="call_vienna", + type="function", + index=0, + function=DeltaFunctionCall( + name="get_weather", + arguments="", + ), + ) + ] + ), + DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=0, + function=DeltaFunctionCall( + arguments='{"location":"Vienna"}', + ), + ) + ] + ), + DeltaMessage( + tool_calls=[ + DeltaToolCall( + id="call_berlin", + type="function", + index=1, + function=DeltaFunctionCall( + name="get_weather", + arguments='{"location":"Berlin"}', + ), + ) + ] + ), + ] + events = await self._collect_events(delta_sequence) + + function_items = [ + event + for event in events + if event.type == "response.output_item.added" + and getattr(event.item, "type", None) == "function_call" + ] + assert len(function_items) == 2 + assert [event.item.name for event in function_items] == [ + "get_weather", + "get_weather", + ] + assert [event.output_index for event in function_items] == [0, 1] + + argument_deltas = [ + event.delta + for event in events + if event.type == "response.function_call_arguments.delta" + ] + assert argument_deltas == [ + '{"location":"Vienna"}', + '{"location":"Berlin"}', + ] + + argument_done = [ + event + for event in events + if event.type == "response.function_call_arguments.done" + ] + assert [event.arguments for event in argument_done] == [ + '{"location":"Vienna"}', + '{"location":"Berlin"}', + ] + assert [event.output_index for event in argument_done] == [0, 1] + + function_done = [ + event + for event in events + if event.type == "response.output_item.done" + and getattr(event.item, "type", None) == "function_call" + ] + assert [event.item.arguments for event in function_done] == [ + '{"location":"Vienna"}', + '{"location":"Berlin"}', + ] + assert [event.output_index for event in function_done] == [0, 1] + + @pytest.mark.skip_global_cleanup + @pytest.mark.asyncio + async def test_auto_tool_choice_first_delta_tool_call_does_not_duplicate_item( + self, monkeypatch + ): + monkeypatch.setattr(envs, "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", False) + + delta_sequence = [ + DeltaMessage( + tool_calls=[ + DeltaToolCall( + id="call_test", + type="function", + index=0, + function=DeltaFunctionCall( + name="get_weather", + arguments="", + ), + ) + ] + ), + DeltaMessage( + tool_calls=[ + DeltaToolCall( + index=0, + function=DeltaFunctionCall( + arguments='{"location":"Berlin"}', + ), + ) + ] + ), + ] + events = await self._collect_events(delta_sequence) + + function_items = [ + event + for event in events + if event.type == "response.output_item.added" + and getattr(event.item, "type", None) == "function_call" + ] + assert len(function_items) == 1 + assert function_items[0].item.name == "get_weather" + + argument_deltas = [ + event.delta + for event in events + if event.type == "response.function_call_arguments.delta" + ] + assert "".join(argument_deltas) == '{"location":"Berlin"}' diff --git a/tests/entrypoints/openai/speech_to_text/test_transcription_inter_chunk_spacing.py b/tests/entrypoints/openai/speech_to_text/test_transcription_inter_chunk_spacing.py new file mode 100644 index 00000000000..1e80d47f82f --- /dev/null +++ b/tests/entrypoints/openai/speech_to_text/test_transcription_inter_chunk_spacing.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ASR inter-chunk spacing: ``asr_inter_chunk_separator`` and transcription +serving (mocked). + +Unit tests cover the helper and ``SupportsTranscription.no_space_languages``. +Integration-style tests exercise ``OpenAIServingTranscription`` streaming and +``create_transcription`` without loading a model. +""" + +from __future__ import annotations + +import json +from collections.abc import AsyncGenerator +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from vllm.config import ModelConfig +from vllm.config.speech_to_text import SpeechToTextConfig +from vllm.entrypoints.openai.engine.protocol import ( + ErrorResponse, + RequestResponseMetadata, +) +from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.speech_to_text.protocol import TranscriptionRequest +from vllm.entrypoints.openai.speech_to_text.serving import OpenAIServingTranscription +from vllm.entrypoints.openai.speech_to_text.speech_to_text import ( + OpenAISpeechToText, + asr_inter_chunk_separator, +) +from vllm.model_executor.models.interfaces import SupportsTranscription +from vllm.outputs import CompletionOutput, RequestOutput + +# --- Unit: helper + protocol ------------------------------------------------- + + +def test_default_no_space_languages_includes_zh_and_ja(): + assert SupportsTranscription.no_space_languages == {"ja", "zh"} + + +@pytest.mark.parametrize( + ("language", "expected_sep"), + [ + ("en", " "), + ("EN", " "), + ("zh", ""), + ("ZH", ""), + ("ja", ""), + (None, " "), + ], +) +def test_asr_inter_chunk_separator_matches_protocol(language, expected_sep): + sep = asr_inter_chunk_separator(language, SupportsTranscription.no_space_languages) + assert sep == expected_sep + + +def test_joined_chunks_english_has_space_between(): + sep = asr_inter_chunk_separator("en", SupportsTranscription.no_space_languages) + assert sep.join(["hello", "world"]) == "hello world" + + +def test_joined_chunks_chinese_has_no_space_between(): + sep = asr_inter_chunk_separator("zh", SupportsTranscription.no_space_languages) + assert sep.join(["你好", "世界"]) == "你好世界" + + +# --- Integration: serving (no model) ----------------------------------------- + + +class _StubTranscriptionModel: + """Minimal stand-in for a SupportsTranscription implementation (no torch).""" + + no_space_languages: set[str] = {"ja", "zh"} + supports_segment_timestamp = False + + @classmethod + def get_speech_to_text_config( + cls, model_config: ModelConfig, task_type: str + ) -> SpeechToTextConfig: + return SpeechToTextConfig( + sample_rate=16000.0, + max_audio_clip_s=5.0, + ) + + @classmethod + def post_process_output(cls, text: str) -> str: + return text + + +def _request_output(text: str) -> RequestOutput: + return RequestOutput( + request_id="rid", + prompt=None, + prompt_token_ids=None, + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text=text, + token_ids=(1, 2, 3), + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + ], + finished=True, + ) + + +def _sse_delta_contents(sse_body: str) -> list[str]: + """Extract ``choices[0].delta.content`` from each ``data:`` line (streaming API).""" + contents: list[str] = [] + for line in sse_body.splitlines(): + if not line.startswith("data: "): + continue + payload = line.removeprefix("data: ").strip() + if payload == "[DONE]": + continue + obj = json.loads(payload) + for choice in obj.get("choices") or []: + delta = choice.get("delta") or {} + if "content" in delta: + contents.append(delta["content"]) + return contents + + +@pytest.mark.asyncio +async def test_transcription_stream_generator_english_inserts_space_between_chunks(): + """Online streaming: first output per audio chunk is prefixed with *separator*.""" + + async def gen_hello() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("hello") + + async def gen_world() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("world") + + serving = OpenAIServingTranscription.__new__(OpenAIServingTranscription) + serving.enable_force_include_usage = False + serving.model_cls = _StubTranscriptionModel + serving.task_type = "transcribe" + request = SimpleNamespace( + model="stub-model", + stream_include_usage=False, + stream_continuous_usage_stats=False, + ) + sep = asr_inter_chunk_separator("en", _StubTranscriptionModel.no_space_languages) + assert sep == " " + + out_lines: list[str] = [] + agen = OpenAIServingTranscription.transcription_stream_generator( + serving, + request=request, + result_generator=[gen_hello(), gen_world()], + request_id="test-req", + request_metadata=RequestResponseMetadata(request_id="test-req"), + audio_duration_s=1.0, + separator=sep, + ) + async for line in agen: + out_lines.append(line) + sse = "".join(out_lines) + combined = "".join(_sse_delta_contents(sse)) + assert combined.strip() == "hello world" + + +@pytest.mark.asyncio +async def test_transcription_stream_generator_chinese_no_space_between_chunks(): + async def gen_a() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("你好") + + async def gen_b() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("世界") + + serving = OpenAIServingTranscription.__new__(OpenAIServingTranscription) + serving.enable_force_include_usage = False + serving.model_cls = _StubTranscriptionModel + serving.task_type = "transcribe" + request = SimpleNamespace( + model="stub-model", + stream_include_usage=False, + stream_continuous_usage_stats=False, + ) + sep = asr_inter_chunk_separator("zh", _StubTranscriptionModel.no_space_languages) + assert sep == "" + + out_lines: list[str] = [] + agen = OpenAIServingTranscription.transcription_stream_generator( + serving, + request=request, + result_generator=[gen_a(), gen_b()], + request_id="test-req-zh", + request_metadata=RequestResponseMetadata(request_id="test-req-zh"), + audio_duration_s=1.0, + separator=sep, + ) + async for line in agen: + out_lines.append(line) + combined = "".join(_sse_delta_contents("".join(out_lines))) + assert combined == "你好世界" + + +@pytest.mark.asyncio +async def test_create_transcription_non_streaming_joins_chunks_by_language(): + """``create_transcription`` uses the same separator logic as the helper.""" + + async def gen_hello() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("hello") + + async def gen_world() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("world") + + engine_client = MagicMock() + engine_client.model_config = MagicMock() + engine_client.model_config.get_diff_sampling_param.return_value = { + "max_tokens": 256, + "temperature": 0.0, + } + engine_client.model_config.max_model_len = 8192 + engine_client.errored = False + engine_client.generate.side_effect = [gen_hello(), gen_world()] + + models = MagicMock(spec=OpenAIServingModels) + models.lora_requests = {} + models.is_base_model.return_value = True + + preprocess_mock = AsyncMock(return_value=([MagicMock(), MagicMock()], 1.0)) + + with ( + patch( + "vllm.model_executor.model_loader.get_model_cls", + return_value=_StubTranscriptionModel, + ), + patch.object(OpenAISpeechToText, "_preprocess_speech_to_text", preprocess_mock), + ): + serving = OpenAIServingTranscription(engine_client, models, request_logger=None) + + req_en = TranscriptionRequest.model_construct( + file=MagicMock(), + model="stub-model", + language="en", + stream=False, + response_format="json", + ) + out_en = await serving.create_transcription( + b"\x00\x00", req_en, raw_request=None + ) + assert not isinstance(out_en, ErrorResponse) + assert out_en.text == "hello world" + + async def gen_nihao() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("你好") + + async def gen_shijie() -> AsyncGenerator[RequestOutput, None]: + yield _request_output("世界") + + engine_client.generate.side_effect = [gen_nihao(), gen_shijie()] + + req_zh = TranscriptionRequest.model_construct( + file=MagicMock(), + model="stub-model", + language="zh", + stream=False, + response_format="json", + ) + out_zh = await serving.create_transcription( + b"\x00\x00", req_zh, raw_request=None + ) + assert not isinstance(out_zh, ErrorResponse) + assert out_zh.text == "你好世界" diff --git a/tests/entrypoints/openai/test_tool_calls_serialization.py b/tests/entrypoints/openai/test_tool_calls_serialization.py new file mode 100644 index 00000000000..cedc2575b80 --- /dev/null +++ b/tests/entrypoints/openai/test_tool_calls_serialization.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for tool_calls Iterable → list materialisation. + +Regression tests for https://github.com/vllm-project/vllm/issues/34792. + +Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral +models because: + 1. The OpenAI Python SDK types tool_calls as Iterable[...] in + ChatCompletionAssistantMessageParam. + 2. Pydantic v2, when validating from Python objects (not from raw JSON), + wraps Iterable fields in a one-shot lazy iterator. + 3. Debug logging called model_dump_json() which consumed that iterator. + 4. The Mistral tokenizer then saw empty tool_calls and raised + "ValueError: Unexpected tool call id ...". +""" + +import pytest + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + + +def _make_tool_call(tc_id: str, name: str, args: str) -> dict: + return { + "id": tc_id, + "type": "function", + "function": {"name": name, "arguments": args}, + } + + +def _make_request(messages: list) -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=messages, + ) + + +def test_tool_calls_list_preserved_after_model_dump(): + """tool_calls in assistant messages must be readable after model_dump_json. + + When the request is built from Python dicts (as in the Anthropic → OpenAI + conversion path), Pydantic v2 previously wrapped the Iterable tool_calls + in a one-shot iterator. model_dump_json() consumed it, leaving subsequent + readers (e.g. the Mistral tokenizer) with an empty sequence. + """ + tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}') + messages = [ + {"role": "user", "content": "What is the weather in Paris?"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + { + "role": "tool", + "tool_call_id": "call_abc123", + "content": '{"temperature": 20}', + }, + ] + + req = _make_request(messages) + + # Simulate debug logging: serialize the model (this was the trigger) + _ = req.model_dump_json() + + # The assistant message must still have accessible tool_calls afterwards + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + tool_calls = assistant_msg.get("tool_calls") + assert tool_calls is not None, "tool_calls must not be None after model_dump_json" + assert isinstance(tool_calls, list), "tool_calls must be a list" + assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json" + + +def test_tool_calls_from_generator_are_materialised(): + """tool_calls passed as a generator must be converted to list on validation.""" + tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}') + + def tool_calls_gen(): + yield tool_call + + messages = [ + {"role": "user", "content": "Search for vllm"}, + { + "role": "assistant", + "content": None, + "tool_calls": tool_calls_gen(), # one-shot generator + }, + ] + + req = _make_request(messages) + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + + # Iterate twice — must not raise or return empty on second pass + tool_calls_first = list(assistant_msg.get("tool_calls", [])) + tool_calls_second = list(assistant_msg.get("tool_calls", [])) + + assert len(tool_calls_first) == 1, "First read must return the tool call" + assert len(tool_calls_second) == 1, "Second read must also return the tool call" + + +def test_tool_calls_list_passthrough(): + """tool_calls already provided as a list must remain a list.""" + tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}') + messages = [ + {"role": "user", "content": "Calculate 2+2"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + ] + + req = _make_request(messages) + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + assert isinstance(assistant_msg.get("tool_calls"), list) + + +def test_messages_without_tool_calls_unaffected(): + """Messages without tool_calls must be handled correctly.""" + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi there!"}, + ] + + req = _make_request(messages) + # None of the messages should have tool_calls injected + for msg in req.messages: + assert isinstance(msg, dict) + assert msg.get("tool_calls") is None or msg.get("tool_calls") == [] + + +@pytest.mark.parametrize("num_tool_calls", [1, 3]) +def test_multiple_tool_calls_materialised(num_tool_calls: int): + """Multiple tool calls in a single message are all preserved.""" + tool_calls = [ + _make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}') + for i in range(num_tool_calls) + ] + messages = [ + {"role": "user", "content": "Do things"}, + {"role": "assistant", "content": None, "tool_calls": iter(tool_calls)}, + ] + + req = _make_request(messages) + assistant_msg = req.messages[1] + assert isinstance(assistant_msg, dict) + + result_tool_calls = assistant_msg.get("tool_calls") + assert isinstance(result_tool_calls, list) + assert len(result_tool_calls) == num_tool_calls + + # Verify after model_dump_json too + _ = req.model_dump_json() + assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls diff --git a/tests/entrypoints/pooling/classify/test_offline.py b/tests/entrypoints/pooling/classify/test_offline.py index f556dd579e6..828443ebc41 100644 --- a/tests/entrypoints/pooling/classify/test_offline.py +++ b/tests/entrypoints/pooling/classify/test_offline.py @@ -110,8 +110,11 @@ def test_score_api(llm: LLM): llm.score("ping", "pong", use_tqdm=False) -@pytest.mark.parametrize("task", ["embed", "token_embed"]) +@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask): - err_msg = "Embedding API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Embedding API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/classify/test_online.py b/tests/entrypoints/pooling/classify/test_online.py index ed295a09b88..848ce4083c5 100644 --- a/tests/entrypoints/pooling/classify/test_online.py +++ b/tests/entrypoints/pooling/classify/test_online.py @@ -469,4 +469,8 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/embed/test_io_processor.py b/tests/entrypoints/pooling/embed/test_io_processor.py index f25911b661f..341ccbd5f0c 100644 --- a/tests/entrypoints/pooling/embed/test_io_processor.py +++ b/tests/entrypoints/pooling/embed/test_io_processor.py @@ -4,6 +4,7 @@ import pytest +from vllm import PoolingParams from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor from vllm.entrypoints.pooling.embed.protocol import ( CohereEmbedContent, @@ -218,6 +219,7 @@ class TestPreProcessCohereOnline: def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]: return PoolingServeContext( request=CohereEmbedRequest(model="test", **request_kwargs), + pooling_params=PoolingParams(), model_name="test", request_id="embd-test", ) @@ -233,13 +235,13 @@ class TestPreProcessCohereOnline: ctx = self._make_context(texts=["hello"]) calls: list[tuple[str, object]] = [] - def preprocess_completion(request, prompt_input, prompt_embeds): + def preprocess_cmpl_online(request, prompt_input, prompt_embeds): calls.append(("completion", prompt_input)) return ["completion"] handler._get_task_instruction_prefix = lambda _input_type: None handler._has_chat_template = lambda: False - handler._preprocess_completion_online = preprocess_completion + handler._preprocess_cmpl_online = preprocess_cmpl_online handler._batch_render_chat = lambda *_args, **_kwargs: ( pytest.fail("text-only request should not require chat rendering") ) @@ -254,7 +256,7 @@ class TestPreProcessCohereOnline: ctx = self._make_context(texts=["hello"], input_type="query") calls: list[tuple[str, object]] = [] - def preprocess_completion(request, prompt_input, prompt_embeds): + def preprocess_cmpl(request, prompt_input, prompt_embeds): calls.append(("completion", prompt_input)) return ["fallback"] @@ -263,7 +265,7 @@ class TestPreProcessCohereOnline: handler._batch_render_chat = lambda *_args, **_kwargs: ( pytest.fail("chat rendering should be skipped without a template") ) - handler._preprocess_completion_online = preprocess_completion + handler._preprocess_cmpl_online = preprocess_cmpl handler._pre_process_cohere_online(ctx) @@ -297,7 +299,7 @@ class TestPreProcessCohereOnline: handler._get_task_instruction_prefix = lambda _input_type: "query: " handler._has_chat_template = lambda: True handler._batch_render_chat = batch_render_chat - handler._preprocess_completion_online = lambda *_args, **_kwargs: ( + handler._preprocess_cmpl_online = lambda *_args, **_kwargs: ( pytest.fail("completion path should be skipped when a template exists") ) diff --git a/tests/entrypoints/pooling/embed/test_offline.py b/tests/entrypoints/pooling/embed/test_offline.py index e8d84ed45e0..1ffeb027b48 100644 --- a/tests/entrypoints/pooling/embed/test_offline.py +++ b/tests/entrypoints/pooling/embed/test_offline.py @@ -107,8 +107,11 @@ def test_pooling_params(llm: LLM): ) -@pytest.mark.parametrize("task", ["token_classify", "classify"]) +@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask): - err_msg = "Classification API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Classification API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/embed/test_online.py b/tests/entrypoints/pooling/embed/test_online.py index dc61244c944..3032645c7b1 100644 --- a/tests/entrypoints/pooling/embed/test_online.py +++ b/tests/entrypoints/pooling/embed/test_online.py @@ -767,4 +767,8 @@ async def test_pooling_not_supported( }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py b/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py index 38146084e37..39251405664 100644 --- a/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py +++ b/tests/entrypoints/pooling/scoring/test_bi_encoder_online.py @@ -411,4 +411,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str): }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py index cb76d74608e..56e83de3f74 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_offline.py @@ -112,6 +112,35 @@ def test_classify(llm): assert len(outputs[0].outputs.data) == 1 +@pytest.mark.skip_global_cleanup +def test_max_tokens_per_doc(llm: LLM): + """Test max_tokens_per_doc via PoolingParams.extra_kwargs (offline).""" + long_doc = "The capital of France is Paris. " * 20 + + # Without truncation + outputs_no_limit = llm.score( + TEXTS_1[0], + long_doc, + use_tqdm=False, + ) + + # With truncation via extra_kwargs + outputs_with_limit = llm.score( + TEXTS_1[0], + long_doc, + pooling_params=PoolingParams(extra_kwargs={"max_tokens_per_doc": 10}), + use_tqdm=False, + ) + + assert len(outputs_no_limit) == 1 + assert len(outputs_with_limit) == 1 + + # Truncated version should have fewer prompt tokens + no_limit_tokens = len(outputs_no_limit[0].prompt_token_ids) + with_limit_tokens = len(outputs_with_limit[0].prompt_token_ids) + assert with_limit_tokens < no_limit_tokens + + def test_pooling_params(llm: LLM): def get_outputs(use_activation): outputs = llm.score( diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py index ebb339263ce..5eaa2d92805 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online.py @@ -471,6 +471,78 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer): assert len(poolings.data[0].data[0]) == 1 +@pytest.mark.asyncio +async def test_rerank_max_tokens_per_doc( + server: RemoteOpenAIServer, +): + """Test that max_tokens_per_doc actually reduces the token count.""" + query = "What is the capital of France?" + # Use a doc that fits within max_model_len=100 (query ~8 tokens + 4 special) + long_doc = "The capital of France is Paris. " * 10 # ~70 tokens + + # Without max_tokens_per_doc + response_no_limit = requests.post( + server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": [long_doc], + "truncate_prompt_tokens": 99, + }, + ) + response_no_limit.raise_for_status() + rerank_no_limit = RerankResponse.model_validate(response_no_limit.json()) + + # With max_tokens_per_doc + response_with_limit = requests.post( + server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": [long_doc], + "max_tokens_per_doc": 10, + }, + ) + response_with_limit.raise_for_status() + rerank_with_limit = RerankResponse.model_validate(response_with_limit.json()) + + assert rerank_with_limit.usage.prompt_tokens < rerank_no_limit.usage.prompt_tokens + + +@pytest.mark.asyncio +async def test_rerank_max_tokens_per_doc_validation( + server: RemoteOpenAIServer, +): + """Test that max_tokens_per_doc validation works correctly.""" + query = "What is the capital of France?" + documents = ["The capital of France is Paris."] + + # Test with max_tokens_per_doc=0 (should succeed — means no truncation) + response = requests.post( + server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": documents, + "max_tokens_per_doc": 0, + }, + ) + response.raise_for_status() + + # Test with invalid max_tokens_per_doc (negative) + response = requests.post( + server.url_for("rerank"), + json={ + "model": MODEL_NAME, + "query": query, + "documents": documents, + "max_tokens_per_doc": -5, + }, + ) + assert response.status_code == 400 + assert "max_tokens_per_doc must be a non-negative integer" in response.text + + @pytest.mark.asyncio @pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"]) async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str): @@ -484,4 +556,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str): }, ) assert response.json()["error"]["type"] == "BadRequestError" - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}") + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = f"Unsupported task: {task!r}" + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/token_classify/test_offline.py b/tests/entrypoints/pooling/token_classify/test_offline.py index f7a74675463..d36761466df 100644 --- a/tests/entrypoints/pooling/token_classify/test_offline.py +++ b/tests/entrypoints/pooling/token_classify/test_offline.py @@ -65,14 +65,17 @@ def test_score_api(llm: LLM): llm.score("ping", "pong", use_tqdm=False) -@pytest.mark.parametrize("task", ["classify", "embed", "token_embed"]) +@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): if task == "classify": with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"): llm.encode(prompt, pooling_task=task, use_tqdm=False) assert "deprecated" in caplog_vllm.text else: - err_msg = "Embedding API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Embedding API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/token_classify/test_online.py b/tests/entrypoints/pooling/token_classify/test_online.py index e91d0bc9a39..39fd378336c 100644 --- a/tests/entrypoints/pooling/token_classify/test_online.py +++ b/tests/entrypoints/pooling/token_classify/test_online.py @@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"]) +@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"]) async def test_pooling_not_supported( server: RemoteOpenAIServer, model_name: str, task: str ): @@ -64,7 +64,8 @@ async def test_pooling_not_supported( }, ) - if task != "classify": - assert response.json()["error"]["type"] == "BadRequestError" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: err_msg = f"Unsupported task: {task!r}" - assert response.json()["error"]["message"].startswith(err_msg) + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/pooling/token_embed/test_offline.py b/tests/entrypoints/pooling/token_embed/test_offline.py index 697f4f81a11..d2e87fbf23e 100644 --- a/tests/entrypoints/pooling/token_embed/test_offline.py +++ b/tests/entrypoints/pooling/token_embed/test_offline.py @@ -62,14 +62,17 @@ def test_token_ids_prompts(llm: LLM): assert outputs[0].outputs.data.shape == (11, 384) -@pytest.mark.parametrize("task", ["embed", "classify", "token_classify"]) +@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"]) def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm): if task == "embed": with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"): llm.encode(prompt, pooling_task=task, use_tqdm=False) assert "deprecated" in caplog_vllm.text else: - err_msg = "Classification API is not supported by this model.+" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: + err_msg = "Classification API is not supported by this model.+" with pytest.raises(ValueError, match=err_msg): llm.encode(prompt, pooling_task=task, use_tqdm=False) diff --git a/tests/entrypoints/pooling/token_embed/test_online.py b/tests/entrypoints/pooling/token_embed/test_online.py index 922c624e98e..048491dac15 100644 --- a/tests/entrypoints/pooling/token_embed/test_online.py +++ b/tests/entrypoints/pooling/token_embed/test_online.py @@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str): @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) -@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"]) +@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"]) async def test_pooling_not_supported( server: RemoteOpenAIServer, model_name: str, task: str ): @@ -87,7 +87,8 @@ async def test_pooling_not_supported( }, ) - if task != "embed": - assert response.json()["error"]["type"] == "BadRequestError" + if task == "plugin": + err_msg = "No IOProcessor plugin installed." + else: err_msg = f"Unsupported task: {task!r}" - assert response.json()["error"]["message"].startswith(err_msg) + assert response.json()["error"]["message"].startswith(err_msg) diff --git a/tests/entrypoints/serve/disagg/test_generate_stream.py b/tests/entrypoints/serve/disagg/test_generate_stream.py index a9ca026306f..76a9df22f69 100644 --- a/tests/entrypoints/serve/disagg/test_generate_stream.py +++ b/tests/entrypoints/serve/disagg/test_generate_stream.py @@ -86,7 +86,6 @@ def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens: serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, @@ -148,7 +147,6 @@ def _mock_engine() -> MagicMock: engine.errored = False engine.model_config = MockModelConfig() engine.input_processor = MagicMock() - engine.io_processor = MagicMock() engine.renderer = _build_renderer(engine.model_config) return engine diff --git a/tests/entrypoints/serve/instrumentator/test_metrics.py b/tests/entrypoints/serve/instrumentator/test_metrics.py index ba4e65977c7..9095f80e20f 100644 --- a/tests/entrypoints/serve/instrumentator/test_metrics.py +++ b/tests/entrypoints/serve/instrumentator/test_metrics.py @@ -182,6 +182,7 @@ async def test_metrics_counts( EXPECTED_METRICS_V1 = [ "vllm:num_requests_running", "vllm:num_requests_waiting", + "vllm:num_requests_waiting_by_reason", "vllm:kv_cache_usage_perc", "vllm:prefix_cache_queries", "vllm:prefix_cache_hits", diff --git a/tests/entrypoints/serve/lora/test_serving_models.py b/tests/entrypoints/serve/lora/test_serving_models.py index f6755f48934..ce9fdcc2bfb 100644 --- a/tests/entrypoints/serve/lora/test_serving_models.py +++ b/tests/entrypoints/serve/lora/test_serving_models.py @@ -34,7 +34,6 @@ async def _async_serving_models_init() -> OpenAIServingModels: mock_model_config.max_model_len = 2048 mock_engine_client.model_config = mock_model_config mock_engine_client.input_processor = MagicMock() - mock_engine_client.io_processor = MagicMock() mock_engine_client.renderer = MagicMock() serving_models = OpenAIServingModels( diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml index d9110efaaad..b0f886a86ad 100644 --- a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16.yaml @@ -8,4 +8,5 @@ server_args: >- --max-model-len 4096 --tensor-parallel-size 8 --enable-expert-parallel + --mamba-backend flashinfer --speculative-config '{"method":"mtp","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml index 50f09731946..71ba7d52f14 100644 --- a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-NVFP4.yaml @@ -8,4 +8,5 @@ server_args: >- --max-model-len 4096 --tensor-parallel-size 2 --enable-expert-parallel + --mamba-backend flashinfer --speculative-config '{"method":"mtp","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml new file mode 100644 index 00000000000..fedb7416960 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k3v4nc.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.78 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_k3v4_nc --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml new file mode 100644 index 00000000000..9717333582b --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-k8v4.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_k8v4 --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml new file mode 100644 index 00000000000..8ece1852625 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t3nc.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.75 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_3bit_nc --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml new file mode 100644 index 00000000000..9b3a14f9b95 --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-4B-TQ-t4nc.yaml @@ -0,0 +1,5 @@ +model_name: "Qwen/Qwen3-4B" +accuracy_threshold: 0.80 +num_questions: 1319 +num_fewshot: 5 +server_args: "--kv-cache-dtype turboquant_4bit_nc --enforce-eager --max-model-len 4096" diff --git a/tests/evals/gsm8k/configs/models-turboquant.txt b/tests/evals/gsm8k/configs/models-turboquant.txt new file mode 100644 index 00000000000..518aac780b9 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-turboquant.txt @@ -0,0 +1,4 @@ +Qwen3-4B-TQ-k8v4.yaml +Qwen3-4B-TQ-t4nc.yaml +Qwen3-4B-TQ-k3v4nc.yaml +Qwen3-4B-TQ-t3nc.yaml diff --git a/tests/kernels/attention/test_lightning_attn.py b/tests/kernels/attention/test_lightning_attn.py index 37fd85ccec0..46757cc10b6 100644 --- a/tests/kernels/attention/test_lightning_attn.py +++ b/tests/kernels/attention/test_lightning_attn.py @@ -122,8 +122,6 @@ def test_linear_decode_forward_triton( dtype: torch.dtype, ): torch.set_default_device("cuda") - torch.manual_seed(42) - torch.cuda.manual_seed_all(42) set_random_seed(42) base = 0.01 q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype) @@ -165,8 +163,6 @@ def test_linear_decode_forward_triton_with_padding( dtype: torch.dtype, ): torch.set_default_device("cuda") - torch.manual_seed(42) - torch.cuda.manual_seed_all(42) set_random_seed(42) batch_size = 4 @@ -229,8 +225,6 @@ def test_lightning_attention_reference( dtype: torch.dtype, ): torch.set_default_device("cuda") - torch.manual_seed(42) - torch.cuda.manual_seed_all(42) set_random_seed(42) base = 0.01 diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py index f9c01f4f1e6..07d15e3b1df 100644 --- a/tests/kernels/core/test_fused_quant_layernorm.py +++ b/tests/kernels/core/test_fused_quant_layernorm.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import ( per_token_group_quant_int8, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed DTYPES = [torch.bfloat16, torch.float] QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()] @@ -180,9 +181,7 @@ def test_rms_norm( device: str, strided_input: bool, ) -> None: - torch.random.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) + set_random_seed(seed) torch.set_default_device(device) torch.accelerator.set_device_index(device) diff --git a/tests/kernels/core/test_minimax_reduce_rms.py b/tests/kernels/core/test_minimax_reduce_rms.py new file mode 100644 index 00000000000..d17a448bd97 --- /dev/null +++ b/tests/kernels/core/test_minimax_reduce_rms.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MiniMax QK RMS-norm: NCCL reference vs Lamport fused kernel.""" + +import pytest +import torch +import torch.nn as nn +from torch.multiprocessing import spawn + +from tests.kernels.utils import opcheck +from tests.utils import ensure_current_vllm_config, init_test_distributed_environment +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port +from vllm.utils.torch_utils import set_random_seed + + +@ensure_current_vllm_config() +def _worker_forward_qk( + local_rank, + world_size, + port, + num_tokens, + hidden_q_full, + hidden_k_full, + dtype, + seed, + eps, +): + """Per-rank worker: compare NCCL allreduce path vs Lamport fused kernel.""" + + if not hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + cleanup_dist_env_and_memory() + return + device = torch.device(f"cuda:{local_rank}") + torch.accelerator.set_device_index(device) + init_test_distributed_environment( + world_size, 1, local_rank, port, local_rank=local_rank + ) + + hq = hidden_q_full // world_size + hk = hidden_k_full // world_size + + q_norm = MiniMaxText01RMSNormTP(hidden_q_full, eps=eps).cuda() + k_norm = MiniMaxText01RMSNormTP(hidden_k_full, eps=eps).cuda() + + set_random_seed(seed) + qw = torch.randn(hidden_q_full, dtype=dtype, device="cuda") + kw = torch.randn(hidden_k_full, dtype=dtype, device="cuda") + q_norm.weight = nn.Parameter(qw[local_rank * hq : (local_rank + 1) * hq]) + k_norm.weight = nn.Parameter(kw[local_rank * hk : (local_rank + 1) * hk]) + + torch.manual_seed(seed + 1000 + local_rank) + qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda") + + q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1) + ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref) + + # Set up Lamport workspace. + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.layers.mamba.lamport_workspace import ( + get_allreduce_workspace, + ) + + workspace = get_allreduce_workspace( + rank=local_rank, + world_size=world_size, + max_tokens=num_tokens, + process_group=get_tp_group().cpu_group, + ) + + opcheck( + torch.ops._C.minimax_allreduce_rms_qk, + ( + qkv.clone(), + q_norm.weight, + k_norm.weight, + workspace, + hq, + hk, + local_rank, + world_size, + eps, + ), + ) + fused_q, fused_k = torch.ops._C.minimax_allreduce_rms_qk( + qkv.clone(), + q_norm.weight, + k_norm.weight, + workspace, + hq, + hk, + local_rank, + world_size, + eps, + ) + _, _, fused_v = qkv.split([hq, hk, hk], dim=-1) + torch.accelerator.synchronize() + + torch.testing.assert_close( + fused_q, + ref_q, + atol=3e-2, + rtol=3e-2, + ) + torch.testing.assert_close(fused_k, ref_k, atol=3e-2, rtol=3e-2) + + cleanup_dist_env_and_memory() + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="CUDA required", +) +@pytest.mark.parametrize("world_size", [2, 4, 8]) +@pytest.mark.parametrize("num_tokens", [1, 128, 333]) +@pytest.mark.parametrize( + "hidden_dims", + [(6144, 1024)], +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("eps", [1e-6]) +@pytest.mark.parametrize("seed", [42]) +def test_minimax_reduce_rms_qk( + world_size, + num_tokens, + hidden_dims, + dtype, + eps, + seed, +): + num_gpus = current_platform.device_count() + if num_gpus < world_size: + pytest.skip(f"Need >= {world_size} GPUs, have {num_gpus}") + hidden_q_full, hidden_k_full = hidden_dims + port = str(get_open_port()) + spawn( + _worker_forward_qk, + args=( + world_size, + port, + num_tokens, + hidden_q_full, + hidden_k_full, + dtype, + seed, + eps, + ), + nprocs=world_size, + join=True, + ) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py new file mode 100644 index 00000000000..96b04f44d22 --- /dev/null +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + FlashInferSSUBackend, + TritonSSUBackend, + get_mamba_ssu_backend, + initialize_mamba_ssu_backend, + selective_state_update, +) +from vllm.utils.torch_utils import set_random_seed + +try: + import flashinfer.mamba # noqa: F401 + + HAS_FLASHINFER = True +except ImportError: + HAS_FLASHINFER = False + + +def test_default_backend_is_triton(): + initialize_mamba_ssu_backend(MambaConfig()) + backend = get_mamba_ssu_backend() + assert isinstance(backend, TritonSSUBackend) + assert backend.name == "triton" + + +def test_explicit_triton_backend(): + initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON)) + backend = get_mamba_ssu_backend() + assert isinstance(backend, TritonSSUBackend) + + +@pytest.mark.skipif(not HAS_FLASHINFER, reason="flashinfer not installed") +def test_flashinfer_backend_init(): + initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.FLASHINFER)) + backend = get_mamba_ssu_backend() + assert isinstance(backend, FlashInferSSUBackend) + assert backend.name == "flashinfer" + + +def test_uninitialized_backend_raises(): + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod + + old = mod._mamba_ssu_backend + mod._mamba_ssu_backend = None + with pytest.raises(RuntimeError, match="not been initialized"): + get_mamba_ssu_backend() + mod._mamba_ssu_backend = old + + +@pytest.mark.skipif(HAS_FLASHINFER, reason="flashinfer is installed") +def test_flashinfer_import_error(): + with pytest.raises(ImportError, match="FlashInfer is required"): + FlashInferSSUBackend(MambaConfig()) + + +def test_triton_basic_call(): + set_random_seed(0) + initialize_mamba_ssu_backend(MambaConfig(backend=MambaBackendEnum.TRITON)) + device = "cuda" + batch_size = 2 + dim = 64 + dstate = 16 + + state = torch.randn(batch_size, dim, dstate, device=device) + x = torch.randn(batch_size, dim, device=device) + out = torch.empty_like(x) + dt = torch.randn(batch_size, dim, device=device) + dt_bias = torch.rand(dim, device=device) - 4.0 + A = -torch.rand(dim, dstate, device=device) + B = torch.randn(batch_size, dstate, device=device) + C = torch.randn(batch_size, dstate, device=device) + D = torch.randn(dim, device=device) + + selective_state_update( + state, + x, + dt, + A, + B, + C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + out=out, + ) + assert not torch.isnan(out).any() diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index a6f3bc35a0b..f07d4c75e75 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -46,6 +46,7 @@ from vllm.utils.import_utils import ( has_deep_gemm, has_mori, ) +from vllm.utils.math_utils import next_power_of_2 from .mk_objects import ( TestMoEQuantConfig, @@ -604,13 +605,6 @@ def make_modular_kernel( vllm_config: VllmConfig, quant_config: FusedMoEQuantConfig, ) -> mk.FusedMoEKernel: - def next_power_of_2(x): - import math - - if x == 0: - return 1 - return 2 ** math.ceil(math.log2(x)) - # make moe config moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make( tp_size_=get_tensor_model_parallel_world_size(), diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index 10a226bcd97..95004fa0ab4 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -126,7 +126,7 @@ def parallel_launch_with_config( world_size: int, worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None], vllm_config: VllmConfig, - env_dict: dict[Any, Any], + env_dict: dict[Any, Any] | None, *args: P.args, **kwargs: P.kwargs, ) -> None: diff --git a/tests/kernels/moe/test_cutedsl_moe.py b/tests/kernels/moe/test_cutedsl_moe.py index 2a6f83695c4..c3e7190ca09 100644 --- a/tests/kernels/moe/test_cutedsl_moe.py +++ b/tests/kernels/moe/test_cutedsl_moe.py @@ -142,7 +142,9 @@ def prepare_inputs( # Initialize the hidden_states_3d with ones instead of empty to avoid nan # issue. hidden_states_3d = torch.ones( - (num_experts, max(masked_m), hidden_states.shape[1]), dtype=hidden_states.dtype + (num_experts, max(masked_m), hidden_states.shape[1]), + dtype=hidden_states.dtype, + device=hidden_states.device, ) for i in range(num_experts): hidden_states_3d[i, : masked_m[i], :] = hidden_states[topk_idx.view(-1) == i] @@ -426,7 +428,7 @@ def test_flashinfer_cutedsl_moe_masked( w1_alpha = 1.0 / (input_global_scale * w1_global_scale) w2_alpha = 1.0 / (a2_global_scale * w2_global_scale) - out = torch.empty_like(hidden_states_3d) + out = torch.empty_like(hidden_states_3d, device=hidden_states.device) # Note: the 1st dim shouldn't be bs wk = torch.empty( num_experts, diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index 6bde13e0ecf..6caa9d8c068 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -29,6 +29,7 @@ from vllm.utils.deep_gemm import ( is_deep_gemm_supported, ) from vllm.utils.import_utils import has_deep_ep, has_deep_gemm +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.workspace import init_workspace_manager @@ -84,14 +85,6 @@ def with_dp_metadata(M: int, world_size: int): yield -def next_power_of_2(x): - import math - - if x == 0: - return 1 - return 2 ** math.ceil(math.log2(x)) - - def make_block_quant_fp8_weights( e: int, n: int, diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index db499b68843..dad25bc3195 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8 from vllm.model_executor.models.llama4 import Llama4MoE from vllm.platforms import current_platform +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed try: @@ -174,6 +175,7 @@ class TestData: routing_method=layer.routing_method_type, activation=activation, device=w13_quantized.device, + max_num_tokens=next_power_of_2(m), ) return TestData( @@ -348,6 +350,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( in_dtype=torch.bfloat16, is_act_and_mul=activation.is_gated, routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), ) kernel = mk.FusedMoEKernel( diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index a3fb474f151..d116a96f58b 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability( @@ -105,6 +106,7 @@ def test_flashinfer_fp4_moe_no_graph( in_dtype=dtype, is_act_and_mul=is_gated_act, routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), ) flashinfer_experts = FusedMoEKernel( diff --git a/tests/kernels/moe/test_gpt_oss_triton_kernels.py b/tests/kernels/moe/test_gpt_oss_triton_kernels.py index 032b4fc047c..c61004acaa3 100644 --- a/tests/kernels/moe/test_gpt_oss_triton_kernels.py +++ b/tests/kernels/moe/test_gpt_oss_triton_kernels.py @@ -25,7 +25,7 @@ from triton_kernels.tensor_details import layout from triton_kernels.testing import assert_close from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config -from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( +from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( triton_kernel_moe_forward, ) from vllm.utils.math_utils import round_up diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index 0ce3d165d0b..0ad2cd06ee3 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ) from vllm.platforms import current_platform from vllm.scalar_type import scalar_types +from vllm.utils.torch_utils import set_random_seed def mxint4_quantize( @@ -134,7 +135,7 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group pytest.importorskip("flashinfer") monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_INT4", "1") - torch.cuda.manual_seed(0) + set_random_seed(0) dtype = torch.bfloat16 @@ -289,7 +290,7 @@ def test_flashinfer_trtllm_mxint4_moe_wrapper(m, n, k, e, topk): flashinfer_trtllm_mxint4_moe, ) - torch.cuda.manual_seed(0) + set_random_seed(0) dtype = torch.bfloat16 a = torch.randn((m, k), device="cuda", dtype=dtype) * 0.5 diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index b071e72dafb..589d90d1eca 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -29,7 +29,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, ) from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config -from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( +from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( OAITritonExperts, UnfusedOAITritonExperts, ) diff --git a/tests/kernels/moe/test_moe.py b/tests/kernels/moe/test_moe.py index 5ec2b106e12..35b21320f82 100644 --- a/tests/kernels/moe/test_moe.py +++ b/tests/kernels/moe/test_moe.py @@ -59,6 +59,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_w from vllm.model_executor.models.mixtral import MixtralMoE from vllm.platforms import current_platform from vllm.scalar_type import ScalarType, scalar_types +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.workspace import init_workspace_manager @@ -1031,7 +1032,7 @@ def test_fused_marlin_moe( act_order: bool, is_k_full: bool, ): - torch.cuda.manual_seed(1) + set_random_seed(1) group_size = group_blocks if group_blocks <= 0 else group_blocks * 16 if c_type == scalar_types.float16: @@ -1131,7 +1132,7 @@ def test_fused_marlin_moe( @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm") @pytest.mark.parametrize("m", [1, 256]) def test_fused_marlin_moe_with_bias(m): - torch.cuda.manual_seed(0) + set_random_seed(0) e, topk = 32, 4 n, k = 2048, 2048 @@ -1213,7 +1214,7 @@ def test_fused_marlin_moe_non_gated( Non-gated activations like relu2 don't have the gate-up projection pattern, so w1 has shape (e, n, k) instead of (e, 2*n, k). """ - torch.cuda.manual_seed(42) + set_random_seed(42) group_size = 16 # NVFP4 group size is_k_full = True @@ -1397,7 +1398,7 @@ def test_cpu_fused_moe_basic( from vllm.model_executor.layers.fused_moe.cpu_fused_moe import CPUFusedMOE device = "cpu" - torch.manual_seed(7) + set_random_seed(7) a = torch.randn((m, k), device=device, dtype=dtype) / 10 w13 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 10 @@ -1469,7 +1470,7 @@ def test_batched_fused_marlin_moe( f"topk={topk}, " f"max_tokens_per_batch={max_tokens_per_batch}" ) - torch.cuda.manual_seed(0) + set_random_seed(0) dtype = torch.bfloat16 quant_dtype = scalar_types.float4_e2m1f @@ -1676,7 +1677,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend( in_dtype=dtype, is_act_and_mul=True, routing_method=RoutingMethodType.Renormalize, - max_num_tokens=m, + max_num_tokens=next_power_of_2(m), ) with set_current_vllm_config(vllm_config): diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 7b31edd3360..aa2948b8e98 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -26,6 +26,7 @@ from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_wei from vllm.config import ( CompilationConfig, ParallelConfig, + SchedulerConfig, VllmConfig, set_current_vllm_config, ) @@ -53,7 +54,7 @@ from vllm.utils.flashinfer import ( has_flashinfer_nvlink_two_sided, ) from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep -from vllm.utils.math_utils import cdiv +from vllm.utils.math_utils import cdiv, next_power_of_2 from vllm.utils.torch_utils import set_random_seed from vllm.v1.worker.workspace import ( init_workspace_manager, @@ -65,8 +66,9 @@ fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype SHAPE_COMBOS = [ (1, 128, 256), (32, 1024, 512), - (222, 2048, 2048), # should be big enough to exercise DP chunking + (222, 2048, 2048), ] +MAX_M = max([x[0] for x in SHAPE_COMBOS]) NUM_EXPERTS = [8, 64] TOP_KS = [2, 6] @@ -112,7 +114,7 @@ BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = { "mori": {None, "fp8", "modelopt_fp8"}, "flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"}, "flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"}, - "deepep_low_latency": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, + "deepep_low_latency": {None, "modelopt_fp8", "modelopt_fp4"}, "deepep_high_throughput": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, "nixl_ep": {None, "fp8", "modelopt_fp8"}, } @@ -363,9 +365,9 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]: ) # routed_input_transform + quantization + high hidden dimensions - # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors. + # TODO: Disable >= 2048 for now due to insane errors. if ( - (config.use_routed_input_transform or config.backend == "deepep_low_latency") + config.use_routed_input_transform and config.quantization is not None and config.k >= 2048 ): @@ -1663,9 +1665,6 @@ def test_moe_layer( verbosity = pytestconfig.getoption("verbose") - test_env = dict() - test_env["VLLM_MOE_DP_CHUNK_SIZE"] = "128" - monkeypatch.setenv("VLLM_MOE_DP_CHUNK_SIZE", "128") if os.environ.get("VLLM_LOGGING_LEVEL") is None: monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR") @@ -1690,7 +1689,11 @@ def test_moe_layer( compilation_config.pass_config.fuse_allreduce_rms = False # for now vllm_config = VllmConfig( - parallel_config=parallel_config, compilation_config=compilation_config + parallel_config=parallel_config, + compilation_config=compilation_config, + scheduler_config=SchedulerConfig.default_factory( + max_num_batched_tokens=next_power_of_2(MAX_M) + ), ) test_configs = generate_valid_test_configs( @@ -1718,7 +1721,7 @@ def test_moe_layer( world_size, _parallel_worker, vllm_config, - test_env, + None, test_configs, verbosity, ) diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index 422d380f8e4..abe473879f1 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -257,6 +257,41 @@ class TestWeightLoadingWithPaddedHiddenSize: assert torch.equal(expert_data_full, loaded_weight) + def test_narrow_shard_dim(self): + """Simulate loading w2 when both hidden_size and intermediate_size + are padded. + """ + padded_hidden = 3072 + original_hidden = 2688 + padded_intermediate = 1024 + original_intermediate = 896 + + expert_data_full = torch.zeros(padded_hidden, padded_intermediate) + loaded_weight = torch.randn(original_hidden, original_intermediate) + + shard_dim = 1 + hidden_dim = FusedMoE._get_hidden_dim(shard_dim=shard_dim, ndim=2) + expert_data = FusedMoE._narrow_expert_data_for_padding( + expert_data_full, + loaded_weight, + hidden_dim=hidden_dim, + shard_dim=shard_dim, + ) + expert_data.copy_(loaded_weight) + + assert torch.equal( + expert_data_full[:original_hidden, :original_intermediate], + loaded_weight, + ) + assert torch.equal( + expert_data_full[original_hidden:, :], + torch.zeros(padded_hidden - original_hidden, padded_intermediate), + ) + assert torch.equal( + expert_data_full[:original_hidden, original_intermediate:], + torch.zeros(original_hidden, padded_intermediate - original_intermediate), + ) + def test_bnb_shape_mismatch_raises(self): """BnB + padded hidden_size should raise via weight_loader.""" from unittest.mock import MagicMock diff --git a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py index 366009dce99..89e23eb0d74 100644 --- a/tests/kernels/moe/test_shared_fused_moe_routed_transform.py +++ b/tests/kernels/moe/test_shared_fused_moe_routed_transform.py @@ -15,7 +15,7 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.platforms import current_platform -from vllm.utils.torch_utils import is_torch_equal_or_newer +from vllm.utils.torch_utils import is_torch_equal_or_newer, set_random_seed class SimpleLinear(nn.Module): @@ -144,8 +144,7 @@ def test_routed_input_transform_inside_vs_outside( rocm_aiter_ops.refresh_env_variables() - torch.manual_seed(42) - torch.cuda.manual_seed(42) + set_random_seed(42) num_experts = 8 top_k = 2 diff --git a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py index ed58db62d6d..6a5fad8262c 100644 --- a/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py +++ b/tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py @@ -14,7 +14,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) from vllm.platforms import current_platform -from vllm.utils.deep_gemm import DeepGemmQuantScaleFMT, has_deep_gemm +from vllm.utils.deep_gemm import ( + DeepGemmQuantScaleFMT, + has_deep_gemm, + transform_sf_into_required_layout, +) from vllm.utils.math_utils import cdiv, round_up from vllm.utils.torch_utils import set_random_seed @@ -256,8 +260,6 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt and current_platform.has_device_capability(100) and scale_fmt == DeepGemmQuantScaleFMT.UE8M0 ): - from deep_gemm import transform_sf_into_required_layout - _q, _s = ref_with_scale_fmt( E, T, diff --git a/tests/kernels/moe/test_zero_expert_moe.py b/tests/kernels/moe/test_zero_expert_moe.py new file mode 100644 index 00000000000..d8f900256ec --- /dev/null +++ b/tests/kernels/moe/test_zero_expert_moe.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for FusedMoE with zero experts. + +Verifies that: +- The ZeroExpertRouter is properly created and used as the layer router. +- A forward pass through FusedMoE with zero experts produces correct output. +- The output decomposes correctly into real expert + zero expert contributions. + +Note: tests generated with Claude. +""" + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.fused_moe.layer import FusedMoE +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) +from vllm.v1.worker.workspace import init_workspace_manager + + +@pytest.fixture +def zero_expert_moe(dist_init, default_vllm_config): + """Create a FusedMoE layer with zero experts.""" + num_experts = 4 + top_k = 2 + # hidden_size must be >= 256 for the zero expert identity kernel to + # produce output (its BLOCK_SIZE=256 causes grid=0 when hidden_dim<256). + hidden_size = 256 + intermediate_size = 512 + zero_expert_num = 1 + + e_score_correction_bias = torch.zeros( + num_experts + zero_expert_num, + dtype=torch.float32, + device="cuda", + ) + + vllm_config = VllmConfig() + vllm_config.compilation_config.static_forward_context = dict() + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + init_workspace_manager(torch.accelerator.current_device_index()) + + layer = FusedMoE( + zero_expert_type="identity", + e_score_correction_bias=e_score_correction_bias, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + params_dtype=torch.bfloat16, + prefix="test_zero_expert_moe", + renormalize=False, + routed_scaling_factor=1.0, + scoring_func="softmax", + ).cuda() + + layer.quant_method.process_weights_after_loading(layer) + + yield layer, vllm_config + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_router_is_zero_expert_router(zero_expert_moe, num_tokens): + """Verify that FusedMoE with zero_expert_type creates a ZeroExpertRouter.""" + layer, _ = zero_expert_moe + assert isinstance(layer.router, ZeroExpertRouter), ( + f"Expected ZeroExpertRouter but got {type(layer.router).__name__}." + ) + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_no_custom_routing_fn(zero_expert_moe, num_tokens): + """Verify that custom_routing_function is not set (routing is handled + by ZeroExpertRouter, not a memoizing closure).""" + layer, _ = zero_expert_moe + assert layer.custom_routing_function is None + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_forward(zero_expert_moe, num_tokens): + """Run a forward pass through FusedMoE with zero experts and verify output shape.""" + layer, vllm_config = zero_expert_moe + + hidden_size = layer.hidden_size + num_experts = 4 + zero_expert_num = 1 + total_experts = num_experts + zero_expert_num + + hidden_states = torch.randn( + num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda" + ) + router_logits = torch.randn( + num_tokens, total_experts, dtype=torch.float32, device="cuda" + ) + + # Initialize weights to small random values to avoid NaN from + # uninitialized memory. + with torch.no_grad(): + for param in layer.parameters(): + if param.dtype.is_floating_point: + param.normal_(0, 0.01) + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + get_forward_context().all_moe_layers = None + output = layer.forward(hidden_states, router_logits) + + assert output.shape == hidden_states.shape, ( + f"Expected output shape {hidden_states.shape}, got {output.shape}" + ) + assert output.dtype == hidden_states.dtype + assert not torch.isnan(output).any(), "Output contains NaN values" + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens): + """Validate that the FusedMoE output equals a plain FusedMoE + output (real experts only) plus the zero expert contribution. + + The key invariant is: + zero_layer.forward(h, r_full) == plain_layer.forward(h, r_real) + + zero_expert_output + + We create a plain FusedMoE layer with the same weights and real-expert-only + router logits, compute the zero expert output via the ZeroExpertRouter, and + verify the sum matches the FusedMoE output. + """ + layer, vllm_config = zero_expert_moe + num_experts = 4 + zero_expert_num = 1 + total_experts = num_experts + zero_expert_num + + hidden_states = torch.randn( + num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + router_logits = torch.randn( + num_tokens, total_experts, dtype=torch.float32, device="cuda" + ) + + with torch.no_grad(): + for param in layer.parameters(): + if param.dtype.is_floating_point: + param.normal_(0, 0.01) + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + get_forward_context().all_moe_layers = None + + # Create a plain FusedMoE layer with the same config but no zero + # experts. Use a separate prefix to avoid collision. + plain_layer = FusedMoE( + num_experts=num_experts, + top_k=layer.top_k, + hidden_size=layer.hidden_size, + intermediate_size=layer.intermediate_size_per_partition, + params_dtype=torch.bfloat16, + prefix="test_zero_expert_moe_plain", + renormalize=False, + scoring_func="softmax", + e_score_correction_bias=layer.e_score_correction_bias, + ).cuda() + + # Share weights from the zero expert layer. + plain_layer.w13_weight.data.copy_(layer.w13_weight.data) + plain_layer.w2_weight.data.copy_(layer.w2_weight.data) + plain_layer.quant_method.process_weights_after_loading(plain_layer) + + # Compute routing via the ZeroExpertRouter. This produces masked + # topk_weights/topk_ids (zero expert entries have weight=0, id=0) + # and stores zero_expert_output as a side effect. + topk_weights, topk_ids = layer.router.select_experts( + hidden_states, router_logits + ) + zero_output = layer.router.zero_expert_output + + # Compute real expert output using the plain layer with the masked + # routing from the ZeroExpertRouter. + real_output = plain_layer.quant_method.apply( + layer=plain_layer, + x=hidden_states, + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_experts_input=None, + ) + + # Get the combined output from the zero expert layer. + full_output = layer.forward(hidden_states, router_logits) + + assert zero_output is not None, "Zero expert output should not be None" + assert not torch.isnan(real_output).any(), "Real expert output has NaN" + assert not torch.isnan(zero_output).any(), "Zero expert output has NaN" + assert not torch.isnan(full_output).any(), "Full output has NaN" + + expected = real_output + zero_output + torch.testing.assert_close( + full_output, + expected, + atol=0, + rtol=0, + msg="FusedMoE output should equal plain FusedMoE output " + "plus zero expert contribution", + ) + + +@pytest.mark.parametrize("num_tokens", [1, 32]) +def test_zero_expert_moe_zero_expert_is_identity(zero_expert_moe, num_tokens): + """Validate zero expert identity behavior. + + When routing strongly favors the zero expert, its contribution should + be a scaled version of hidden_states (identity operation). We verify + this by manually computing the expected zero expert output from the + routing weights and comparing against what the router produces. + """ + layer, vllm_config = zero_expert_moe + num_experts = 4 + zero_expert_num = 1 + total_experts = num_experts + zero_expert_num + + hidden_states = torch.randn( + num_tokens, layer.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + # Strongly bias toward the zero expert (index 4). + router_logits = torch.full( + (num_tokens, total_experts), -10.0, dtype=torch.float32, device="cuda" + ) + router_logits[:, num_experts] = 10.0 # zero expert gets high logit + + with torch.no_grad(): + for param in layer.parameters(): + if param.dtype.is_floating_point: + param.normal_(0, 0.01) + + with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config): + get_forward_context().all_moe_layers = None + + # Run routing to get topk_weights/topk_ids before masking. + from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, + ) + + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + e_score_correction_bias=layer.router.e_score_correction_bias.data, + topk=layer.top_k, + renormalize=layer.router.renormalize, + scoring_func=layer.router.scoring_func, + ) + + # Manually compute expected zero expert identity output: + # For each token, sum routing weights assigned to zero expert slots, + # then multiply by hidden_states. + zero_mask = topk_ids >= num_experts + zero_weight_per_token = (topk_weights * zero_mask.float()).sum( + dim=-1, keepdim=True + ) + expected_zero_output = (hidden_states.float() * zero_weight_per_token).to( + hidden_states.dtype + ) + + # Run routing directly to trigger zero expert computation + # without going through the runner (which consumes the output). + layer.router.select_experts(hidden_states, router_logits) + actual_zero_output = layer.router.zero_expert_output + + assert actual_zero_output is not None + assert zero_mask.any(), ( + "With high zero expert logit, at least some slots should route " + "to the zero expert" + ) + + torch.testing.assert_close( + actual_zero_output, + expected_zero_output, + atol=1e-3, + rtol=1e-3, + msg="Zero expert identity output should equal " + "hidden_states * sum(zero_expert_weights)", + ) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index 8763ad68351..c9c5c97b26d 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -69,6 +69,7 @@ def make_dummy_moe_config( in_dtype=in_dtype, device="cuda", routing_method=RoutingMethodType.TopK, + max_num_tokens=512, ) diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py index 045bc63de90..db00d74bd6a 100644 --- a/tests/kernels/quantization/test_mxfp4_triton_ep.py +++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py @@ -29,18 +29,22 @@ class TestTritonMoeForwardExpertMap: torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None ) + from vllm.utils.import_utils import import_triton_kernels + + import_triton_kernels() + with ( patch("triton_kernels.topk.topk") as mock_topk, patch( - "vllm.model_executor.layers.fused_moe." + "vllm.model_executor.layers.fused_moe.experts." "gpt_oss_triton_kernels_moe.make_routing_data" ) as mock_make_routing, patch( - "vllm.model_executor.layers.fused_moe." + "vllm.model_executor.layers.fused_moe.experts." "gpt_oss_triton_kernels_moe.triton_kernel_fused_experts" ) as mock_fused_experts, ): - from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501 + from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 triton_kernel_moe_forward, ) diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index e3b934722b9..5447a43ffb9 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -6,6 +6,7 @@ import pytest import torch from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils +from vllm.platforms import current_platform @pytest.mark.parametrize( @@ -48,6 +49,118 @@ def test_per_token_group_quant_fp8( assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01) +@pytest.mark.parametrize( + "num_tokens,hidden_dim,group_size", + [ + # No padding: mn=4 (mult of 4), groups_per_row=56 (mult of 4) + (4, 7168, 128), + # MN padding only: mn=1, tma_aligned_mn=4 + (1, 7168, 128), + # MN padding only: mn=3, tma_aligned_mn=4 + (3, 7168, 128), + # K padding only: groups_per_row=5 (5%4=1) + (4, 640, 128), + # K padding only: groups_per_row=6 (6%4=2) + (4, 768, 128), + # Single packed column, no padding: k_num_packed=1, mn%4=0 + (4, 384, 128), + # Both MN and K padding + (1, 384, 128), + (3, 640, 128), + # Larger shapes with no padding + (64, 7168, 128), + (128, 14336, 128), + # Larger shapes with padding + (127, 7168, 128), + (253, 640, 128), + # Non-power-of-2 group size + (4, 768, 96), # 768/96=8 groups, no padding + (3, 768, 96), # 768/96=8 groups, MN padding + (4, 480, 96), # 480/96=5 groups, K padding + (1, 480, 96), # both MN and K padding + ], +) +@pytest.mark.parametrize("poisoned_scales", [False, True]) +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" +) +def test_per_token_group_quant_fp8_packed( + num_tokens, hidden_dim, group_size, poisoned_scales +): + """Test the packed DeepGEMM quantization kernel against the Triton + reference (row-major, UE8M0 scales).""" + + device = "cuda" + torch.manual_seed(42) + + x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 + + mn = num_tokens + groups_per_row = hidden_dim // group_size + k_num_packed = (groups_per_row + 3) // 4 + tma_aligned_mn = ((mn + 3) // 4) * 4 + num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn + + if poisoned_scales: + # Call the kernel with poisoned scale buffer to + # ensure padded indices are correctly zeroed. + fp8_dtype = torch.float8_e4m3fn + finfo = torch.finfo(fp8_dtype) + out_q = torch.empty_like(x, dtype=fp8_dtype) + out_s_packed = torch.empty_strided( + (mn, k_num_packed), + (1, tma_aligned_mn), + device=device, + dtype=torch.int32, + ) + torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).fill_(0x7F7F7F7F) + torch.ops._C.per_token_group_fp8_quant_packed( + x, + out_q, + out_s_packed, + group_size, + 1e-10, + finfo.min, + finfo.max, + ) + else: + out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( + x, + group_size=group_size, + use_ue8m0=True, + ) + + # Triton reference (row-major float32 scales, UE8M0) + with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( + x, + group_size, + use_ue8m0=True, + ) + + # Quantized values must match. + assert torch.equal(out_q, ref_q), "Quantized output mismatch" + + # Verify packed scales (valid exponents + padding zeros). + ref_s_flat = ref_s.reshape(mn, groups_per_row) + ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF + + expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu") + for row in range(mn): + for g in range(groups_per_row): + pack_col = g // 4 + pos = g % 4 + idx = pack_col * tma_aligned_mn + row + expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8) + + actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu() + assert torch.equal(actual, expected), ( + f"Packed scale storage mismatch.\n" + f"First diff at index " + f"{(actual != expected).nonzero(as_tuple=True)[0][0].item()}" + ) + + @pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)]) @pytest.mark.parametrize("group_size", [64, 128]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") diff --git a/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py b/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py new file mode 100644 index 00000000000..a9b35a4ea64 --- /dev/null +++ b/tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end smoke test for CT W4A16 models on ROCm. + +This validates that a real compressed-tensors W4A16 model can run inference +end-to-end (which will exercise the Triton W4A16 kernel when selected). + +Run `pytest tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py`. +""" + +import pytest + +from vllm.platforms import current_platform + + +@pytest.mark.parametrize( + "model_path", + [ + # Listed in tests/weight_loading/models.txt + "nm-testing/tinyllama-oneshot-w4a16-group128-v2", + ], +) +@pytest.mark.parametrize("max_tokens", [32]) +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Should only run on ROCm") +def test_rocm_compressed_tensors_w4a16_e2e( + vllm_runner, example_prompts, model_path, max_tokens +): + # Use fp16 activations for maximum compatibility. + # gpu_memory_utilization lowered to work on shared nodes. + with vllm_runner( + model_path, dtype="float16", gpu_memory_utilization=0.3 + ) as vllm_model: + # If the W4A16 kernel is broken, this will typically throw. + vllm_model.generate_greedy(example_prompts, max_tokens=max_tokens) diff --git a/tests/kernels/quantization/test_triton_w4a16.py b/tests/kernels/quantization/test_triton_w4a16.py new file mode 100644 index 00000000000..6502f524429 --- /dev/null +++ b/tests/kernels/quantization/test_triton_w4a16.py @@ -0,0 +1,304 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the ROCm Triton W4A16 GEMM kernel. + +Run `pytest tests/kernels/quantization/test_triton_w4a16.py`. +""" + +import importlib + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +# This test module is ROCm/Triton specific. Avoid import-time failures on +# non-ROCm or environments without Triton by skipping early. +if not current_platform.is_rocm(): + pytest.skip("ROCm only", allow_module_level=True) + +pytest.importorskip("triton") + +device = "cuda" + +triton_w4a16_module = importlib.import_module( + "vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16" +) +triton_w4a16_gemm = triton_w4a16_module.triton_w4a16_gemm +TritonW4A16LinearKernel = triton_w4a16_module.TritonW4A16LinearKernel + + +def _pack_int4_along_n(w_int4_kn: torch.Tensor) -> torch.Tensor: + """Pack int4 values along N: [K, N] -> [K, N//8] int32.""" + assert w_int4_kn.dtype == torch.int32 + K, N = w_int4_kn.shape + assert N % 8 == 0 + shifts = torch.arange(8, device=w_int4_kn.device, dtype=torch.int32) * 4 + return torch.sum( + (w_int4_kn.view(K, N // 8, 8) & 0xF) << shifts, + dim=2, + dtype=torch.int32, + ).contiguous() + + +def _unpack_int4_along_n(w_packed_kn8: torch.Tensor) -> torch.Tensor: + """Unpack int4 values along N: [K, N//8] -> [K, N] int32.""" + assert w_packed_kn8.dtype == torch.int32 + K, N8 = w_packed_kn8.shape + shifts = torch.arange(8, device=w_packed_kn8.device, dtype=torch.int32) * 4 + nibbles = (w_packed_kn8.unsqueeze(-1) >> shifts) & 0xF + return nibbles.reshape(K, N8 * 8) + + +def _pack_int4_along_k_to_ckpt(w_int4_kn: torch.Tensor) -> torch.Tensor: + """Pack int4 values along K into CT checkpoint layout: [K,N] -> [N, K//8].""" + assert w_int4_kn.dtype == torch.int32 + K, N = w_int4_kn.shape + assert K % 8 == 0 + out = torch.zeros((N, K // 8), dtype=torch.int32, device=w_int4_kn.device) + for i in range(8): + out |= (w_int4_kn[i::8, :].t() & 0xF) << (i * 4) + return out.contiguous() + + +def _w4a16_reference( + a_mk: torch.Tensor, + b_packed_kn8: torch.Tensor, + scales_gn: torch.Tensor, + *, + group_size: int, + qzeros_gn8: torch.Tensor | None, + zp_bias: int, +) -> torch.Tensor: + """Reference implementation for W4A16. + + a_mk: [M,K] fp16/bf16 + b_packed_kn8: [K, N//8] int32, N-packed int4 weights + scales_gn: [K//G, N] fp16/bf16 + qzeros_gn8: [K//G, N//8] int32, N-packed int4 zeros, or None + """ + assert a_mk.dtype in (torch.float16, torch.bfloat16) + assert b_packed_kn8.dtype == torch.int32 + assert scales_gn.dtype == a_mk.dtype + + M, K = a_mk.shape + N = b_packed_kn8.shape[1] * 8 + assert b_packed_kn8.shape[0] == K + + assert group_size > 0 and K % group_size == 0 + G = group_size + num_groups = K // G + assert scales_gn.shape == (num_groups, N) + + w_int4 = _unpack_int4_along_n(b_packed_kn8) # [K,N] + if qzeros_gn8 is None: + z_full = torch.full((K, N), zp_bias, dtype=torch.int32, device=a_mk.device) + else: + assert qzeros_gn8.shape == (num_groups, N // 8) + z_gn = _unpack_int4_along_n(qzeros_gn8) # [G,N] in groups + z_full = z_gn.repeat_interleave(G, dim=0) # [K,N] + + s_full = scales_gn.repeat_interleave(G, dim=0).to(torch.float32) # [K,N] + w_fp = (w_int4 - z_full).to(torch.float32) * s_full # [K,N] + + out = a_mk.to(torch.float32) @ w_fp # [M,N] + return out.to(a_mk.dtype) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "M,K,N,G,has_zp", + [ + (1, 256, 256, 32, False), + (17, 256, 512, 32, False), + (32, 512, 256, 64, False), + (33, 512, 512, 128, False), + (64, 1024, 256, 256, False), + (128, 256, 1024, 32, True), + (64, 512, 512, 64, True), + ], +) +def test_triton_w4a16_gemm_matches_reference(dtype, M, K, N, G, has_zp): + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + if N % 8 != 0 or K % G != 0: + pytest.skip("Invalid test shape") + + set_random_seed(0) + + a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype) + w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + b_packed = _pack_int4_along_n(w_int4) + + scales = (0.05 * torch.rand((K // G, N), device=device, dtype=torch.float32)).to( + dtype + ) + + qzeros = None + if has_zp: + zeros_int4 = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32) + qzeros = _pack_int4_along_n(zeros_int4) + + out = triton_w4a16_gemm( + a=a, + b_q=b_packed, + scales=scales, + qzeros=qzeros, + group_size=G, + zp_bias=8, + ) + ref = _w4a16_reference( + a, + b_packed, + scales, + group_size=G, + qzeros_gn8=qzeros, + zp_bias=8, + ) + + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_triton_w4a16_gemm_requires_contiguous_inputs(): + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + + set_random_seed(0) + M, K, N, G = 32, 256, 256, 32 + a = torch.randn((K, M), device=device, dtype=torch.float16).t() # non-contiguous + w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + b_packed = _pack_int4_along_n(w_int4) + scales = torch.rand((K // G, N), device=device, dtype=torch.float16) + + with pytest.raises(AssertionError): + triton_w4a16_gemm( + a=a, + b_q=b_packed, + scales=scales, + qzeros=None, + group_size=G, + zp_bias=8, + ) + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_triton_w4a16_process_weights_after_loading_repacks_layout(): + if not torch.cuda.is_available(): + pytest.skip("CUDA/HIP device not available") + + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import ( + ensure_model_parallel_initialized, + init_distributed_environment, + ) + from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( + MPLinearLayerConfig, + ) + from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedColumnParameter, + PackedvLLMParameter, + ) + from vllm.scalar_type import scalar_types + + with set_current_vllm_config(VllmConfig()): + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method="tcp://127.0.0.1:0", + local_rank=0, + ) + ensure_model_parallel_initialized(1, 1) + + set_random_seed(0) + + # Small-but-nontrivial shapes. + K, N = 256, 256 + G = 32 + assert K % 8 == 0 and N % 8 == 0 and K % G == 0 + + # Build a canonical int4 weight grid then pack into the CT checkpoint layout. + w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32) + w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn) # [N, K//8] + + # Scales in CT checkpoint layout for WNA16: [N, K//G] + scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16) + + # Asymmetric case: zero points in CT checkpoint layout [N//8, K//G] (N-packed) + zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32) + zeros_packed_gn8 = _pack_int4_along_n(zeros_int4_gn) # [K//G, N//8] + zeros_ckpt_n8kg = zeros_packed_gn8.t().contiguous() # [N//8, K//G] + + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=scalar_types.uint4, # asymmetric + act_type=torch.float16, + group_size=G, + zero_points=True, + has_g_idx=False, + ) + kernel = TritonW4A16LinearKernel( + config, + w_q_param_name="weight_packed", + w_s_param_name="weight_scale", + w_zp_param_name="weight_zero_point", + w_gidx_param_name=None, + ) + + # Build dummy layer with vLLM parameter wrappers. + weight_loader = lambda *args, **kwargs: None + + class DummyLayer(torch.nn.Module): + pass + + layer = DummyLayer() + layer.register_parameter( + "weight_packed", + PackedvLLMParameter( + data=w_ckpt_nk8, + weight_loader=weight_loader, + input_dim=1, + output_dim=0, + packed_factor=8, + packed_dim=1, + ), + ) + layer.register_parameter( + "weight_scale", + GroupQuantScaleParameter( + data=scales_ckpt_nkg, + weight_loader=weight_loader, + input_dim=1, + output_dim=0, + ), + ) + layer.register_parameter( + "weight_zero_point", + PackedColumnParameter( + data=zeros_ckpt_n8kg, + weight_loader=weight_loader, + output_dim=0, + packed_factor=8, + packed_dim=0, + ), + ) + + kernel.process_weights_after_loading(layer) + + # Expected transformed layouts. + expected_w_kn8 = _pack_int4_along_n(w_int4_kn) # [K, N//8] + expected_scales_gn = scales_ckpt_nkg.t().contiguous() # [K//G, N] + expected_zeros_gn8 = zeros_ckpt_n8kg.t().contiguous() # [K//G, N//8] + + assert tuple(layer.weight_packed.shape) == (K, N // 8) + assert tuple(layer.weight_scale.shape) == (K // G, N) + assert tuple(layer.weight_zero_point.shape) == (K // G, N // 8) + + torch.testing.assert_close(layer.weight_packed, expected_w_kn8) + torch.testing.assert_close(layer.weight_scale, expected_scales_gn) + torch.testing.assert_close(layer.weight_zero_point, expected_zeros_gn8) diff --git a/tests/kernels/quantization/test_w4a16_kernel_selection.py b/tests/kernels/quantization/test_w4a16_kernel_selection.py new file mode 100644 index 00000000000..f0696191d3f --- /dev/null +++ b/tests/kernels/quantization/test_w4a16_kernel_selection.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for W4A16 kernel selection logic (ROCm). + +Run `pytest tests/kernels/quantization/test_w4a16_kernel_selection.py`. +""" + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + MPLinearLayerConfig, + choose_mp_linear_kernel, +) +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4b8(): + # int4 weights, 16-bit activations (CT W4A16 typical config). + K, N = 1024, 256 + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=scalar_types.uint4b8, # symmetric int4 (bias=8) + act_type=torch.float16, + group_size=128, + zero_points=False, + has_g_idx=False, + ) + + kernel_type = choose_mp_linear_kernel(config) + assert kernel_type.__name__ == "TritonW4A16LinearKernel" + + +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") +def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4_asymmetric(): + # Asymmetric int4 weights should also be supported (explicit zero points). + K, N = 512, 512 + config = MPLinearLayerConfig( + full_weight_shape=(K, N), + partition_weight_shape=(K, N), + weight_type=scalar_types.uint4, # asymmetric int4 (explicit zeros) + act_type=torch.bfloat16, + group_size=64, + zero_points=True, + has_g_idx=False, + ) + + kernel_type = choose_mp_linear_kernel(config) + assert kernel_type.__name__ == "TritonW4A16LinearKernel" diff --git a/tests/kernels/test_fused_quant_activation.py b/tests/kernels/test_fused_quant_activation.py index 2670f224d7c..0696ebb8d55 100644 --- a/tests/kernels/test_fused_quant_activation.py +++ b/tests/kernels/test_fused_quant_activation.py @@ -7,6 +7,7 @@ import vllm._custom_ops as ops from tests.kernels.utils import opcheck from vllm.model_executor.layers.activation import SiluAndMul from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed DTYPES = [torch.bfloat16, torch.float16] QUANT_DTYPES = [current_platform.fp8_dtype()] @@ -49,9 +50,7 @@ def test_silu_and_mul( seed: int, device: str, ) -> None: - torch.random.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) + set_random_seed(seed) torch.set_default_device(device) layer = SiluAndMul() diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index b97a9a0ea27..20944a9111e 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -43,6 +43,13 @@ def cleanup_fixture(should_do_global_cleanup_after_test: bool): cleanup_dist_env_and_memory(shutdown_ray=True) +@pytest.fixture +def maybe_enable_lora_dual_stream(monkeypatch: pytest.MonkeyPatch): + if current_platform.is_cuda(): + monkeypatch.setenv("VLLM_LORA_ENABLE_DUAL_STREAM", "1") + yield + + @pytest.fixture def dist_init(): from tests.utils import ensure_current_vllm_config diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 2a37abac6d7..c2b4f551564 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -521,8 +521,10 @@ def test_linear_replicated( punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) assert check_punica_wrapper(punica_wrapper) - def create_random_linear_replicated_layer(): - linear = ReplicatedLinear(4096, 4096, bias=False, params_dtype=torch.float16) + def create_random_linear_replicated_layer(idx: int = 0): + linear = ReplicatedLinear( + 4096, 4096, bias=False, params_dtype=torch.float16, prefix=f"layer_{idx}" + ) linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ReplicatedLinearWithLoRA(linear) @@ -539,7 +541,7 @@ def test_linear_replicated( set_random_seed(i) id_to_index = get_random_id_to_index(num_loras, max_loras) - linear, lora_linear = create_random_linear_replicated_layer() + linear, lora_linear = create_random_linear_replicated_layer(i) assert torch.equal(linear.weight, lora_linear.weight) lora_linear.set_mapping(punica_wrapper) lora_dict, _ = populate_loras( @@ -629,10 +631,14 @@ def test_linear_parallel( punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) assert check_punica_wrapper(punica_wrapper) - def create_random_linear_parallel_layer(): + def create_random_linear_parallel_layer(idx: int = 0): if orientation == "row": linear = RowParallelLinear( - 4096, 4096, bias=False, params_dtype=torch.float16 + 4096, + 4096, + bias=False, + params_dtype=torch.float16, + prefix=f"layer_{idx}", ) linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ( @@ -642,7 +648,11 @@ def test_linear_parallel( ) else: linear = ColumnParallelLinear( - 4096, 4096, bias=False, params_dtype=torch.float16 + 4096, + 4096, + bias=False, + params_dtype=torch.float16, + prefix=f"layer_{idx}", ) linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ( @@ -664,7 +674,7 @@ def test_linear_parallel( set_random_seed(i) id_to_index = get_random_id_to_index(num_loras, max_loras) - linear, lora_linear = create_random_linear_parallel_layer() + linear, lora_linear = create_random_linear_parallel_layer(i) assert torch.equal(linear.weight, lora_linear.weight) lora_linear.set_mapping(punica_wrapper) lora_dict, _ = populate_loras( @@ -754,10 +764,14 @@ def test_column_parallel_packed( punica_wrapper = get_punica_wrapper(8192, 256, device, lora_config=lora_config) assert check_punica_wrapper(punica_wrapper) - def create_column_parallel_packed_layer(): + def create_column_parallel_packed_layer(idx: int = 0): if repeats == 2: linear = MergedColumnParallelLinear( - 4096, [4096] * repeats, bias=False, params_dtype=torch.float16 + 4096, + [4096] * repeats, + bias=False, + params_dtype=torch.float16, + prefix=f"layer_{idx}", ) linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ( @@ -767,7 +781,12 @@ def test_column_parallel_packed( ) elif repeats == 3: linear = QKVParallelLinear( - 4096, 64, 32, bias=False, params_dtype=torch.float16 + 4096, + 64, + 32, + bias=False, + params_dtype=torch.float16, + prefix=f"layer_{idx}", ) linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ( @@ -777,7 +796,12 @@ def test_column_parallel_packed( ) else: linear = QKVParallelLinear( - 4096, 64, 32, bias=False, params_dtype=torch.float16 + 4096, + 64, + 32, + bias=False, + params_dtype=torch.float16, + prefix=f"layer_{idx}", ) linear.weight.data = torch.rand_like(linear.weight.data) lora_linear = ( @@ -810,7 +834,7 @@ def test_column_parallel_packed( id_to_index = get_random_id_to_index(num_loras, max_loras) - linear, lora_linear = create_column_parallel_packed_layer() + linear, lora_linear = create_column_parallel_packed_layer(i) assert torch.equal(linear.weight, lora_linear.weight) lora_linear.set_mapping(punica_wrapper) lora_dict, sublora_dict = populate_loras( @@ -902,10 +926,14 @@ def test_merged_column_parallel_variable_slice( output_sizes = [1024 + i * 256 for i in range(num_slices)] total_output = sum(output_sizes) - def create_layer(): + def create_layer(idx: int = 0): # Create linear layer linear = MergedColumnParallelLinear( - 4096, output_sizes, bias=False, params_dtype=torch.float16 + 4096, + output_sizes, + bias=False, + params_dtype=torch.float16, + prefix=f"layer_{idx}", ) linear.weight.data = torch.rand_like(linear.weight.data) @@ -917,7 +945,7 @@ def test_merged_column_parallel_variable_slice( for i in range(NUM_RANDOM_SEEDS): set_random_seed(i) id_to_index = get_random_id_to_index(num_loras, max_loras) - linear, lora_linear = create_layer() + linear, lora_linear = create_layer(i) lora_linear.set_mapping(punica_wrapper) # Populate LoRA weights diff --git a/tests/lora/test_lora_checkpoints.py b/tests/lora/test_lora_checkpoints.py index e6816e83da0..7c263e2a227 100644 --- a/tests/lora/test_lora_checkpoints.py +++ b/tests/lora/test_lora_checkpoints.py @@ -5,7 +5,9 @@ import pytest from vllm.lora.lora_model import LoRAModel from vllm.lora.peft_helper import PEFTHelper +from vllm.lora.utils import parse_fine_tuned_lora_name from vllm.model_executor.models.baichuan import BaiChuanBaseForCausalLM +from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM from vllm.model_executor.models.utils import WeightsMapper lora_lst = ["baichuan7B", "baichuan7B-zero", "baichuan7B-zero-regex", "chatglm3-6b"] @@ -128,3 +130,24 @@ def test_lora_weights_mapping(baichuan_lora_files): for name in lora_model.loras: assert name.startswith(hf_to_vllm_mapper.orig_to_new_prefix["model."]) assert ".baichuan_layers." in name + + +def test_gemma4_lora_weights_mapping(): + mapper = Gemma4ForCausalLM.hf_to_vllm_mapper + name = "base_model.model.model.language_model.layers.9.mlp.down_proj.lora_A.weight" + assert parse_fine_tuned_lora_name(name, mapper) == ( + "model.layers.9.mlp.down_proj", + True, + ) + + +def test_gemma4_moe_lora_weights_mapping(): + mapper = Gemma4ForCausalLM.hf_to_vllm_mapper + name = ( + "base_model.model.model.language_model.layers.9.moe.experts." + "gate_up_proj.lora_B.weight" + ) + assert parse_fine_tuned_lora_name(name, mapper) == ( + "model.layers.9.moe.gate_up_proj", + False, + ) diff --git a/tests/lora/test_olmoe_tp.py b/tests/lora/test_olmoe_tp.py index 492716b4645..0b477062205 100644 --- a/tests/lora/test_olmoe_tp.py +++ b/tests/lora/test_olmoe_tp.py @@ -110,7 +110,7 @@ def generate_and_test( ) -def test_olmoe_lora(olmoe_lora_files): +def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream): # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI, # Otherwise, the lora-test will fail due to CUDA OOM. llm = vllm.LLM( @@ -141,7 +141,9 @@ def test_olmoe_lora_mixed(olmoe_lora_files): generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None]) -def test_olmoe_lora_mixed_random(olmoe_lora_files, tmp_path): +def test_olmoe_lora_mixed_random( + olmoe_lora_files, tmp_path, maybe_enable_lora_dual_stream +): # Create a dummy LoRA with random weights based on the real one random_lora_path = tmp_path / "random_lora" shutil.copytree(olmoe_lora_files, random_lora_path) diff --git a/tests/lora/test_qwen35_densemodel_lora.py b/tests/lora/test_qwen35_densemodel_lora.py index 665fb99de0f..a9ee5fac8cb 100644 --- a/tests/lora/test_qwen35_densemodel_lora.py +++ b/tests/lora/test_qwen35_densemodel_lora.py @@ -312,7 +312,9 @@ def _assert_qwen35_text_vl_and_mixed_lora( @create_new_process_for_each_test() -def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files): +def test_qwen35_text_lora( + qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream +): llm = vllm.LLM( model=MODEL_PATH, max_model_len=4096, @@ -335,7 +337,9 @@ def test_qwen35_text_lora(qwen35_text_lora_files, qwen35_vl_lora_files): @multi_gpu_test(num_gpus=4) -def test_qwen35_text_lora_tp4(qwen35_text_lora_files, qwen35_vl_lora_files): +def test_qwen35_text_lora_tp4( + qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream +): llm = vllm.LLM( model=MODEL_PATH, max_model_len=4096, diff --git a/tests/model_executor/test_eagle_quantization.py b/tests/model_executor/test_eagle_quantization.py index 519a48cae52..481715da9cd 100644 --- a/tests/model_executor/test_eagle_quantization.py +++ b/tests/model_executor/test_eagle_quantization.py @@ -10,9 +10,10 @@ from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig from vllm.model_executor.models.utils import get_draft_quant_config from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type DEVICES = ( - [f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)] - if current_platform.is_cuda_alike() + [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))] + if not current_platform.is_cpu() else ["cpu"] ) diff --git a/tests/model_executor/test_ernie45_vl_mrope.py b/tests/model_executor/test_ernie45_vl_mrope.py new file mode 100644 index 00000000000..8344115ba34 --- /dev/null +++ b/tests/model_executor/test_ernie45_vl_mrope.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass + +import pytest +import torch + +from vllm.model_executor.models.ernie45_vl import ( + Ernie4_5_VLMoeForConditionalGeneration, +) +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldElem, + MultiModalKwargsItem, + PlaceholderRange, +) + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True, scope="module") +def _force_cpu_default_device(): + original = torch.get_default_device() + torch.set_default_device("cpu") + yield + torch.set_default_device(original) + + +@dataclass +class DummyConfig: + spatial_conv_size: int = 2 + temporal_conv_size: int = 2 + + +def make_model(config: DummyConfig) -> Ernie4_5_VLMoeForConditionalGeneration: + model = object.__new__(Ernie4_5_VLMoeForConditionalGeneration) + model.config = config + return model + + +def make_mm_feature( + *, + modality: str, + offset: int, + length: int, + grid_thw: tuple[int, int, int], +) -> MultiModalFeatureSpec: + field_name = "image_grid_thw" if modality == "image" else "video_grid_thw" + return MultiModalFeatureSpec( + data=MultiModalKwargsItem( + { + field_name: MultiModalFieldElem( + data=torch.tensor(grid_thw), + field=None, # HACK. + ), + } + ), + modality=modality, + identifier="DUMMY", + mm_position=PlaceholderRange(offset=offset, length=length), + ) + + +def test_get_mrope_input_positions_text_only(): + model = make_model(DummyConfig()) + + positions, delta = model.get_mrope_input_positions( + input_tokens=[11, 12, 13, 14, 15], + mm_features=[], + ) + + expected = torch.tensor( + [ + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + [0, 1, 2, 3, 4], + ] + ) + + assert torch.equal(positions, expected) + assert delta == 0 + + +def test_get_mrope_input_positions_single_image(): + model = make_model(DummyConfig()) + mm_features = [ + make_mm_feature( + modality="image", + offset=1, + length=4, + grid_thw=(1, 4, 4), + ) + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=[10, 20, 21, 22, 23, 30, 31], + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4], + [0, 1, 1, 2, 2, 3, 4], + [0, 1, 2, 1, 2, 3, 4], + ] + ) + + assert torch.equal(positions, expected) + assert delta == -2 + + +def test_get_mrope_input_positions_interleaved_image_and_video(): + model = make_model(DummyConfig()) + mm_features = [ + make_mm_feature( + modality="image", + offset=1, + length=4, + grid_thw=(1, 4, 4), + ), + make_mm_feature( + modality="video", + offset=7, + length=2, + grid_thw=(2, 4, 2), + ), + ] + + positions, delta = model.get_mrope_input_positions( + input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 50, 51], + mm_features=mm_features, + ) + + expected = torch.tensor( + [ + [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 8], + [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8], + [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 8], + ] + ) + + assert torch.equal(positions, expected) + assert delta == -2 diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index f831c8dfcd1..770a3fa5385 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -1,16 +1,40 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import types +from types import SimpleNamespace +from unittest.mock import patch import pytest import torch from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter pytestmark = pytest.mark.cpu_test +_REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer" + + +def _capturer_with_buffer( + *, + max_tokens: int = 8, + num_layers: int = 4, + num_experts_per_tok: int = 2, + dp_rank: int = 0, +) -> RoutedExpertsCapturer: + c = RoutedExpertsCapturer() + c.dp_rank = dp_rank + c._device_buffer = torch.full( + (max_tokens, num_layers, num_experts_per_tok), + -1, + dtype=torch.int32, + ) + return c + class DummyRouter(BaseRouter): @property @@ -159,3 +183,61 @@ def test_gpu_model_runner_binding_stage(monkeypatch): assert callable(dummy_module.router.capture_fn) dummy_module.router.capture_fn(torch.tensor([[9, 10]])) assert len(capturer.calls) == 1 + + +def test_routed_experts_capturer_single_dp_no_metadata(): + """dp_metadata is None: capture writes the full topk_ids rows.""" + capturer = _capturer_with_buffer(dp_rank=0) + topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) + ctx = SimpleNamespace(dp_metadata=None) + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): + capturer.capture(layer_id=0, topk_ids=topk) + assert torch.equal(capturer._device_buffer[:3, 0, :], topk) + assert capturer._device_buffer[3, 0, 0].item() == -1 + + +def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): + """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" + capturer = _capturer_with_buffer(dp_rank=1) + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) + ctx = SimpleNamespace( + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ) + # Concatenated order: rank0 rows then rank1 rows. + topk = torch.tensor( + [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32 + ) + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): + capturer.capture(layer_id=0, topk_ids=topk) + want = topk[2:5] + assert torch.equal(capturer._device_buffer[:3, 0, :], want) + + +def test_routed_experts_capturer_dp_modular_local_tokens(): + """n == token_num_per_dp: topk is already local to this DP rank.""" + capturer = _capturer_with_buffer(dp_rank=1) + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) + ctx = SimpleNamespace( + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ) + topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32) + with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): + capturer.capture(layer_id=0, topk_ids=topk) + assert torch.equal(capturer._device_buffer[:3, 0, :], topk) + + +def test_routed_experts_capturer_dp_unexpected_batch_raises(): + """Mismatch between topk batch dim and DP layout: fail fast.""" + capturer = _capturer_with_buffer(dp_rank=0) + num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) + ctx = SimpleNamespace( + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ) + # total=5, local=2: n=1 matches neither naive (5) nor modular (2). + topk = torch.tensor([[1, 2]], dtype=torch.int32) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + pytest.raises(AssertionError, match="unexpected topk_ids batch dim"), + ): + capturer.capture(layer_id=0, topk_ids=topk) + assert capturer._device_buffer[0, 0, 0].item() == -1 diff --git a/tests/models/language/pooling/test_jina_reranker_v3.py b/tests/models/language/pooling/test_jina_reranker_v3.py new file mode 100644 index 00000000000..dcce6d5bd4a --- /dev/null +++ b/tests/models/language/pooling/test_jina_reranker_v3.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 +import pytest +import requests +import torch +import torch.nn.functional as F + +from tests.utils import RemoteOpenAIServer +from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse +from vllm.entrypoints.pooling.scoring.protocol import ScoreResponse + +model_name = "jinaai/jina-reranker-v3" +query = "What are the health benefits of green tea?" +documents = [ + "Green tea contains antioxidants called catechins that may help reduce inflammation and protect cells from damage.", + "El precio del café ha aumentado un 20% este año debido a problemas en la cadena de suministro.", + "Studies show that drinking green tea regularly can improve brain function and boost metabolism.", + "Basketball is one of the most popular sports in the United States.", + "绿茶富含儿茶素等抗氧化剂，可以降低心脏病风险，还有助于控制体重。", + "Le thé vert est riche en antioxydants et peut améliorer la fonction cérébrale.", +] + +EMBEDDING_SIZE = 512 +REFERENCE_1_VS_1 = [ + 0.345703125, + -0.10498046, + 0.314453125, + -0.1376953125, + 0.3398437500, + 0.2539062, +] +REFERENCE_1_VS_N = [ + 0.294921875, + -0.16015625, + 0.189453125, + -0.1708984375, + 0.2255859375, + 0.1640625, +] +TOL = 0.01 + + +def test_offline(vllm_runner): + with vllm_runner(model_name, runner="pooling") as llm_runner: + llm = llm_runner.get_llm() + _test_offline_1_v_1(llm) + _test_offline_1_v_n(llm) + _test_offline_n_v_n(llm) + _test_offline_token_embed_illegal_inputs(llm) + assert llm.model_config.embedding_size == EMBEDDING_SIZE + + +def test_online(): + with RemoteOpenAIServer(model_name, ["--runner", "pooling"]) as server: + _test_online_1_v_1(server) + _test_online_1_v_n(server) + _test_online_n_v_n(server) + _test_online_token_embed_illegal_inputs(server) + + +def _test_offline_1_v_1(llm): + # test llm.score + outputs = llm.score(query, documents[0]) + assert len(outputs) == 1 + assert outputs[0].outputs.score == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + # test llm.encode + outputs = llm.encode(documents[:1] + [query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + assert embeds.shape[0] == 2 + assert embeds.shape[-1] == EMBEDDING_SIZE + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + +def _test_offline_1_v_n(llm): + # test llm.score + outputs = llm.score(query, documents) + assert len(outputs) == len(documents) + + for expected, output in zip(REFERENCE_1_VS_N, outputs): + actual = output.outputs.score + assert actual == pytest.approx(expected, abs=TOL) + + # test llm.encode + outputs = llm.encode(documents + [query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + assert embeds.shape[0] == len(documents) + 1 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + assert len(scores) == len(documents) + for expected, actual in zip(REFERENCE_1_VS_N, scores): + assert actual == pytest.approx(expected, abs=TOL) + + +def _test_offline_n_v_n(llm): + # test llm.score + outputs = llm.score([query] * len(documents), documents) + assert len(outputs) == len(documents) + + for expected, output in zip(REFERENCE_1_VS_1, outputs): + actual = output.outputs.score + assert actual == pytest.approx(expected, abs=TOL) + + # test llm.encode + for doc, expected in zip(documents, REFERENCE_1_VS_1): + outputs = llm.encode([doc, query], pooling_task="token_embed") + embeds = outputs[0].outputs.data.float() + assert embeds.shape[0] == 2 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert scores[0] == pytest.approx(expected, abs=TOL) + + +def _test_offline_token_embed_illegal_inputs(llm): + with pytest.raises( + ValueError, match="The JinaForRanking model requires at least 2 inputs." + ): + llm.encode([query], pooling_task="token_embed") + + with pytest.raises( + ValueError, match="The JinaForRanking model only supports text as input." + ): + llm.encode([1, 2, 3], pooling_task="token_embed") + + +def _get_scores(server, query, document): + score_response = requests.post( + server.url_for("score"), + json={ + "model": model_name, + "queries": query, + "documents": document, + }, + ) + + score_response.raise_for_status() + score = ScoreResponse.model_validate(score_response.json()) + + return [d.score for d in score.data] + + +def _get_embeds(server, prompts: list[str]): + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "input": prompts, + "encoding_format": "float", + }, + ) + response.raise_for_status() + poolings = PoolingResponse.model_validate(response.json()) + + return torch.as_tensor([d.data for d in poolings.data][0]).float() + + +def _test_online_1_v_1(server): + # test scoring api + scores = _get_scores(server, query, documents[0]) + assert len(scores) == 1 + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + # test pooling api + embeds = _get_embeds(server, [documents[0], query]) + assert embeds.shape[0] == 2 + assert embeds.shape[-1] == EMBEDDING_SIZE + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert scores[0] == pytest.approx(REFERENCE_1_VS_1[0], abs=TOL) + + +def _test_online_1_v_n(server): + # test scoring api + scores = _get_scores(server, query, documents) + assert len(scores) == len(documents) + + for expected, actual in zip(REFERENCE_1_VS_N, scores): + assert actual == pytest.approx(expected, abs=TOL) + + # test pooling api + embeds = _get_embeds(server, documents + [query]) + assert embeds.shape[0] == len(documents) + 1 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + assert len(scores) == len(documents) + for expected, actual in zip(REFERENCE_1_VS_N, scores): + assert actual == pytest.approx(expected, abs=TOL) + + +def _test_online_n_v_n(server): + # test scoring api + scores = _get_scores(server, [query] * len(documents), documents) + assert len(scores) == len(documents) + + for expected, actual in zip(REFERENCE_1_VS_1, scores): + assert actual == pytest.approx(expected, abs=TOL) + + # test pooling api + for doc, expected in zip(documents, REFERENCE_1_VS_1): + embeds = _get_embeds(server, [doc, query]) + assert embeds.shape[0] == 2 + + doc_embeds = embeds[:-1] + query_embeds = embeds[-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + assert len(scores) == 1 + assert scores[0] == pytest.approx(expected, abs=TOL) + + +def _test_online_token_embed_illegal_inputs(server): + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "input": [query], + "encoding_format": "float", + }, + ) + assert response.json()["error"]["message"].startswith( + "The JinaForRanking model requires at least 2 inputs." + ) + + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "input": [1, 2, 3], + "encoding_format": "float", + }, + ) + assert response.json()["error"]["message"].startswith( + "The JinaForRanking model only supports text as input." + ) + + response = requests.post( + server.url_for("pooling"), + json={ + "model": model_name, + "task": "token_embed", + "messages": [ + { + "role": "user", + "content": "The cat sat on the mat.", + } + ], + "encoding_format": "float", + }, + ) + assert response.json()["error"]["message"].startswith( + "The JinaForRanking does not support chat Request." + ) diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py new file mode 100644 index 00000000000..7b4e1fd0349 --- /dev/null +++ b/tests/models/language/pooling/test_max_tokens_per_doc.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for max_tokens_per_doc and max_tokens_per_query. +""" + +import json +import os +from dataclasses import dataclass + +import pytest +import requests + +from tests.utils import VLLM_PATH, RemoteOpenAIServer +from vllm.entrypoints.pooling.scoring.protocol import RerankResponse + +os.environ["VLLM_LOGGING_LEVEL"] = "WARNING" + +TEMPLATE_DIR = str(VLLM_PATH / "examples/pooling/score/template") + +long_query = "What is the capital of France?" * 20 +long_doc = "The capital of France is Paris. " * 20 + + +@dataclass +class TestConfig: + model: str + args: list[str] + without_truncated_prompt_tokens: int + with_max_tokens_per_query_prompt_tokens: int + with_max_tokens_per_doc_prompt_tokens: int + with_max_tokens_per_query_and_doc_prompt_tokens: int + + +RERANK_CONFIGS = [ + # 1. cross-encoder + TestConfig( + model="jinaai/jina-reranker-v2-base-multilingual", + args=[ + "--enforce-eager", + "--max-model-len", + "1024", + "--trust-remote-code", + ], + without_truncated_prompt_tokens=284, + with_max_tokens_per_query_prompt_tokens=154, + with_max_tokens_per_doc_prompt_tokens=154, + with_max_tokens_per_query_and_doc_prompt_tokens=24, + ), + # 2. cross-encoder + score template + TestConfig( + model="Qwen/Qwen3-Reranker-0.6B", + args=[ + "--enforce-eager", + "--max-model-len", + "1024", + "--hf-overrides", + json.dumps( + { + "architectures": ["Qwen3ForSequenceClassification"], + "classifier_from_token": ["no", "yes"], + "is_original_qwen3_reranker": True, + } + ), + "--chat-template", + os.path.join(TEMPLATE_DIR, "qwen3_reranker.jinja"), + ], + without_truncated_prompt_tokens=352, + with_max_tokens_per_query_prompt_tokens=223, + with_max_tokens_per_doc_prompt_tokens=221, + with_max_tokens_per_query_and_doc_prompt_tokens=92, + ), + # 3. bi-encoder + TestConfig( + model="intfloat/multilingual-e5-small", + args=[ + "--enforce-eager", + "--max-model-len", + "512", + "--trust-remote-code", + ], + without_truncated_prompt_tokens=286, + with_max_tokens_per_query_prompt_tokens=156, + with_max_tokens_per_doc_prompt_tokens=155, + with_max_tokens_per_query_and_doc_prompt_tokens=25, + ), + # 4. late-interaction + TestConfig( + model="answerdotai/answerai-colbert-small-v1", + args=[ + "--enforce-eager", + "--max-model-len", + "512", + "--trust-remote-code", + ], + without_truncated_prompt_tokens=285, + with_max_tokens_per_query_prompt_tokens=155, + with_max_tokens_per_doc_prompt_tokens=155, + with_max_tokens_per_query_and_doc_prompt_tokens=25, + ), + # 5. jinaai/jina-reranker-v3 + TestConfig( + model="jinaai/jina-reranker-v3", + args=[ + "--enforce-eager", + "--max-model-len", + "1024", + "--trust-remote-code", + ], + without_truncated_prompt_tokens=567, + with_max_tokens_per_query_prompt_tokens=308, + with_max_tokens_per_doc_prompt_tokens=436, + with_max_tokens_per_query_and_doc_prompt_tokens=177, + ), +] + + +@pytest.fixture(scope="module", params=RERANK_CONFIGS, ids=lambda c: c.model) +def server(request): + config: TestConfig = request.param + with RemoteOpenAIServer(config.model, config.args) as remote_server: + yield config, remote_server + + +def test_without_truncated(server): + """Test that max_tokens_per_doc truncates documents correctly.""" + config, remote_server = server + + response = requests.post( + remote_server.url_for("rerank"), + json={"model": config.model, "query": long_query, "documents": [long_doc]}, + ) + response.raise_for_status() + rerank = RerankResponse.model_validate(response.json()) + + assert rerank.id is not None + assert rerank.results is not None + assert len(rerank.results) == 1 + assert rerank.usage.prompt_tokens == config.without_truncated_prompt_tokens + + +def test_max_tokens_per_query(server): + """Test that max_tokens_per_doc truncates documents correctly.""" + config, remote_server = server + + response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": config.model, + "query": long_query, + "documents": [long_doc], + "max_tokens_per_query": 10, + }, + ) + response.raise_for_status() + rerank = RerankResponse.model_validate(response.json()) + + assert rerank.id is not None + assert rerank.results is not None + assert len(rerank.results) == 1 + assert rerank.usage.prompt_tokens == config.with_max_tokens_per_query_prompt_tokens + + +def test_max_tokens_per_doc(server): + """Test that max_tokens_per_doc truncates documents correctly.""" + config, remote_server = server + + response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": config.model, + "query": long_query, + "documents": [long_doc], + "max_tokens_per_doc": 10, + }, + ) + response.raise_for_status() + rerank = RerankResponse.model_validate(response.json()) + + assert rerank.id is not None + assert rerank.results is not None + assert len(rerank.results) == 1 + assert rerank.usage.prompt_tokens == config.with_max_tokens_per_doc_prompt_tokens + + +def test_max_tokens_per_query_and_doc(server): + """Test that max_tokens_per_doc truncates documents correctly.""" + config, remote_server = server + + response = requests.post( + remote_server.url_for("rerank"), + json={ + "model": config.model, + "query": long_query, + "documents": [long_doc], + "max_tokens_per_query": 10, + "max_tokens_per_doc": 10, + }, + ) + response.raise_for_status() + rerank = RerankResponse.model_validate(response.json()) + + assert rerank.id is not None + assert rerank.results is not None + assert len(rerank.results) == 1 + assert ( + rerank.usage.prompt_tokens + == config.with_max_tokens_per_query_and_doc_prompt_tokens + ) diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 42511f22f58..be71f7918ec 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -1,25 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import random -import numpy as np import pytest import torch from transformers import AutoModelForTokenClassification from tests.models.utils import softmax from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed @pytest.fixture(autouse=True) def seed_everything(): """Seed all random number generators for reproducibility.""" seed = 0 - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) + set_random_seed(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False yield diff --git a/tests/models/multimodal/generation/test_memory_leak.py b/tests/models/multimodal/generation/test_memory_leak.py new file mode 100644 index 00000000000..743a71f928f --- /dev/null +++ b/tests/models/multimodal/generation/test_memory_leak.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import gc +import random +import string +import sys +import weakref + +import pytest +import torch + +from tests.models.registry import HF_EXAMPLE_MODELS +from vllm import LLM, SamplingParams +from vllm.distributed import cleanup_dist_env_and_memory +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.platforms import current_platform +from vllm.utils.mem_utils import KiB_bytes, MiB_bytes, format_mib + +MODEL_NAME = "Qwen/Qwen3-VL-4B-Instruct" +RANDOM_PREFIX_LEN = 100 +TEST_IMAGE_NAMES = [ + "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", + "Grayscale_8bits_palette_sample_image.png", +] +MAX_MODEL_LEN = 8192 +REQUESTS_PER_ROUND = 4 +WARMUP_ROUNDS = 1 +MEASURED_ROUNDS = 16 +GPU_GROWTH_THRESHOLD_MIB = 0 +CPU_PEAK_GROWTH_THRESHOLD_MIB = 0 + +SAMPLING_PARAMS = SamplingParams( + temperature=0.0, + max_tokens=16, +) + + +def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]: + # Avoid obscuring memory leaks because of prefix caching + random_text = "".join(random.choices(string.ascii_uppercase, k=RANDOM_PREFIX_LEN)) + + return [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": f"Ignore this random string: {random_text}", + }, + {"type": "image_url", "image_url": {"url": image_url}}, + { + "type": "text", + "text": "Describe this image in one short sentence.", + }, + ], + } + ] + + +def _build_request_batch( + image_urls: list[str], +) -> list[list[ChatCompletionMessageParam]]: + return [ + _make_messages(image_urls[i % len(image_urls)]) + for i in range(REQUESTS_PER_ROUND) + ] + + +def _ru_maxrss_bytes() -> int | None: + try: + import resource + except ImportError: + return None + + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + if rss <= 0: + return 0 + + # Linux reports kilobytes, macOS reports bytes. + return rss if sys.platform == "darwin" else rss * KiB_bytes + + +def _gpu_used_bytes() -> int: + torch.accelerator.synchronize() + free_bytes, total_bytes = current_platform.mem_get_info() + return int(total_bytes - free_bytes) + + +def _format_mib(num_bytes: int | None) -> str: + if num_bytes is None: + return "n/a" + + return f"{format_mib(num_bytes)} MiB" + + +@pytest.fixture(scope="function") +def llm(monkeypatch): + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + # pytest caches the fixture so we use weakref.proxy to + # enable garbage collection + llm_kwargs = dict( + model=MODEL_NAME, + enforce_eager=True, + max_model_len=MAX_MODEL_LEN, + max_num_seqs=REQUESTS_PER_ROUND, + limit_mm_per_prompt={"image": 1}, + seed=0, + disable_log_stats=True, + gpu_memory_utilization=0.8, + ) + if current_platform.is_rocm(): + llm_kwargs["attention_backend"] = "TRITON_ATTN" + + llm = LLM(**llm_kwargs) + + yield weakref.proxy(llm) + + del llm + + cleanup_dist_env_and_memory() + + +@pytest.mark.core_model +@pytest.mark.parametrize("image_urls", [TEST_IMAGE_NAMES], indirect=True) +def test_no_memory_leak(llm, image_urls: list[str]) -> None: + model_info = HF_EXAMPLE_MODELS.find_hf_info(MODEL_NAME) + model_info.check_available_online(on_fail="skip") + model_info.check_transformers_version(on_fail="skip") + + request_batch = _build_request_batch(image_urls) + + # Establish a warmup baseline after model load and the first multimodal + # requests complete. Later rounds should remain near this steady state. + for _ in range(WARMUP_ROUNDS): + outputs = llm.chat(request_batch, sampling_params=SAMPLING_PARAMS) + assert len(outputs) == len(request_batch) + assert llm.llm_engine.get_num_unfinished_requests() == 0 + del outputs + + gc.collect() + warmup_gpu = _gpu_used_bytes() + warmup_cpu_peak = _ru_maxrss_bytes() + + post_warmup_gpu_samples: list[int] = [] + post_warmup_cpu_peak_samples: list[int] = [] + + for _ in range(MEASURED_ROUNDS): + outputs = llm.chat(request_batch, sampling_params=SAMPLING_PARAMS) + assert len(outputs) == len(request_batch) + assert llm.llm_engine.get_num_unfinished_requests() == 0 + del outputs + + gc.collect() + post_warmup_gpu_samples.append(_gpu_used_bytes()) + cpu_peak = _ru_maxrss_bytes() + if cpu_peak is not None: + post_warmup_cpu_peak_samples.append(cpu_peak) + + gpu_growth = max(post_warmup_gpu_samples) - warmup_gpu + gpu_threshold = GPU_GROWTH_THRESHOLD_MIB * MiB_bytes + + assert gpu_growth <= gpu_threshold, ( + "Qwen3-VL GPU memory kept growing after warmup. " + f"warmup_baseline={_format_mib(warmup_gpu)}, " + f"post_warmup_samples={[_format_mib(x) for x in post_warmup_gpu_samples]}, " + f"gpu_growth={_format_mib(gpu_growth)}, " + f"gpu_threshold={GPU_GROWTH_THRESHOLD_MIB} MiB" + ) + + if warmup_cpu_peak is not None and post_warmup_cpu_peak_samples: + cpu_peak_growth = max(post_warmup_cpu_peak_samples) - warmup_cpu_peak + cpu_threshold = CPU_PEAK_GROWTH_THRESHOLD_MIB * MiB_bytes + + assert cpu_peak_growth <= cpu_threshold, ( + "Qwen3-VL CPU peak RSS kept growing after warmup. " + f"warmup_ru_maxrss={_format_mib(warmup_cpu_peak)}, " + f"post_warmup_ru_maxrss={[_format_mib(x) for x in post_warmup_cpu_peak_samples]}, " # noqa: E501 + f"cpu_peak_growth={_format_mib(cpu_peak_growth)}, " + f"cpu_peak_threshold={CPU_PEAK_GROWTH_THRESHOLD_MIB} MiB" + ) diff --git a/tests/models/multimodal/pooling/test_colmodernvbert.py b/tests/models/multimodal/pooling/test_colmodernvbert.py index 6e9dce7ab2e..efeb3195b15 100644 --- a/tests/models/multimodal/pooling/test_colmodernvbert.py +++ b/tests/models/multimodal/pooling/test_colmodernvbert.py @@ -15,10 +15,6 @@ from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score MODEL_NAME = "ModernVBERT/colmodernvbert-merged" COLBERT_DIM = 128 DTYPE = "half" -# Fixme: -# Update colmodernvbert code to support the latest HF version -# and remove revision set. -REVISION = "4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee" # ----------------------------------------------------------------------- @@ -30,7 +26,6 @@ def test_colmodernvbert_text_token_embed(vllm_runner): """Text query produces per-token embeddings with shape (seq_len, 128).""" with vllm_runner( MODEL_NAME, - revision=REVISION, runner="pooling", dtype=DTYPE, enforce_eager=True, @@ -54,7 +49,6 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner): with vllm_runner( MODEL_NAME, - revision=REVISION, runner="pooling", dtype=DTYPE, enforce_eager=True, @@ -72,7 +66,6 @@ def test_colmodernvbert_text_late_interaction(vllm_runner): with vllm_runner( MODEL_NAME, - revision=REVISION, runner="pooling", dtype=DTYPE, enforce_eager=True, @@ -99,7 +92,6 @@ def test_colmodernvbert_image_token_embed(vllm_runner, image_assets): """Image input produces per-token embeddings including vision tokens.""" with vllm_runner( MODEL_NAME, - revision=REVISION, runner="pooling", dtype=DTYPE, enforce_eager=True, diff --git a/tests/models/multimodal/pooling/test_intern_vit.py b/tests/models/multimodal/pooling/test_intern_vit.py index cd457c62c0a..c3f7c81b78b 100644 --- a/tests/models/multimodal/pooling/test_intern_vit.py +++ b/tests/models/multimodal/pooling/test_intern_vit.py @@ -7,6 +7,7 @@ from huggingface_hub import snapshot_download from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory +from vllm.platforms import current_platform from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from ....conftest import ImageTestAssets @@ -15,6 +16,8 @@ from ....conftest import ImageTestAssets # dynamic_module and trust_remote_code for hf_runner DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] +DEVICE_TYPE = current_platform.device_type + @torch.inference_mode() def run_intern_vit_test( @@ -39,9 +42,9 @@ def run_intern_vit_test( hf_model = AutoModel.from_pretrained( model, dtype=torch_dtype, trust_remote_code=True - ).to("cuda") + ).to(DEVICE_TYPE) hf_outputs_per_image = [ - hf_model(pixel_value.to("cuda")).last_hidden_state + hf_model(pixel_value.to(DEVICE_TYPE)).last_hidden_state for pixel_value in pixel_values ] @@ -53,9 +56,10 @@ def run_intern_vit_test( del hf_model cleanup_dist_env_and_memory() - vllm_model = vllm_model.to("cuda", torch_dtype) + vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype) vllm_outputs_per_image = [ - vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values + vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE)) + for pixel_value in pixel_values ] del vllm_model cleanup_dist_env_and_memory() diff --git a/tests/models/multimodal/pooling/test_radio.py b/tests/models/multimodal/pooling/test_radio.py index 86b5b1b5d1f..fcab077fbba 100644 --- a/tests/models/multimodal/pooling/test_radio.py +++ b/tests/models/multimodal/pooling/test_radio.py @@ -8,6 +8,7 @@ from transformers import AutoConfig, AutoModel, CLIPImageProcessor from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.models.radio import RadioModel +from vllm.platforms import current_platform from vllm.transformers_utils.configs.radio import RadioConfig from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE @@ -17,6 +18,8 @@ from ....conftest import ImageTestAssets # dynamic_module and trust_remote_code for hf_runner DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"] +DEVICE_TYPE = current_platform.device_type + @torch.inference_mode() def run_radio_test( @@ -51,7 +54,7 @@ def run_radio_test( config=hf_config, dtype=torch_dtype, trust_remote_code=True, - ).to("cuda") + ).to(DEVICE_TYPE) hf_model.eval() # A HF model has image normalization as a part of model's forward @@ -62,7 +65,7 @@ def run_radio_test( hf_model.make_preprocessor_external() hf_outputs_per_image = [ - hf_model(pixel_value.to("cuda")) for pixel_value in pixel_values + hf_model(pixel_value.to(DEVICE_TYPE)) for pixel_value in pixel_values ] vllm_config = RadioConfig( @@ -71,10 +74,11 @@ def run_radio_test( ) vllm_model = RadioModel(vllm_config) vllm_model.load_weights(hf_model.state_dict()) - vllm_model = vllm_model.to("cuda", torch_dtype) + vllm_model = vllm_model.to(DEVICE_TYPE, torch_dtype) vllm_outputs_per_image = [ - vllm_model(pixel_values=pixel_value.to("cuda")) for pixel_value in pixel_values + vllm_model(pixel_values=pixel_value.to(DEVICE_TYPE)) + for pixel_value in pixel_values ] del vllm_model, hf_model cleanup_dist_env_and_memory() diff --git a/tests/models/registry.py b/tests/models/registry.py index 61f8958d11b..9c15decd8fd 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -416,6 +416,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "MiniMaxAI/MiniMax-M2", trust_remote_code=True, ), + "Ministral3ForCausalLM": _HfExamplesInfo("mistralai/Ministral-3-3B-Instruct-2512"), "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"), "MistralLarge3ForCausalLM": _HfExamplesInfo( "mistralai/Mistral-Large-3-675B-Instruct-2512-NVFP4" @@ -645,10 +646,10 @@ _LATE_INTERACTION_EXAMPLE_MODELS = { trust_remote_code=True, hf_overrides={"architectures": ["ColBERTLfm2Model"]}, ), + "JinaForRanking": _HfExamplesInfo("jinaai/jina-reranker-v3"), # [Multimodal] "ColModernVBertForRetrieval": _HfExamplesInfo( "ModernVBERT/colmodernvbert-merged", - revision="4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee", ), "ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"), "ColQwen3": _HfExamplesInfo( @@ -814,9 +815,16 @@ _MULTIMODAL_EXAMPLE_MODELS = { trust_remote_code=True, revision="refs/pr/17", ), + "Exaone4_5_ForConditionalGeneration": _HfExamplesInfo( + "LGAI-EXAONE/EXAONE-4.5-33B", + min_transformers_version="5.6.0", + ), "FireRedASR2ForConditionalGeneration": _HfExamplesInfo( "allendou/FireRedASR2-LLM-vllm", ), + "FireRedLIDForConditionalGeneration": _HfExamplesInfo( + "PatchyTisa/FireRedLID-vllm", + ), "FunASRForConditionalGeneration": _HfExamplesInfo( "allendou/Fun-ASR-Nano-2512-vllm", ), @@ -1030,6 +1038,50 @@ _MULTIMODAL_EXAMPLE_MODELS = { }, trust_remote_code=True, ), + # NemotronH_Nano_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 + # Use the same registry test as NemotronH_Nano_VL_V2 above + "NemotronH_Nano_Omni_Reasoning_V3": _HfExamplesInfo( + "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + max_model_len=4096, + use_original_num_layers=True, + hf_overrides={ + "vision_config": PretrainedConfig( + args={ + "min_num_patches": 1, + "max_num_patches": 12, + "model": "vit_huge_patch16_224", + }, + video_temporal_patch_size=2, + ), + "text_config": { + "num_hidden_layers": 2, + "hybrid_override_pattern": "M*", + }, + }, + trust_remote_code=True, + ), + # NemotronH_Super_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 as well + # Use the same registry test as NemotronH_Nano_VL_V2 above + "NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo( + "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16", + max_model_len=4096, + use_original_num_layers=True, + hf_overrides={ + "vision_config": PretrainedConfig( + args={ + "min_num_patches": 1, + "max_num_patches": 12, + "model": "vit_huge_patch16_224", + }, + video_temporal_patch_size=2, + ), + "text_config": { + "num_hidden_layers": 2, + "hybrid_override_pattern": "M*", + }, + }, + trust_remote_code=True, + ), "OpenCUAForConditionalGeneration": _HfExamplesInfo( "xlangai/OpenCUA-7B", trust_remote_code=True ), @@ -1307,6 +1359,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = { min_transformers_version="5.1.0", enable_prefix_caching=False, ), + "Exaone4_5_MTP": _HfExamplesInfo( + "LGAI-EXAONE/EXAONE-4.5-33B", + speculative_model="LGAI-EXAONE/EXAONE-4.5-33B", + min_transformers_version="5.6.0", + ), "ExtractHiddenStatesModel": _HfExamplesInfo( "Qwen/Qwen3-8B", speculative_method="extract_hidden_states", diff --git a/tests/models/test_adapters.py b/tests/models/test_adapters.py new file mode 100644 index 00000000000..7b1815998f2 --- /dev/null +++ b/tests/models/test_adapters.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for model adapter weight loading (adapters.py).""" + +import pytest +import torch + +from vllm.model_executor.models.adapters import _create_pooling_model_cls +from vllm.model_executor.models.utils import AutoWeightsLoader, StageMissingLayer + +pytestmark = pytest.mark.cpu_test + + +class SimpleInnerModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.embed = torch.nn.Linear(4, 8, bias=False) + self.layer0 = torch.nn.Linear(8, 8, bias=False) + self.layer1 = torch.nn.Linear(8, 8, bias=False) + self.norm = torch.nn.Linear(8, 4, bias=False) + + def load_weights(self, weights): + params = dict(self.named_parameters()) + loaded = set() + for name, tensor in weights: + if name in params: + params[name].data.copy_(tensor) + loaded.add(name) + return loaded + + +class SimpleModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = SimpleInnerModel() + self.lm_head = torch.nn.Linear(8, 16, bias=False) + + def load_weights(self, weights): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + +class PackedWeightInnerModel(torch.nn.Module): + """Remaps q_proj/k_proj into a fused qkv_proj (Qwen2/Llama pattern).""" + + def __init__(self): + super().__init__() + self.qkv_proj = torch.nn.Linear(4, 16, bias=False) + self.out = torch.nn.Linear(8, 4, bias=False) + + def load_weights(self, weights): + params = dict(self.named_parameters()) + loaded = set() + for name, tensor in weights: + if name == "q_proj.weight": + params["qkv_proj.weight"].data[:8].copy_(tensor) + loaded.add("qkv_proj.weight") + elif name == "k_proj.weight": + params["qkv_proj.weight"].data[8:].copy_(tensor) + loaded.add("qkv_proj.weight") + elif name in params: + params[name].data.copy_(tensor) + loaded.add(name) + return loaded + + +class PackedWeightModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = PackedWeightInnerModel() + self.lm_head = torch.nn.Linear(4, 8, bias=False) + + def load_weights(self, weights): + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + +def _buffer_reusing_iterator(weight_dict): + """Yield weights through a shared buffer overwritten each step. + + Mimics ``runai_model_streamer`` with ``RUNAI_STREAMER_MEMORY_LIMIT=0``. + """ + buf = None + for name, tensor in weight_dict.items(): + if buf is None or buf.numel() < tensor.numel(): + buf = torch.empty(tensor.numel(), dtype=tensor.dtype) + view = buf[: tensor.numel()].view(tensor.shape) + view.copy_(tensor) + yield name, view + + +def _make_pooling_model(base_cls=SimpleModel): + PoolingModel = _create_pooling_model_cls(base_cls) + model = base_cls() + model.__class__ = PoolingModel + model.lm_head = StageMissingLayer("output", model.lm_head) + return model + + +def _make_reference_weights(): + torch.manual_seed(42) + return { + "model.embed.weight": torch.randn(8, 4), + "model.layer0.weight": torch.randn(8, 8), + "model.layer1.weight": torch.randn(8, 8), + "model.norm.weight": torch.randn(4, 8), + "lm_head.weight": torch.randn(16, 8), + } + + +def _make_packed_reference_weights(): + torch.manual_seed(42) + return { + "model.q_proj.weight": torch.randn(8, 4), + "model.k_proj.weight": torch.randn(8, 4), + "model.out.weight": torch.randn(4, 8), + "lm_head.weight": torch.randn(8, 4), + } + + +def _load_and_compare(model, ref, expected): + for p in model.parameters(): + p.data.zero_() + model.load_weights(_buffer_reusing_iterator(ref)) + for name, param in model.named_parameters(): + assert torch.equal(param.data, expected[name]), name + + +def test_pooling_load_weights_with_buffer_reuse(): + """Ensure ModelForPooling.load_weights works with buffer-reusing iterators.""" + ref = _make_reference_weights() + + ground_truth = SimpleModel() + ground_truth.load_weights(ref.items()) + expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()} + + _load_and_compare(_make_pooling_model(), ref, expected) + + +def test_pooling_load_weights_clones_probed_weights(): + """Ensure probed weights survive buffer reuse during packed remapping.""" + ref = _make_packed_reference_weights() + + ground_truth = PackedWeightModel() + ground_truth.load_weights(ref.items()) + expected = {n: p.data.clone() for n, p in ground_truth.named_parameters()} + + _load_and_compare(_make_pooling_model(PackedWeightModel), ref, expected) diff --git a/tests/models/test_utils.py b/tests/models/test_utils.py index 3d719940e3a..8d47b443657 100644 --- a/tests/models/test_utils.py +++ b/tests/models/test_utils.py @@ -10,6 +10,8 @@ from vllm.model_executor.models.utils import ( ) from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + class ModuleWithBatchNorm(torch.nn.Module): def __init__(self): @@ -174,8 +176,12 @@ class raise_if_cuda_sync: @pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") def test_merge_multimodal_embeddings_no_sync(): - inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0") - multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")] + inputs_embeds = torch.zeros( + [5, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0" + ) + multimodal_embeddings = [ + torch.ones([3, 10], dtype=torch.bfloat16, device=f"{DEVICE_TYPE}:0") + ] is_multimodal = torch.tensor([True, False, True, True, False], device="cpu") with raise_if_cuda_sync(): _merge_multimodal_embeddings( diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index bd8c85e95bd..6b95d9e346d 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -28,6 +28,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso CompressedTensorsW4A16Fp4, CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8, + CompressedTensorsW8A8Mxfp8, CompressedTensorsW8A16Fp8, CompressedTensorsWNA16, ) @@ -632,3 +633,38 @@ def test_get_quant_method_returns_none_for_unmatched_parallel_lm_head(): assert method is None, ( f"Expected None for unmatched ParallelLMHead, got {type(method).__name__}" ) + + +@pytest.mark.skipif( + not current_platform.is_cuda() or not current_platform.has_device_capability(75), + reason="MXFP8 requires Turing (sm_75+) or newer.", +) +def test_compressed_tensors_mxfp8_moe_setup(vllm_runner): + """Verify MXFP8 scheme, dtypes, and generation for a MoE model.""" + model_path = "AliEdalati97/Qwen3-30B-A3B-MXFP8" + with vllm_runner( + model_path, + enforce_eager=True, + load_format="dummy", + hf_overrides={"num_hidden_layers": 4}, + ) as llm: + + def check_model(model): + from vllm.model_executor.layers.fused_moe import FusedMoE + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_mxfp8 import ( # noqa: E501 + CompressedTensorsW8A8Mxfp8MoEMethod, + ) + + layer = model.model.layers[0] + + qkv = layer.self_attn.qkv_proj + assert isinstance(qkv.quant_method, CompressedTensorsLinearMethod) + assert isinstance(qkv.scheme, CompressedTensorsW8A8Mxfp8) + + experts = layer.mlp.experts + assert isinstance(experts, FusedMoE) + assert isinstance(experts.quant_method, CompressedTensorsW8A8Mxfp8MoEMethod) + + llm.apply_model(check_model) + output = llm.generate_greedy("Hello my name is", max_tokens=4) + assert output diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py index 650bf714a07..5520dc1747a 100644 --- a/tests/quantization/test_cpu_wna16.py +++ b/tests/quantization/test_cpu_wna16.py @@ -12,6 +12,7 @@ MODELS = [ "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx "RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp + "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", ] DTYPE = ["bfloat16"] diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 4209d59ba28..4ba0e5d3dad 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -24,6 +24,8 @@ from vllm.model_executor.layers.quantization.fp8 import ( from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + MODELS = [ "neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV", # The checkpoint below was removed from the HF. @@ -314,7 +316,7 @@ def test_scaled_fp8_quant(dtype) -> None: # Note that we use a shape % 4 != 0 to cover edge cases, # because scaled_fp8_quant is vectorized by 4. - x = (torch.randn(size=(11, 11), device="cuda") * 13).to(dtype) + x = (torch.randn(size=(11, 11), device=DEVICE_TYPE) * 13).to(dtype) # Dynamic quantization ref_y, inv_scale = ops.scaled_fp8_quant(x, None) @@ -338,7 +340,9 @@ def test_scaled_fp8_quant(dtype) -> None: # non-contiguous input with padding m, n, padded_stride = 975, 512, 576 - padded_tensor = (torch.randn(size=(m, padded_stride), device="cuda") * 13).to(dtype) + padded_tensor = (torch.randn(size=(m, padded_stride), device=DEVICE_TYPE) * 13).to( + dtype + ) x_nc = padded_tensor[:, :n] # shape (m, n) with stride (padded_stride, 1) assert not x_nc.is_contiguous() @@ -409,7 +413,7 @@ def test_fp8_reloading( # Set model config as model_config.dtype is required in Fp8LinearMethod. default_vllm_config.model_config = ModelConfig() - with torch.device("cuda:0"): + with torch.device(f"{DEVICE_TYPE}:0"): config = Fp8Config( is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, weight_block_size=weight_block_size, diff --git a/tests/quantization/test_per_token_kv_cache.py b/tests/quantization/test_per_token_kv_cache.py index 3e660e6b00d..254e284efb5 100644 --- a/tests/quantization/test_per_token_kv_cache.py +++ b/tests/quantization/test_per_token_kv_cache.py @@ -25,11 +25,13 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from vllm.v1.kv_cache_interface import KVQuantMode, is_quantized_kv_cache +DEVICE_TYPE = current_platform.device_type + # Skip entire module if no CUDA/ROCm GPU available pytestmark = [ pytest.mark.skipif( - not current_platform.is_cuda_alike(), - reason="Per-token-head KV cache tests require CUDA or ROCm GPU.", + current_platform.is_cpu(), + reason="Per-token-head KV cache tests require GPU.", ), ] @@ -166,7 +168,7 @@ def test_reshape_and_cache_per_token_head( ) set_random_seed(seed) - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) num_blocks = (num_tokens + block_size - 1) // block_size + 4 @@ -260,7 +262,7 @@ def test_per_token_head_round_trip_accuracy( triton_reshape_and_cache_flash_per_token_head_quant, ) - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) set_random_seed(42) num_blocks = (num_tokens + block_size - 1) // block_size + 2 @@ -323,7 +325,7 @@ def test_per_token_head_negative_slot_skipped(qcfg: QuantConfig): triton_reshape_and_cache_flash_per_token_head_quant, ) - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) num_tokens = 4 num_heads = 2 head_size = 64 @@ -430,7 +432,7 @@ def test_triton_unified_attention_per_token_head_scale( from vllm.utils.math_utils import next_power_of_2 from vllm.v1.attention.ops.triton_unified_attention import unified_attention - torch.set_default_device("cuda") + torch.set_default_device(DEVICE_TYPE) set_random_seed(0) num_seqs = len(seq_lens) diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index afb0437f5b3..9eca6cda083 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -22,6 +22,9 @@ from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501 QuarkW8A8Fp8, QuarkW8A8Int8, ) +from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501 + QuarkW8A8Int8MoEMethod, +) from vllm.platforms import current_platform from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch @@ -33,6 +36,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse( importlib.metadata.version("amd-quark") ) >= version.parse(QUARK_MXFP4_MIN_VERSION) +DEVICE_TYPE = current_platform.device_type + if QUARK_MXFP4_AVAILABLE: from quark.torch.export.nn.modules.realquantizer import StaticScaledRealQuantizer from quark.torch.kernel import mx as mx_kernel @@ -126,6 +131,34 @@ def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp): assert output +@pytest.mark.parametrize("tp", [1]) +def test_quark_int8_w8a8_moe(vllm_runner, tp): + """Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model.""" + model_path = "nameistoken/tiny-qwen3-moe-w8a8-int8-quark" + with vllm_runner( + model_path, + enforce_eager=True, + tensor_parallel_size=tp, + gpu_memory_utilization=0.1, + ) as llm: + + def check_model(model): + layer = model.model.layers[0] + # MoE experts should use QuarkW8A8Int8MoEMethod + moe = layer.mlp.experts + assert isinstance(moe.quant_method, QuarkW8A8Int8MoEMethod), ( + f"Expected QuarkW8A8Int8MoEMethod, got {type(moe.quant_method)}" + ) + # Non-MoE linear layers should use QuarkW8A8Int8 + qkv_proj = layer.self_attn.qkv_proj + assert isinstance(qkv_proj.scheme, QuarkW8A8Int8) + + llm.apply_model(check_model) + + output = llm.generate_greedy("Hello", max_tokens=4) + assert output + + def test_quark_fp8_parity(vllm_runner): quark_model_id = "amd-quark/llama-tiny-fp8-quark-quant-method" fp8_model_id = "amd-quark/llama-tiny-fp8-quant-method" @@ -278,7 +311,7 @@ def test_mxfp4_fused_qdq_match_quark(float_dtype: torch.dtype, scalings: list[in torch.manual_seed(0) hidden_size = 64 * 32 - inp = (torch.rand(1, hidden_size, dtype=float_dtype, device="cuda") - 0.5) * 2 + inp = (torch.rand(1, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2 for i in range(hidden_size // 32): inp[:, i * 32 : (i + 1) * 32] = ( inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)] @@ -322,15 +355,15 @@ def test_mxfp4_dequant_kernel_match_quark( reorder=False, real_quantized=True, float_dtype=float_dtype, - device="cuda", + device=DEVICE_TYPE, ) - observer = qspec.observer_cls(qspec, device="cuda") + observer = qspec.observer_cls(qspec, device=DEVICE_TYPE) hidden_size = 512 shape = (11008, hidden_size) - w = (torch.rand(shape, device="cuda", dtype=float_dtype) - 0.5) * 2 + w = (torch.rand(shape, device=DEVICE_TYPE, dtype=float_dtype) - 0.5) * 2 # Make it so that different groups have different scales. for i in range(hidden_size // 32): @@ -342,7 +375,7 @@ def test_mxfp4_dequant_kernel_match_quark( scale, _ = observer._calculate_qparams() weight_quantizer.scale = scale - w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to("cuda") + w_mxfp4 = weight_quantizer.to_real_quantize_params(w).to(DEVICE_TYPE) weight_quantizer.maybe_convert_and_transpose_scale() scale = weight_quantizer.scale diff --git a/tests/quantization/test_torchao.py b/tests/quantization/test_torchao.py index fb794baa53f..8efc6742a2d 100644 --- a/tests/quantization/test_torchao.py +++ b/tests/quantization/test_torchao.py @@ -8,6 +8,7 @@ import torch from vllm.model_executor.model_loader import get_model_loader from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type DTYPE = ["bfloat16"] TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None @@ -33,7 +34,7 @@ def test_pre_quantized_model(vllm_runner): @pytest.mark.parametrize( "pt_load_map_location", [ - "cuda:0", + f"{DEVICE_TYPE}:0", # {"": "cuda"}, ], ) @@ -60,7 +61,7 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner): model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", enforce_eager=True, ) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=4) @@ -81,7 +82,7 @@ def test_opt_125m_awq_int4wo_model_loading_with_params(vllm_runner): model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", ) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=4) @@ -112,7 +113,7 @@ def test_online_quant_config_dict_json(vllm_runner, enable_pickle): with vllm_runner( model_name=model_name, dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", quantization="torchao", hf_overrides=hf_overrides, enforce_eager=True, @@ -158,7 +159,7 @@ def test_online_quant_config_file(vllm_runner): with vllm_runner( model_name=model_name, dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", quantization="torchao", hf_overrides=hf_overrides, enforce_eager=True, @@ -248,7 +249,7 @@ def test_opt_125m_module_fqn_to_config_regex_model(vllm_runner): torch._dynamo.reset() model_name = "torchao-testing/opt-125m-ModuleFqnToConfig-v1-regex-0.14.0.dev" with vllm_runner( - model_name=model_name, dtype="bfloat16", pt_load_map_location="cuda:0" + model_name=model_name, dtype="bfloat16", pt_load_map_location=f"{DEVICE_TYPE}:0" ) as llm: output = llm.generate_greedy(["The capital of France is"], max_tokens=4) @@ -278,7 +279,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel(vllm_runner, monkeypat model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", enforce_eager=True, ) as llm: @@ -357,7 +358,7 @@ def test_opt_125m_int4wo_model_running_preshuffled_kernel_online_quant( model_name=model_name, quantization="torchao", dtype="bfloat16", - pt_load_map_location="cuda:0", + pt_load_map_location=f"{DEVICE_TYPE}:0", hf_overrides=hf_overrides, enforce_eager=True, ) as llm: diff --git a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py new file mode 100644 index 00000000000..88c9e5f867c --- /dev/null +++ b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + align_trtllm_fp4_moe_hidden_dim_for_fi, +) + + +def test_align_trtllm_fp4_moe_hidden_dim_noop(): + w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256) + w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32) + w2 = torch.arange(2 * 512 * 4, dtype=torch.uint8).reshape(2, 512, 4) + w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1) + + out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = ( + align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) + ) + + assert padded_hidden == 512 + assert out_w13 is w13 + assert out_w13_scale is w13_scale + assert out_w2 is w2 + assert out_w2_scale is w2_scale + + +def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple(): + hidden_dim = 2688 + padded_hidden_dim = 2816 + + w13 = torch.arange(2 * 12 * (hidden_dim // 2), dtype=torch.uint8).reshape( + 2, 12, hidden_dim // 2 + ) + w13_scale = torch.arange(2 * 12 * (hidden_dim // 16), dtype=torch.uint8).reshape( + 2, 12, hidden_dim // 16 + ) + + w2 = torch.arange(2 * hidden_dim * 6, dtype=torch.uint8).reshape(2, hidden_dim, 6) + w2_scale = torch.arange(2 * hidden_dim * 2, dtype=torch.uint8).reshape( + 2, hidden_dim, 2 + ) + + out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = ( + align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) + ) + + assert out_hidden_dim == padded_hidden_dim + assert out_w13.shape == (2, 12, padded_hidden_dim // 2) + assert out_w13_scale.shape == (2, 12, padded_hidden_dim // 16) + assert out_w2.shape == (2, padded_hidden_dim, 6) + assert out_w2_scale.shape == (2, padded_hidden_dim, 2) + + torch.testing.assert_close(out_w13[:, :, : hidden_dim // 2], w13) + torch.testing.assert_close(out_w13_scale[:, :, : hidden_dim // 16], w13_scale) + torch.testing.assert_close(out_w2[:, :hidden_dim, :], w2) + torch.testing.assert_close(out_w2_scale[:, :hidden_dim, :], w2_scale) + + assert torch.count_nonzero(out_w13[:, :, hidden_dim // 2 :]) == 0 + assert torch.count_nonzero(out_w13_scale[:, :, hidden_dim // 16 :]) == 0 + assert torch.count_nonzero(out_w2[:, hidden_dim:, :]) == 0 + assert torch.count_nonzero(out_w2_scale[:, hidden_dim:, :]) == 0 diff --git a/tests/quantization/test_turboquant.py b/tests/quantization/test_turboquant.py new file mode 100644 index 00000000000..78c137e6762 --- /dev/null +++ b/tests/quantization/test_turboquant.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for TurboQuant KV-cache quantization. + +Run: .venv/bin/python -m pytest tests/quantization/test_turboquant.py -v +""" + +import math + +import pytest +import torch + +from vllm.model_executor.layers.quantization.turboquant.centroids import ( + get_centroids, + solve_lloyd_max, +) +from vllm.model_executor.layers.quantization.turboquant.config import ( + TQ_PRESETS, + TurboQuantConfig, +) +from vllm.model_executor.layers.quantization.turboquant.quantizer import ( + generate_wht_signs, +) +from vllm.utils.math_utils import next_power_of_2 + +# ============================================================================ +# Helpers +# ============================================================================ + +ALL_PRESETS = list(TQ_PRESETS.keys()) + + +def _assert_strictly_sorted(seq, name="sequence"): + for i in range(len(seq) - 1): + assert seq[i] < seq[i + 1], f"{name} not sorted at index {i}" + + +def _is_power_of_2(n: int) -> bool: + return n > 0 and next_power_of_2(n) == n + + +# Expected concrete values for each preset at head_dim=128. +# fmt: off +PRESET_EXPECTED = { + "turboquant_k8v4": dict( + key_fp8=True, key_quant_bits=8, + key_mse_bits=0, value_quant_bits=4, + mse_bits=4, n_centroids=16, centroid_bits=4, + norm_correction=False, + key_packed_size=128, value_packed_size=68, + slot_size=196, slot_size_aligned=196, + ), + "turboquant_4bit_nc": dict( + key_fp8=False, key_quant_bits=4, + key_mse_bits=4, value_quant_bits=4, + mse_bits=4, n_centroids=16, centroid_bits=4, + norm_correction=True, + key_packed_size=66, value_packed_size=68, + slot_size=134, slot_size_aligned=134, + ), + "turboquant_k3v4_nc": dict( + key_fp8=False, key_quant_bits=3, + key_mse_bits=3, value_quant_bits=4, + mse_bits=3, n_centroids=8, centroid_bits=3, + norm_correction=True, + key_packed_size=50, value_packed_size=68, + slot_size=118, slot_size_aligned=118, + ), + "turboquant_3bit_nc": dict( + key_fp8=False, key_quant_bits=3, + key_mse_bits=3, value_quant_bits=3, + mse_bits=3, n_centroids=8, centroid_bits=3, + norm_correction=True, + key_packed_size=50, value_packed_size=52, + slot_size=102, slot_size_aligned=102, + ), +} +# fmt: on + + +# ============================================================================ +# Config tests (CPU-only, no dependencies beyond config.py) +# ============================================================================ + + +class TestTurboQuantConfig: + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_preset_parses(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert isinstance(cfg, TurboQuantConfig) + + def test_invalid_preset_raises(self): + with pytest.raises(ValueError, match="Unknown TurboQuant"): + TurboQuantConfig.from_cache_dtype("turboquant_invalid", head_dim=128) + + # ---- Per-preset concrete value checks (table-driven) ---- + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_key_mode(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.key_fp8 is exp["key_fp8"] + assert cfg.key_quant_bits == exp["key_quant_bits"] + assert cfg.key_mse_bits == exp["key_mse_bits"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_value_mode(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.value_quant_bits == exp["value_quant_bits"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_bits_and_centroids(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.mse_bits == exp["mse_bits"] + assert cfg.n_centroids == exp["n_centroids"] + assert cfg.centroid_bits == exp["centroid_bits"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_norm_correction(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.norm_correction is PRESET_EXPECTED[preset]["norm_correction"] + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_packed_sizes(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + exp = PRESET_EXPECTED[preset] + assert cfg.key_packed_size == exp["key_packed_size"] + assert cfg.value_packed_size == exp["value_packed_size"] + assert cfg.slot_size == exp["slot_size"] + assert cfg.slot_size_aligned == exp["slot_size_aligned"] + + # ---- Cross-preset structural invariants ---- + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_slot_equals_key_plus_value(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_padded_slot_is_even(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.slot_size_aligned >= cfg.slot_size + assert cfg.slot_size_aligned % 2 == 0, ( + f"slot_size_aligned={cfg.slot_size_aligned} is not even" + ) + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_key_value_packed_sizes_positive(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.key_packed_size > 0 + assert cfg.value_packed_size > 0 + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_n_centroids_is_2_to_mse_bits(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.n_centroids == 2**cfg.mse_bits + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_centroid_bits_always_positive(self, preset): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + assert cfg.centroid_bits > 0 + + @pytest.mark.parametrize("preset", ALL_PRESETS) + def test_mse_key_or_fp8_exclusive(self, preset): + """Each preset is either FP8 keys or MSE keys, never both.""" + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + if cfg.key_fp8: + assert cfg.key_mse_bits == 0 + assert cfg.key_quant_bits == 8 + else: + assert cfg.key_mse_bits > 0 + assert cfg.key_quant_bits in (3, 4) + + @pytest.mark.parametrize("preset", ALL_PRESETS) + @pytest.mark.parametrize("head_dim", [64, 96, 128, 256]) + def test_all_presets_all_head_dims(self, preset, head_dim): + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=head_dim) + assert cfg.head_dim == head_dim + assert cfg.slot_size == cfg.key_packed_size + cfg.value_packed_size + assert cfg.slot_size_aligned >= cfg.slot_size + assert cfg.slot_size_aligned % 2 == 0 + + # ---- Boundary skip layers ---- + + def test_boundary_skip_layers_basic(self): + layers = TurboQuantConfig.get_boundary_skip_layers(32) + assert layers == ["0", "1", "30", "31"] + + def test_boundary_skip_layers_zero(self): + assert TurboQuantConfig.get_boundary_skip_layers(32, 0) == [] + + def test_boundary_skip_layers_small_model(self): + layers = TurboQuantConfig.get_boundary_skip_layers(4) + assert layers == ["0", "1", "2", "3"] + + def test_boundary_skip_layers_cap_at_half(self): + layers = TurboQuantConfig.get_boundary_skip_layers(8, 10) + assert len(layers) == 8 + + +# ============================================================================ +# Centroids tests (CPU-only) +# ============================================================================ + + +class TestCentroids: + @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)]) + def test_centroids_shape(self, bits, expected_n): + c = get_centroids(128, bits) + assert c.shape == (expected_n,) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_sorted(self, bits): + _assert_strictly_sorted(get_centroids(128, bits), "centroids") + + def test_centroids_cached(self): + c1 = get_centroids(128, 3) + c2 = get_centroids(128, 3) + assert c1 is c2, "get_centroids should return cached object" + + def test_centroids_different_dims_not_identical(self): + c64 = get_centroids(64, 3) + c128 = get_centroids(128, 3) + assert not torch.equal(c64, c128) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_symmetric_around_zero(self, bits): + """N(0, 1/d) is symmetric, so centroids should be ~symmetric.""" + c = get_centroids(128, bits) + assert abs(c.mean().item()) < 0.01, "Centroids not centered near 0" + assert abs(c[0].item() + c[-1].item()) < 0.01 + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_within_4sigma(self, bits): + """All centroids should be within ~4 sigma of N(0, 1/d).""" + sigma = math.sqrt(1.0 / 128) + c = get_centroids(128, bits) + for i, val in enumerate(c): + assert abs(val.item()) < 4 * sigma, ( + f"Centroid {i}={val:.6f} outside 4*sigma={4 * sigma:.6f}" + ) + + +class TestLloydMax: + @pytest.mark.parametrize("bits,expected_n", [(2, 4), (3, 8), (4, 16)]) + def test_solve_shapes(self, bits, expected_n): + centroids, boundaries = solve_lloyd_max(128, bits) + assert centroids.shape == (expected_n,) + assert boundaries.shape == (expected_n - 1,) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_centroids_sorted(self, bits): + centroids, _ = solve_lloyd_max(128, bits) + _assert_strictly_sorted(centroids, "centroids") + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_boundaries_sorted(self, bits): + _, boundaries = solve_lloyd_max(128, bits) + _assert_strictly_sorted(boundaries, "boundaries") + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_boundaries_between_centroids(self, bits): + """Each boundary must lie between its adjacent centroids.""" + centroids, boundaries = solve_lloyd_max(128, bits) + for i in range(len(boundaries)): + assert centroids[i] < boundaries[i] < centroids[i + 1], ( + f"Boundary {i}={boundaries[i]:.6f} not between " + f"c[{i}]={centroids[i]:.6f} and c[{i + 1}]={centroids[i + 1]:.6f}" + ) + + @pytest.mark.parametrize("bits", [2, 3, 4]) + def test_boundaries_are_midpoints(self, bits): + """Lloyd-Max boundaries are midpoints of adjacent centroids.""" + centroids, boundaries = solve_lloyd_max(128, bits) + for i in range(len(boundaries)): + expected = (centroids[i] + centroids[i + 1]) / 2.0 + assert abs(boundaries[i].item() - expected.item()) < 1e-6 + + def test_solve_deterministic(self): + c1, b1 = solve_lloyd_max(128, 3) + c2, b2 = solve_lloyd_max(128, 3) + assert torch.equal(c1, c2) + assert torch.equal(b1, b2) + + def test_solve_dtype_float32(self): + centroids, boundaries = solve_lloyd_max(128, 3) + assert centroids.dtype == torch.float32 + assert boundaries.dtype == torch.float32 + + @pytest.mark.parametrize("bits", [3, 4]) + def test_centroids_match_scipy_reference(self, bits): + """Verify _trapz(n=200) centroids match scipy.integrate.quad reference. + + This ensures our scipy-free trapezoid integration doesn't silently + drift from the published Lloyd-Max quality. + """ + pytest.importorskip("scipy") + from scipy.integrate import quad + + d = 128 + sigma2 = 1.0 / d + sigma = math.sqrt(sigma2) + + def pdf(x): + return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp( + -x * x / (2 * sigma2) + ) + + n_levels = 2**bits + lo, hi = -3.5 * sigma, 3.5 * sigma + ref_centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] + for _ in range(200): + boundaries = [ + (ref_centroids[i] + ref_centroids[i + 1]) / 2.0 + for i in range(n_levels - 1) + ] + edges = [lo * 3] + boundaries + [hi * 3] + new_centroids = [] + for i in range(n_levels): + a, b = edges[i], edges[i + 1] + num, _ = quad(lambda x: x * pdf(x), a, b) + den, _ = quad(pdf, a, b) + new_centroids.append(num / den if den > 1e-15 else ref_centroids[i]) + if ( + max(abs(new_centroids[i] - ref_centroids[i]) for i in range(n_levels)) + < 1e-10 + ): + break + ref_centroids = new_centroids + + # Compare our _trapz centroids against scipy reference + our_centroids, _ = solve_lloyd_max(d, bits) + ref_t = torch.tensor(ref_centroids, dtype=torch.float32) + max_err = (our_centroids - ref_t).abs().max().item() + # _trapz(n=200) has ~O(h^2) error vs adaptive quad; 1e-3 is tight + # enough to catch regression while allowing trapezoid approximation. + assert max_err < 1e-3, ( + f"d={d}, bits={bits}: max centroid error vs scipy = {max_err:.2e}" + ) + + +# ============================================================================ +# Rotation matrix tests (GPU required) +# ============================================================================ + +CUDA_AVAILABLE = torch.cuda.is_available() + + +def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor: + """Haar-distributed random orthogonal matrix via QR (test/benchmark only).""" + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + G = torch.randn(d, d, generator=gen, device="cpu", dtype=torch.float32) + Q, R = torch.linalg.qr(G) + diag_sign = torch.sign(torch.diag(R)) + diag_sign[diag_sign == 0] = 1.0 + Q = Q * diag_sign.unsqueeze(0) + return Q.to(device) + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +class TestRotationMatrix: + """Tests for the QR-based rotation (standalone benchmarks only).""" + + @pytest.mark.parametrize("dim", [64, 96, 128, 256]) + def test_rotation_matrix_shape_and_orthogonal(self, dim): + Pi = generate_rotation_matrix(dim, seed=42, device="cuda") + assert Pi.shape == (dim, dim) + eye = Pi @ Pi.T + assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), ( + f"Pi not orthogonal for dim={dim}" + ) + + def test_rotation_matrix_deterministic(self): + Pi1 = generate_rotation_matrix(128, seed=42) + Pi2 = generate_rotation_matrix(128, seed=42) + assert torch.equal(Pi1, Pi2) + + def test_rotation_matrix_different_seeds(self): + Pi1 = generate_rotation_matrix(128, seed=42) + Pi2 = generate_rotation_matrix(128, seed=99) + assert not torch.equal(Pi1, Pi2) + + def test_rotation_matrix_det_is_pm1(self): + """Orthogonal matrix determinant must be +1 or -1.""" + Pi = generate_rotation_matrix(128, seed=42, device="cuda") + det = torch.linalg.det(Pi) + assert abs(abs(det.item()) - 1.0) < 1e-4 + + +# ============================================================================ +# WHT rotation tests (serving path: generate_wht_signs + _build_hadamard) +# ============================================================================ + + +def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor: + """Reproduce the serving-path Hadamard construction.""" + H = torch.tensor([[1.0]]) + while H.shape[0] < d: + H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0) + return (H / math.sqrt(d)).to(torch.device(device)) + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +class TestWHTRotation: + """Tests for the WHT rotation actually used in serving.""" + + @pytest.mark.parametrize("dim", [64, 128, 256]) + def test_wht_orthonormal(self, dim): + """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I.""" + signs = generate_wht_signs(dim, seed=42, device="cuda") + H = _build_hadamard(dim, "cuda") + PiT = (signs.unsqueeze(1) * H).contiguous() + eye = PiT @ PiT.T + assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), ( + f"WHT rotation not orthonormal for dim={dim}" + ) + + @pytest.mark.parametrize("dim", [64, 128, 256]) + def test_wht_self_inverse(self, dim): + """PiT should be self-inverse: PiT @ PiT = I (up to sign flip).""" + signs = generate_wht_signs(dim, seed=42, device="cuda") + H = _build_hadamard(dim, "cuda") + PiT = (signs.unsqueeze(1) * H).contiguous() + Pi = PiT.T.contiguous() + # Pi @ PiT should be identity (rotation then inverse) + result = Pi @ PiT + assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), ( + f"WHT rotation not self-inverse for dim={dim}" + ) + + def test_wht_signs_deterministic(self): + """Same seed must produce identical signs.""" + s1 = generate_wht_signs(128, seed=42) + s2 = generate_wht_signs(128, seed=42) + assert torch.equal(s1, s2) + + def test_wht_signs_different_seeds(self): + """Different seeds must produce different signs.""" + s1 = generate_wht_signs(128, seed=42) + s2 = generate_wht_signs(128, seed=99) + assert not torch.equal(s1, s2) + + def test_wht_signs_are_pm1(self): + """All sign values must be exactly +1 or -1.""" + signs = generate_wht_signs(128, seed=42) + assert torch.all(signs.abs() == 1.0) + + +# ============================================================================ +# Store → Decode round-trip test (GPU + Triton required) +# ============================================================================ + + +@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available") +class TestStoreDecodeRoundTrip: + """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref.""" + + @pytest.mark.parametrize( + "preset", + ["turboquant_k8v4", "turboquant_4bit_nc"], + ) + def test_single_token_roundtrip(self, preset): + """Store 1 token, decode with query=key, check attention output. + + For a single token with query=key, attention output should equal + the value (softmax over single key = 1.0). Quantization error + means we check cosine similarity rather than exact equality. + """ + from vllm.model_executor.layers.quantization.turboquant.centroids import ( + solve_lloyd_max, + ) + from vllm.v1.attention.ops.triton_turboquant_decode import ( + triton_turboquant_decode_attention, + ) + from vllm.v1.attention.ops.triton_turboquant_store import ( + triton_turboquant_store, + ) + + cfg = TurboQuantConfig.from_cache_dtype(preset, head_dim=128) + D = 128 + Hk = 4 # num_kv_heads + Hq = 4 # num_q_heads (no GQA for simplicity) + B = 1 # single token + block_size = 16 + num_blocks = 1 + + device = torch.device("cuda") + + # Generate rotation + signs = generate_wht_signs(D, seed=42, device=device) + H = _build_hadamard(D, "cuda") + PiT = (signs.unsqueeze(1) * H).contiguous().float() + Pi = PiT.T.contiguous() + + # Generate centroids + centroids, _ = solve_lloyd_max(D, cfg.centroid_bits) + centroids = centroids.float().to(device) + c_sorted, _ = centroids.sort() + midpoints = ((c_sorted[:-1] + c_sorted[1:]) / 2).to(device) + + # Random K, V + torch.manual_seed(123) + key = torch.randn(B, Hk, D, device=device, dtype=torch.float16) + value = torch.randn(B, Hk, D, device=device, dtype=torch.float16) + + # Allocate KV cache + padded_slot = cfg.slot_size_aligned + kv_cache = torch.zeros( + num_blocks, + block_size, + Hk, + padded_slot, + device=device, + dtype=torch.uint8, + ) + slot_mapping = torch.tensor([0], device=device, dtype=torch.int32) + + # Store + triton_turboquant_store( + key, + value, + kv_cache, + slot_mapping, + PiT, + midpoints, + mse_bits=cfg.key_mse_bits, + key_packed_size=cfg.key_packed_size, + value_quant_bits=cfg.effective_value_quant_bits, + key_fp8=cfg.key_fp8, + ) + + # Decode: use key as query so attention = softmax([1]) * V = V + query = key.expand(B, Hq, D).contiguous().to(torch.float16) + block_table = torch.tensor([[0]], device=device, dtype=torch.int32) + seq_lens = torch.tensor([1], device=device, dtype=torch.int32) + + output = triton_turboquant_decode_attention( + query=query, + kv_cache=kv_cache, + block_table=block_table, + seq_lens=seq_lens, + Pi=Pi, + centroids=centroids, + scale=1.0 / math.sqrt(D), + mse_bits=cfg.key_mse_bits, + key_packed_size=cfg.key_packed_size, + value_quant_bits=cfg.effective_value_quant_bits, + key_fp8=cfg.key_fp8, + norm_correction=cfg.norm_correction, + PiT=PiT, + max_num_kv_splits=4, + ) + + # With single KV, output should approximate the stored value. + # Check per-head cosine similarity > threshold. + out_fp32 = output.float() + val_fp32 = value.expand(B, Hq, D).float() + for h in range(Hq): + cos_sim = torch.nn.functional.cosine_similarity( + out_fp32[0, h].unsqueeze(0), + val_fp32[0, h].unsqueeze(0), + ).item() + # FP8 keys should be very accurate; MSE keys have more error + threshold = 0.95 if cfg.key_fp8 else 0.85 + assert cos_sim > threshold, ( + f"Preset {preset} head {h}: cosine_sim={cos_sim:.4f} < {threshold}" + ) diff --git a/tests/standalone_tests/pytorch_nightly_dependency.sh b/tests/standalone_tests/pytorch_nightly_dependency.sh index 92820b269f9..3a23b1c824f 100644 --- a/tests/standalone_tests/pytorch_nightly_dependency.sh +++ b/tests/standalone_tests/pytorch_nightly_dependency.sh @@ -28,8 +28,8 @@ uv pip freeze | grep -E '^torch|^torchvision|^torchaudio' | sort > before.txt echo "Before:" cat before.txt -echo ">>> Installing requirements/nightly_torch_test.txt" -uv pip install --quiet -r requirements/nightly_torch_test.txt +echo ">>> Installing requirements/test/nightly-torch.txt" +uv pip install --quiet -r requirements/test/nightly-torch.txt echo ">>> Capturing torch-related versions after requirements install" uv pip freeze | grep -E '^torch|^torchvision|^torchaudio' | sort > after.txt @@ -40,7 +40,7 @@ echo ">>> Comparing versions" if diff before.txt after.txt; then echo "torch version not overridden." else - echo "torch version overridden by nightly_torch_test.txt, \ + echo "torch version overridden by test/nightly-torch.txt, \ if the dependency is not triggered by the pytorch nightly test,\ please add the dependency to the list 'white_list' in tools/pre_commit/generate_nightly_torch_test.py" exit 1 diff --git a/tests/test_config.py b/tests/test_config.py index 312f4ce5af2..41d34a6cb06 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4,12 +4,14 @@ import logging import os from dataclasses import MISSING, Field, asdict, dataclass, field +from types import SimpleNamespace from unittest.mock import patch import pydantic import pytest from pydantic import ValidationError +import vllm.config.vllm as vllm_config_module from vllm.compilation.backends import VllmBackend from vllm.config import ( CompilationConfig, @@ -32,6 +34,8 @@ from vllm.config.vllm import ( ) from vllm.platforms import current_platform +DEVICE_TYPE = current_platform.device_type + def test_compile_config_repr_succeeds(): # setup: VllmBackend mutates the config object @@ -45,6 +49,81 @@ def test_compile_config_repr_succeeds(): assert "inductor_passes" in val +@pytest.mark.skip_global_cleanup +def test_with_hf_config_populates_missing_architectures_from_causal_lm_mapping( + monkeypatch, +): + monkeypatch.setattr( + vllm_config_module, + "replace", + lambda self, **kwargs: SimpleNamespace(**kwargs), + ) + cfg = SimpleNamespace( + model_config=SimpleNamespace( + is_multimodal_model=False, + hf_config=SimpleNamespace(), + get_model_arch_config=lambda: "arch-config", + ) + ) + hf_config = SimpleNamespace(model_type="mistral", architectures=None) + + updated = VllmConfig.with_hf_config(cfg, hf_config) + + assert updated.model_config.hf_config.architectures == ["MistralForCausalLM"] + assert hf_config.architectures is None + + +@pytest.mark.skip_global_cleanup +def test_with_hf_config_preserves_explicit_architectures_override(monkeypatch): + monkeypatch.setattr( + vllm_config_module, + "replace", + lambda self, **kwargs: SimpleNamespace(**kwargs), + ) + cfg = SimpleNamespace( + model_config=SimpleNamespace( + is_multimodal_model=False, + hf_config=SimpleNamespace(), + get_model_arch_config=lambda: "arch-config", + ) + ) + hf_config = SimpleNamespace(model_type="mistral", architectures=None) + + updated = VllmConfig.with_hf_config( + cfg, + hf_config, + architectures=["Ministral3ForCausalLM"], + ) + + assert updated.model_config.hf_config.architectures == ["Ministral3ForCausalLM"] + + +@pytest.mark.skip_global_cleanup +def test_with_hf_config_leaves_unknown_model_type_without_architectures( + monkeypatch, +): + monkeypatch.setattr( + vllm_config_module, + "replace", + lambda self, **kwargs: SimpleNamespace(**kwargs), + ) + cfg = SimpleNamespace( + model_config=SimpleNamespace( + is_multimodal_model=False, + hf_config=SimpleNamespace(), + get_model_arch_config=lambda: "arch-config", + ) + ) + hf_config = SimpleNamespace( + model_type="not_a_real_model", + architectures=None, + ) + + updated = VllmConfig.with_hf_config(cfg, hf_config) + + assert updated.model_config.hf_config.architectures is None + + def test_async_scheduling_with_pipeline_parallelism_is_allowed(): cfg = VllmConfig( scheduler_config=SchedulerConfig( @@ -427,8 +506,8 @@ def test_generation_config_loading(): @pytest.mark.parametrize( "pt_load_map_location", [ - "cuda", - {"": "cuda"}, + DEVICE_TYPE, + {"": DEVICE_TYPE}, ], ) def test_load_config_pt_load_map_location(pt_load_map_location): diff --git a/tests/test_fxgraphcache_pickle_patch.py b/tests/test_fxgraphcache_pickle_patch.py new file mode 100644 index 00000000000..8a3f395267c --- /dev/null +++ b/tests/test_fxgraphcache_pickle_patch.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the FxGraphCachePickler.dumps ValueError patch in env_override.py. + +Validates that _apply_fxgraphcache_pickle_patch correctly wraps a pickler's +dumps method to convert ValueError into a bypass exception, without affecting +other exception types or normal return values. +""" + +import pytest + +from vllm.env_override import _apply_fxgraphcache_pickle_patch + + +class _BypassStub(Exception): + """Stand-in for BypassFxGraphCache in unit tests.""" + + +class TestApplyFxgraphcachePicklePatch: + def test_valueerror_converted_to_bypass(self): + class Pickler: + def dumps(self, obj): + raise ValueError("can't serialize blocked layout") + + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + + with pytest.raises(_BypassStub, match="Failed to pickle cache key"): + Pickler().dumps(object()) + + def test_original_valueerror_chained(self): + class Pickler: + def dumps(self, obj): + raise ValueError("bad tensor layout") + + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + + with pytest.raises(_BypassStub) as exc_info: + Pickler().dumps(object()) + + cause = exc_info.value.__cause__ + assert isinstance(cause, ValueError) + assert str(cause) == "bad tensor layout" + + def test_non_valueerror_propagates(self): + class Pickler: + def dumps(self, obj): + raise TypeError("unexpected type") + + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + + with pytest.raises(TypeError, match="unexpected type"): + Pickler().dumps(object()) + + def test_normal_return_preserved(self): + sentinel = b"serialized-graph-key" + + class Pickler: + def dumps(self, obj): + return sentinel + + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + + assert Pickler().dumps(object()) is sentinel + + def test_idempotent(self): + class Pickler: + def dumps(self, obj): + return b"ok" + + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + first_dumps = Pickler.dumps + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + + assert Pickler.dumps is first_dumps + + def test_sentinel_attribute_set(self): + class Pickler: + def dumps(self, obj): + return b"ok" + + assert not hasattr(Pickler.dumps, "_vllm_patched") + assert not getattr(Pickler, "_vllm_fxgraph_dumps_patched", False) + + _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) + + assert Pickler.dumps._vllm_patched is True # type: ignore[attr-defined] + assert Pickler._vllm_fxgraph_dumps_patched is True # type: ignore[attr-defined] + + +def test_patch_applied_in_current_environment(): + """Integration: verify patch state matches current torch version.""" + from torch._inductor.codecache import FxGraphCachePickler + + from vllm.utils.torch_utils import is_torch_equal_or_newer + + should_be_patched = is_torch_equal_or_newer( + "2.10.0" + ) and not is_torch_equal_or_newer("2.11.0") + + assert getattr(FxGraphCachePickler, "_vllm_fxgraph_dumps_patched", False) == ( + should_be_patched + ) + assert hasattr(FxGraphCachePickler.dumps, "_vllm_patched") == should_be_patched diff --git a/tests/tool_parsers/test_gemma4_tool_parser.py b/tests/tool_parsers/test_gemma4_tool_parser.py index 54f1e7eed8f..16450c6a8d0 100644 --- a/tests/tool_parsers/test_gemma4_tool_parser.py +++ b/tests/tool_parsers/test_gemma4_tool_parser.py @@ -85,6 +85,14 @@ class TestParseGemma4Args: result = _parse_gemma4_args("flag:false") assert result == {"flag": False} + def test_null_value(self): + # Bare `null` must parse as None (Python), not the string "null". + # Without this, tool_choice=auto would emit `{"param": "null"}` + # instead of `{"param": null}` for nullable tool parameters. + result = _parse_gemma4_args("param:null") + assert result == {"param": None} + assert json.dumps(result) == '{"param": null}' + def test_mixed_types(self): result = _parse_gemma4_args( 'name:<|"|>test<|"|>,count:42,active:true,score:3.14' diff --git a/tests/tool_parsers/test_glm47_moe_tool_parser.py b/tests/tool_parsers/test_glm47_moe_tool_parser.py index ebcd4e8d42e..5e5501e4abf 100644 --- a/tests/tool_parsers/test_glm47_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm47_moe_tool_parser.py @@ -117,28 +117,24 @@ class TestGlm47ExtractToolCalls: def _reset(parser): - parser._buffer = "" - parser._in_tool_call = False parser.current_tool_name_sent = False - parser._current_tool_name = None - parser._pending_key = None - parser._streaming_string_value = False parser.prev_tool_call_arr = [] parser.current_tool_id = -1 parser.streamed_args_for_tool = [] parser._tool_call_ids = [] - parser._args_started = [] - parser._args_closed = [] - parser._seen_keys = [] + parser._sent_content_idx = 0 class TestGlm47Streaming: def test_no_args(self, glm47_tool_parser, mock_request): _reset(glm47_tool_parser) - for chunk in ["", "get_current_date", ""]: + chunks = ["", "get_current_date", ""] + current_text = "" + for chunk in chunks: + current_text += chunk glm47_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], @@ -149,10 +145,7 @@ class TestGlm47Streaming: def test_with_args(self, glm47_tool_parser, mock_request): _reset(glm47_tool_parser) - # Split chunks so that the incremental string streaming path - # processes the value, its closing tag, and the tool-call closing - # tag in separate calls. - for chunk in [ + chunks = [ "", "get_weather\n", "city", @@ -160,14 +153,18 @@ class TestGlm47Streaming: "Beijing", "", "", - ]: + ] + current_text = "" + for chunk in chunks: + current_text += chunk glm47_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], request=mock_request, ) - assert glm47_tool_parser.prev_tool_call_arr[0]["arguments"]["city"] == "Beijing" + args = json.loads(glm47_tool_parser.prev_tool_call_arr[0]["arguments"]) + assert args["city"] == "Beijing" diff --git a/tests/tool_parsers/test_glm4_moe_tool_parser.py b/tests/tool_parsers/test_glm4_moe_tool_parser.py index dbfce204fde..9f430b7814f 100644 --- a/tests/tool_parsers/test_glm4_moe_tool_parser.py +++ b/tests/tool_parsers/test_glm4_moe_tool_parser.py @@ -357,81 +357,69 @@ meaningwhile, I will also check the weather in Shanghai. def test_streaming_basic_functionality(glm4_moe_tool_parser, mock_request): """Test basic streaming functionality.""" - # Reset streaming state - glm4_moe_tool_parser.current_tool_name_sent = False - glm4_moe_tool_parser.prev_tool_call_arr = [] - glm4_moe_tool_parser.current_tool_id = -1 - glm4_moe_tool_parser.streamed_args_for_tool = [] + _reset_streaming_state(glm4_moe_tool_parser) - # Test with a simple tool call current_text = """get_weather city Beijing """ - # Mock token IDs for testing - tool_call_start_id = glm4_moe_tool_parser.tool_call_start_token_id or 12345 - tool_call_end_id = glm4_moe_tool_parser.tool_call_end_token_id or 12346 - result = glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", current_text=current_text, - delta_text="", + delta_text=current_text, previous_token_ids=[], - current_token_ids=[tool_call_start_id, tool_call_end_id], - delta_token_ids=[tool_call_end_id], + current_token_ids=[], + delta_token_ids=[], request=mock_request, ) - # The result behavior depends on the streaming state - # This test mainly ensures no exceptions are thrown - assert result is None or hasattr(result, "tool_calls") or hasattr(result, "content") + # Should return tool call with name and arguments in one shot + assert result is not None + assert result.tool_calls is not None + assert len(result.tool_calls) >= 1 def test_streaming_no_tool_calls(glm4_moe_tool_parser, mock_request): """Test streaming when there are no tool calls.""" + _reset_streaming_state(glm4_moe_tool_parser) + current_text = "This is just regular text without any tool calls." result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="This is just regular text", + previous_text="", current_text=current_text, - delta_text=" without any tool calls.", + delta_text=current_text, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], request=mock_request, ) - # Should return the delta text as content + # Should return content assert result is not None - assert hasattr(result, "content") - assert result.content == " without any tool calls." + assert result.content == current_text def test_streaming_with_content_before_tool_calls(glm4_moe_tool_parser, mock_request): """Test streaming when there's content before tool calls.""" - # Reset streaming state - glm4_moe_tool_parser.current_tool_name_sent = False - glm4_moe_tool_parser.prev_tool_call_arr = [] - glm4_moe_tool_parser.current_tool_id = -1 - glm4_moe_tool_parser.streamed_args_for_tool = [] + _reset_streaming_state(glm4_moe_tool_parser) - current_text = "I will help you get the weather" + current_text = "I will help you get the weather." result = glm4_moe_tool_parser.extract_tool_calls_streaming( - previous_text="I will help you", + previous_text="", current_text=current_text, - delta_text="get the weather.", + delta_text=current_text, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], request=mock_request, ) - # Should return content when no tool call tokens are detected + # Should return content before the tag assert result is not None - assert hasattr(result, "content") - assert result.content == "get the weather." + assert result.content == "I will help you get the weather." def test_extract_tool_calls_special_characters(glm4_moe_tool_parser, mock_request): @@ -479,26 +467,19 @@ def test_extract_tool_calls_incomplete_tool_call(glm4_moe_tool_parser, mock_requ def _reset_streaming_state(parser): """Helper to reset parser streaming state.""" - parser._buffer = "" - parser._in_tool_call = False parser.current_tool_name_sent = False - parser._current_tool_name = None - parser._pending_key = None - parser._streaming_string_value = False parser.prev_tool_call_arr = [] parser.current_tool_id = -1 parser.streamed_args_for_tool = [] parser._tool_call_ids = [] - parser._args_started = [] - parser._args_closed = [] - parser._seen_keys = [] + parser._sent_content_idx = 0 def test_streaming_incremental_string_value(glm4_moe_tool_parser, mock_request): """Test incremental streaming of string argument values.""" _reset_streaming_state(glm4_moe_tool_parser) - # Simulate streaming a tool call character by character + # Simulate streaming a tool call chunk by chunk chunks = [ "", "get_weather\n", @@ -511,30 +492,31 @@ def test_streaming_incremental_string_value(glm4_moe_tool_parser, mock_request): ] collected_fragments = [] + current_text = "" for chunk in chunks: + current_text += chunk result = glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], request=mock_request, ) - if result is not None and hasattr(result, "tool_calls") and result.tool_calls: + if result is not None and result.tool_calls: for tc in result.tool_calls: - if hasattr(tc, "function") and tc.function: - func = tc.function - if isinstance(func, dict): - if func.get("arguments"): - collected_fragments.append(func["arguments"]) - if func.get("name"): - collected_fragments.append(f"name:{func['name']}") - else: - if func.arguments: - collected_fragments.append(func.arguments) - if func.name: - collected_fragments.append(f"name:{func.name}") + func = tc.function + if isinstance(func, dict): + if func.get("arguments"): + collected_fragments.append(func["arguments"]) + if func.get("name"): + collected_fragments.append(f"name:{func['name']}") + else: + if func.arguments: + collected_fragments.append(func.arguments) + if func.name: + collected_fragments.append(f"name:{func.name}") # Verify we got incremental streaming of the argument value assert len(collected_fragments) > 0 @@ -547,11 +529,11 @@ def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request): """Test that empty tool calls don't cause infinite loops.""" _reset_streaming_state(glm4_moe_tool_parser) - # Empty tool call should be handled gracefully + current_text = "" result = glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", - delta_text="", + current_text=current_text, + delta_text=current_text, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], @@ -561,60 +543,52 @@ def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request): # Should not hang and should return something (None or content) # The key is that this completes without hanging assert result is None or hasattr(result, "content") or hasattr(result, "tool_calls") - # State should be properly reset - assert glm4_moe_tool_parser.current_tool_id == -1 def test_streaming_prev_tool_call_arr_updates(glm4_moe_tool_parser, mock_request): - """Test that prev_tool_call_arr contains parsed dict after tool call.""" + """Test that prev_tool_call_arr is populated incrementally.""" _reset_streaming_state(glm4_moe_tool_parser) - # Stream a complete tool call - name_only = {"name": "get_weather", "arguments": {}} - name_and_args = {"name": "get_weather", "arguments": {"city": "Beijing"}} chunks = [ - # Delta, expected streamed_args_for_tool, expected prev_tool_call_arr - ("get_weather\n", "", name_only), - ("city", "", name_only), - ("Beijing", '{"city": "Beijing"', name_only), - # Note: arguments are only updated when the tool call is complete. - ("", '{"city": "Beijing"}', name_and_args), + "get_weather\n", + "city", + "Beijing", + "", ] - for chunk, exp_streamed, exp_prev_tc in chunks: + current_text = "" + for chunk in chunks: + current_text += chunk glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], request=mock_request, ) - assert glm4_moe_tool_parser.streamed_args_for_tool[0] == exp_streamed - assert glm4_moe_tool_parser.prev_tool_call_arr[0] == exp_prev_tc - # After the tool call completes, prev_tool_call_arr should have parsed dict + # After the tool call completes, prev_tool_call_arr should be populated assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1 tool_entry = glm4_moe_tool_parser.prev_tool_call_arr[0] assert tool_entry.get("name") == "get_weather" - # arguments should be a dict, not a string - args = tool_entry.get("arguments") - assert isinstance(args, dict), f"Expected dict, got {type(args)}" - assert args.get("city") == "Beijing" - # Test equivalence of prev_tool_call_arr and streamed_args_for_tool - # Simulates logic in chat_completion/serving.py:chat_completion_stream_generator - tool_call_json = json.dumps(tool_entry.get("arguments", {})) - streamed_content = glm4_moe_tool_parser.streamed_args_for_tool[0] - assert tool_call_json.startswith(streamed_content) + # arguments is a JSON string in the re-parse approach + args_str = tool_entry.get("arguments") + assert isinstance(args_str, str), f"Expected str, got {type(args_str)}" + parsed = json.loads(args_str) + assert parsed["city"] == "Beijing" + + # streamed_args_for_tool should match prev_tool_call_arr arguments + streamed = glm4_moe_tool_parser.streamed_args_for_tool[0] + assert streamed == args_str def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_request): """Test streaming multiple sequential tool calls.""" _reset_streaming_state(glm4_moe_tool_parser) - # Stream two tool calls chunks = [ "get_weather\n", "city", @@ -626,10 +600,12 @@ def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_req "", ] + current_text = "" for chunk in chunks: + current_text += chunk glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], @@ -639,15 +615,16 @@ def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_req # Should have two tool calls in prev_tool_call_arr assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 2 - assert glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"]["city"] == "Beijing" - assert glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"]["city"] == "Shanghai" + args0 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"]) + args1 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"]) + assert args0["city"] == "Beijing" + assert args1["city"] == "Shanghai" def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request): """Test that special characters in string values are properly escaped.""" _reset_streaming_state(glm4_moe_tool_parser) - # String with characters that need JSON escaping chunks = [ "send_message\n", "message", @@ -655,10 +632,12 @@ def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request): "", ] + current_text = "" for chunk in chunks: + current_text += chunk glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], @@ -669,10 +648,8 @@ def test_streaming_json_escape_in_string(glm4_moe_tool_parser, mock_request): # The streamed_args_for_tool should contain valid JSON assert len(glm4_moe_tool_parser.streamed_args_for_tool) == 1 args_json = glm4_moe_tool_parser.streamed_args_for_tool[0] - # Should be parseable as JSON parsed = json.loads(args_json) assert "message" in parsed - # The value should preserve the special characters assert '"' in parsed["message"] or "world" in parsed["message"] @@ -749,27 +726,27 @@ if __name__ == "__main__": # Count argument fragments fragment_count = 0 + current_text = "" for chunk in chunks: + current_text += chunk result = glm4_moe_tool_parser.extract_tool_calls_streaming( previous_text="", - current_text="", + current_text=current_text, delta_text=chunk, previous_token_ids=[], current_token_ids=[], delta_token_ids=[], request=request, ) - if result is not None and hasattr(result, "tool_calls") and result.tool_calls: + if result is not None and result.tool_calls: for tc in result.tool_calls: - if hasattr(tc, "function") and tc.function: - func = tc.function - args = ( - func.get("arguments") - if isinstance(func, dict) - else getattr(func, "arguments", None) - ) - if args: - fragment_count += 1 + func = tc.function + if isinstance(func, dict): + args = func.get("arguments") + else: + args = getattr(func, "arguments", None) + if args: + fragment_count += 1 # For true incremental streaming, we expect many fragments (10+) # Old buffered implementation would give only 1-3 fragments @@ -927,3 +904,432 @@ def test_unicode_characters_preserved(glm4_moe_tool_parser, mock_request): parsed_args = json.loads(raw_args) assert parsed_args["greeting"] == "你好世界" assert parsed_args["emoji"] == "🎉" + + +def test_streaming_multi_token_chunks(glm4_moe_tool_parser, mock_request): + """Test that multi-token chunks (stream_interval > 1) are handled correctly. + + With stream_interval > 1 or MTP, multiple XML tags arrive in one delta. + The old buffer-based parser could only return one delta per call, losing + data on the final output. The re-parse approach handles this correctly. + """ + _reset_streaming_state(glm4_moe_tool_parser) + + # Simulate stream_interval=3: chunks contain multiple XML tags + chunks = [ + "get_weather\ncityBei", + "jing", + "", + ] + + current_text = "" + for chunk in chunks: + current_text += chunk + glm4_moe_tool_parser.extract_tool_calls_streaming( + previous_text="", + current_text=current_text, + delta_text=chunk, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=mock_request, + ) + + # All data should be captured despite multi-token chunks + assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1 + args = json.loads(glm4_moe_tool_parser.streamed_args_for_tool[0]) + assert args["city"] == "Beijing" + + +def test_streaming_entire_tool_call_at_once(glm4_moe_tool_parser, mock_request): + """Test that a complete tool call arriving in one delta works. + + This simulates the extreme MTP case where all tokens arrive at once. + """ + _reset_streaming_state(glm4_moe_tool_parser) + + full_text = ( + "get_weather\n" + "city" + "Beijing" + "" + ) + + result = glm4_moe_tool_parser.extract_tool_calls_streaming( + previous_text="", + current_text=full_text, + delta_text=full_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=mock_request, + ) + + # Should emit tool call with complete arguments in one shot + assert result is not None + assert result.tool_calls is not None + + # Verify final state + assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1 + args = json.loads(glm4_moe_tool_parser.streamed_args_for_tool[0]) + assert args["city"] == "Beijing" + + +def test_streaming_content_between_tool_calls_multi_token( + glm4_moe_tool_parser, mock_request +): + """Test content between tool calls with multi-token chunks.""" + _reset_streaming_state(glm4_moe_tool_parser) + + # Deliver everything at once — worst case for the old buffer parser + full_text = ( + "I will check.\n" + "get_weather\n" + "city" + "Beijing" + "" + "\nAlso Shanghai.\n" + "get_weather\n" + "city" + "Shanghai" + "" + ) + + # First call with partial text (content only) + partial = "I will check.\n" + result1 = glm4_moe_tool_parser.extract_tool_calls_streaming( + previous_text="", + current_text=partial, + delta_text=partial, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=mock_request, + ) + assert result1 is not None + assert result1.content == "I will check.\n" + + # Second call with everything + glm4_moe_tool_parser.extract_tool_calls_streaming( + previous_text="", + current_text=full_text, + delta_text=full_text[len(partial) :], + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=mock_request, + ) + + # Should have both tool calls + assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 2 + args0 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[0]["arguments"]) + args1 = json.loads(glm4_moe_tool_parser.prev_tool_call_arr[1]["arguments"]) + assert args0["city"] == "Beijing" + assert args1["city"] == "Shanghai" + + +def test_streaming_multi_token_with_multiple_args(glm4_moe_tokenizer): + """Test multi-token streaming with multiple arguments of mixed types.""" + tools = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="calculate", + parameters={ + "type": "object", + "properties": { + "operation": {"type": "string"}, + "a": {"type": "number"}, + "b": {"type": "number"}, + }, + }, + ), + ), + ] + parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) + request = ChatCompletionRequest( + model=MODEL, + messages=[], + tools=tools, + ) + + # All arguments arrive in two big chunks (simulates stream_interval=5) + chunks = [ + "calculate\noperationadda", + "42b3.14", + ] + + current_text = "" + for chunk in chunks: + current_text += chunk + parser.extract_tool_calls_streaming( + previous_text="", + current_text=current_text, + delta_text=chunk, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=request, + ) + + args = json.loads(parser.streamed_args_for_tool[0]) + assert args["operation"] == "add" + assert args["a"] == 42 + assert args["b"] == 3.14 + + +def _simulate_streaming(tokenizer, parser, request, text, stream_interval=1): + """Simulate streaming with a given stream_interval. + + Tokens are batched into chunks of ``stream_interval`` tokens, + mimicking how the output processor delivers them. + Returns a list of non-None DeltaMessages. + """ + tokens = tokenizer.encode(text) + previous_text = "" + deltas = [] + for i in range(0, len(tokens), stream_interval): + chunk_ids = tokens[i : i + stream_interval] + delta_text = tokenizer.decode(chunk_ids) + current_text = previous_text + delta_text + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=chunk_ids, + request=request, + ) + previous_text = current_text + if delta is not None: + deltas.append(delta) + return deltas + + +def _collect_from_deltas(deltas): + """Reconstruct tool call names/args and content from a delta stream.""" + tools: dict[int, dict] = {} + content_parts: list[str] = [] + for d in deltas: + if d.content: + content_parts.append(d.content) + if d.tool_calls: + for tc in d.tool_calls: + func = tc.function + if isinstance(func, dict): + name = func.get("name") + args = func.get("arguments") + else: + name = getattr(func, "name", None) + args = getattr(func, "arguments", None) + idx = tc.index + if idx not in tools: + tools[idx] = {"name": None, "args_fragments": []} + if name: + tools[idx]["name"] = name + if args: + tools[idx]["args_fragments"].append(args) + return content_parts, tools + + +@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) +def test_stream_interval_single_tool_call(glm4_moe_tokenizer, stream_interval): + """Tool call streaming produces correct name + args at any interval.""" + tools = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ), + ), + ] + parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + text = ( + "get_weather\n" + "city" + "Beijing" + "" + ) + + deltas = _simulate_streaming( + glm4_moe_tokenizer, parser, request, text, stream_interval + ) + _, tools_found = _collect_from_deltas(deltas) + + assert 0 in tools_found + assert tools_found[0]["name"] == "get_weather" + args_json = "".join(tools_found[0]["args_fragments"]) + parsed = json.loads(args_json) + assert parsed == {"city": "Beijing"} + + +@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) +def test_stream_interval_multiple_tool_calls(glm4_moe_tokenizer, stream_interval): + """Multiple sequential tool calls with correct indices at any interval.""" + tools = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ), + ), + ] + parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + text = ( + "get_weather\n" + "city" + "Beijing" + "" + "get_weather\n" + "city" + "Shanghai" + "" + ) + + deltas = _simulate_streaming( + glm4_moe_tokenizer, parser, request, text, stream_interval + ) + _, tools_found = _collect_from_deltas(deltas) + + assert 0 in tools_found and 1 in tools_found + args0 = json.loads("".join(tools_found[0]["args_fragments"])) + args1 = json.loads("".join(tools_found[1]["args_fragments"])) + assert args0 == {"city": "Beijing"} + assert args1 == {"city": "Shanghai"} + + +@pytest.mark.parametrize("stream_interval", [1, 2, 3, 5, 8]) +def test_stream_interval_content_then_tool_call(glm4_moe_tokenizer, stream_interval): + """Content before a tool call is fully emitted before tool deltas.""" + tools = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ), + ), + ] + parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + text = ( + "I will check the weather for you.\n" + "get_weather\n" + "city" + "Beijing" + "" + ) + + deltas = _simulate_streaming( + glm4_moe_tokenizer, parser, request, text, stream_interval + ) + content_parts, tools_found = _collect_from_deltas(deltas) + + # Content must be present and precede tool calls + full_content = "".join(content_parts) + assert "I will check the weather" in full_content + + # Tool call must be correct + assert 0 in tools_found + assert tools_found[0]["name"] == "get_weather" + args = json.loads("".join(tools_found[0]["args_fragments"])) + assert args == {"city": "Beijing"} + + +def test_stream_interval_extreme_single_chunk(glm4_moe_tokenizer): + """Extreme MTP: entire output arrives in one chunk (interval=9999).""" + tools = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ), + ), + ] + parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + text = ( + "Here is the weather.\n" + "get_weather\n" + "city" + "Beijing" + "" + ) + + deltas = _simulate_streaming( + glm4_moe_tokenizer, parser, request, text, stream_interval=9999 + ) + content_parts, tools_found = _collect_from_deltas(deltas) + + assert "Here is the weather" in "".join(content_parts) + assert 0 in tools_found + assert tools_found[0]["name"] == "get_weather" + args = json.loads("".join(tools_found[0]["args_fragments"])) + assert args == {"city": "Beijing"} + + +@pytest.mark.parametrize("stream_interval", [1, 2, 5]) +def test_stream_interval_content_between_tool_calls( + glm4_moe_tokenizer, stream_interval +): + """Content between tool calls must be emitted, not silently dropped.""" + tools = [ + ChatCompletionToolsParam( + function=FunctionDefinition( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ), + ), + ] + parser = Glm4MoeModelToolParser(glm4_moe_tokenizer, tools=tools) + request = ChatCompletionRequest(model=MODEL, messages=[], tools=tools) + + text = ( + "Checking Beijing.\n" + "get_weather\n" + "city" + "Beijing" + "" + "\nAlso Shanghai.\n" + "get_weather\n" + "city" + "Shanghai" + "" + ) + + deltas = _simulate_streaming( + glm4_moe_tokenizer, parser, request, text, stream_interval + ) + content_parts, tools_found = _collect_from_deltas(deltas) + + full_content = "".join(content_parts) + # Both prefix and inter-tool-call content must appear + assert "Checking Beijing" in full_content + assert "Also Shanghai" in full_content + + # Both tool calls must be correct + assert 0 in tools_found and 1 in tools_found + args0 = json.loads("".join(tools_found[0]["args_fragments"])) + args1 = json.loads("".join(tools_found[1]["args_fragments"])) + assert args0 == {"city": "Beijing"} + assert args1 == {"city": "Shanghai"} diff --git a/tests/transformers_utils/test_utils.py b/tests/transformers_utils/test_utils.py index 485c2efff77..94dd014c929 100644 --- a/tests/transformers_utils/test_utils.py +++ b/tests/transformers_utils/test_utils.py @@ -81,6 +81,25 @@ class TestIsRemoteGGUF: assert not is_remote_gguf("repo/model:INVALID_M") assert not is_remote_gguf("repo/model:Q9_K_M") + def test_is_remote_gguf_nonstandard_quant_type(self): + """Test is_remote_gguf with non-standard quant types containing + a known GGML type.""" + # Non-standard quant types with known GGML type after prefix + assert is_remote_gguf("unsloth/Qwen3.5-35B-A3B-GGUF:UD-Q4_K_XL") + assert is_remote_gguf("user/Model:UD-Q4_K_M") + assert is_remote_gguf("user/SomeModel:Custom-Q8_0") + + # Exact GGML type after prefix (no suffix stripping needed) + assert is_remote_gguf("user/Model-GGUF:UD-IQ4_NL") + assert is_remote_gguf("user/Model-GGUF:UD-Q8_0") + + # Completely unknown quant types should still fail + assert not is_remote_gguf("repo/model:TOTALLY-RANDOM") + assert not is_remote_gguf("user/Model:UD-INVALID") + + # No dash separator → not recognized as prefixed + assert not is_remote_gguf("repo/model:UDIQ4NL") + def test_is_remote_gguf_without_colon(self): """Test is_remote_gguf without colon.""" assert not is_remote_gguf("repo/model") @@ -143,6 +162,14 @@ class TestSplitRemoteGGUF: assert repo_id == "repo/model" assert quant_type == "Q3_K_S" + def test_split_remote_gguf_nonstandard_quant_type(self): + """Test split_remote_gguf with non-standard quant types in GGUF repos.""" + repo_id, quant_type = split_remote_gguf( + "unsloth/Qwen3.5-35B-A3B-GGUF:UD-Q4_K_XL" + ) + assert repo_id == "unsloth/Qwen3.5-35B-A3B-GGUF" + assert quant_type == "UD-Q4_K_XL" + def test_split_remote_gguf_with_path_object(self): """Test split_remote_gguf with Path object.""" repo_id, quant_type = split_remote_gguf(Path("unsloth/Qwen3-0.6B-GGUF:IQ1_S")) diff --git a/tests/utils.py b/tests/utils.py index 7af72cb730b..cfa5f525f5d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1261,9 +1261,6 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non @functools.wraps(func) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None: - # Make the process the leader of its own process group - # to avoid sending SIGTERM to the parent process - os.setpgrp() from _pytest.outcomes import Skipped # Create a unique temporary file to store exception info from child @@ -1283,6 +1280,9 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non pid = os.fork() print(f"Fork a new process to run a test {pid}") if pid == 0: + # Make the child process the leader of its own process group + # to avoid sending SIGTERM to the parent process + os.setpgrp() # Parent process responsible for deleting, don't delete # in child. delete_after.pop_all() @@ -1322,14 +1322,12 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non else: os._exit(0) else: - pgid = os.getpgid(pid) + # After setpgrp(), the child's pgid equals its pid + pgid = pid _pid, _exitcode = os.waitpid(pid, 0) - # ignore SIGTERM signal itself - old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN) - # kill all child processes - os.killpg(pgid, signal.SIGTERM) - # restore the signal handler - signal.signal(signal.SIGTERM, old_signal_handler) + # kill all child processes - but they may already have exited cleanly + with contextlib.suppress(ProcessLookupError): + os.killpg(pgid, signal.SIGTERM) if _exitcode != 0: # Try to read the exception from the child process exc_info = {} diff --git a/tests/utils_/test_mem_utils.py b/tests/utils_/test_mem_utils.py index 4067b025781..c56d6a44e7d 100644 --- a/tests/utils_/test_mem_utils.py +++ b/tests/utils_/test_mem_utils.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock, patch + import torch from vllm_test_utils.monitor import monitor @@ -61,3 +63,62 @@ def test_memory_profiling(): del weights lib.cudaFree(handle1) lib.cudaFree(handle2) + + +def test_memory_snapshot_uses_psutil_on_integrated_gpu(): + """On integrated (UMA) GPUs, free_memory should come from psutil.""" + mock_cuda_free = 40 * 1024**3 + mock_cuda_total = 120 * 1024**3 + mock_psutil_available = 100 * 1024**3 + + with ( + patch("vllm.utils.mem_utils.current_platform") as mock_platform, + patch("vllm.utils.mem_utils.psutil") as mock_psutil, + ): + mock_platform.mem_get_info.return_value = ( + mock_cuda_free, + mock_cuda_total, + ) + mock_platform.is_integrated_gpu.return_value = True + mock_platform.memory_stats.return_value = { + "allocated_bytes.all.peak": 0, + } + mock_platform.memory_reserved.return_value = 0 + mock_platform.current_device = lambda: "cuda:0" + + mock_vmem = MagicMock() + mock_vmem.available = mock_psutil_available + mock_psutil.virtual_memory.return_value = mock_vmem + + snapshot = MemorySnapshot(device="cuda:0") + + assert snapshot.free_memory == mock_psutil_available + assert snapshot.total_memory == mock_cuda_total + mock_psutil.virtual_memory.assert_called_once() + + +def test_memory_snapshot_uses_cuda_on_discrete_gpu(): + """On discrete GPUs, free_memory should come from CUDA mem_get_info.""" + mock_cuda_free = 70 * 1024**3 + mock_cuda_total = 80 * 1024**3 + + with ( + patch("vllm.utils.mem_utils.current_platform") as mock_platform, + patch("vllm.utils.mem_utils.psutil") as mock_psutil, + ): + mock_platform.mem_get_info.return_value = ( + mock_cuda_free, + mock_cuda_total, + ) + mock_platform.is_integrated_gpu.return_value = False + mock_platform.memory_stats.return_value = { + "allocated_bytes.all.peak": 0, + } + mock_platform.memory_reserved.return_value = 0 + mock_platform.current_device = lambda: "cuda:0" + + snapshot = MemorySnapshot(device="cuda:0") + + assert snapshot.free_memory == mock_cuda_free + assert snapshot.total_memory == mock_cuda_total + mock_psutil.virtual_memory.assert_not_called() diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index e821e47172c..c05e979a9ff 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -153,7 +153,6 @@ def test_prefix_caching_for_prefill_dedup(): same_prompt=True, block_size=BLOCK_SIZE, ) - requests_copy = requests.copy() # Two requests with the same prompt. req0 = requests.pop(0) @@ -167,26 +166,31 @@ def test_prefix_caching_for_prefill_dedup(): # Make sure prefix caching de-duplicates the prompts in the same step, # so all the blocks except the last are shared between the two requests. assert len(sched_output.num_scheduled_tokens) == 2 - num_blocks = num_prompt_tokens // BLOCK_SIZE - assert req0.num_cached_tokens == 0 - assert req1.num_cached_tokens >= num_blocks * BLOCK_SIZE + assert sched_output.num_scheduled_tokens[req0.request_id] == num_prompt_tokens + assert ( + sched_output.num_scheduled_tokens[req1.request_id] + == num_prompt_tokens % BLOCK_SIZE + ) sched_outputs.append(scheduler.schedule()) while sched_outputs: + added_req = None if requests: - scheduler.add_request(requests.pop(0)) + added_req = requests.pop(0) + scheduler.add_request(added_req) sched_output = sched_outputs.popleft() model_runner_output = _make_model_runner_output(sched_output) scheduler.update_from_output(sched_output, model_runner_output) sched_output = scheduler.schedule() if sched_output.num_scheduled_tokens: sched_outputs.append(sched_output) + if added_req: + assert ( + sched_output.num_scheduled_tokens[added_req.request_id] + == num_prompt_tokens % BLOCK_SIZE + ) - # Other requests scheduled after the two requests should also get - # prefix cache hit. assert scheduler.get_num_unfinished_requests() == 0 - for req in requests_copy[1:]: - assert req.num_cached_tokens >= num_blocks * BLOCK_SIZE def test_prefix_caching_for_multi_turn(): @@ -243,12 +247,15 @@ def test_prefix_caching_for_multi_turn(): # Schedule the next-turn requests. for req in next_turn_requests: scheduler.add_request(req) - sched_outputs.append(scheduler.schedule()) + sched_output = scheduler.schedule() + sched_outputs.append(sched_output) # Make sure the next-turn requests get prefix cache hit by the previous # requests. for req in next_turn_requests: - assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE + assert sched_output.num_scheduled_tokens[req.request_id] == ( + req.num_prompt_tokens % BLOCK_SIZE + ) def test_abort_request_when_structured_output_fsm_cannot_advance(): diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index d8ecf28cbed..046f04e0c79 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -10,6 +10,7 @@ import torch import vllm.v1.core.kv_cache_utils as kv_cache_utils from vllm.config import ModelConfig, SchedulerConfig, VllmConfig +from vllm.config.kv_events import KVEventsConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -2137,3 +2138,30 @@ def test_unify_hybrid_kv_cache_specs(): with pytest.raises(ValueError): kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec) + + +def test_hma_not_disabled_when_kv_events_enabled(): + """ + Test enabling KV events must not force disable_hybrid_kv_cache_manager to True. + + This test guards against that regression by verifying that a VllmConfig + with kv_events_config set still resolves disable_hybrid_kv_cache_manager + to False (i.e. HMA remains enabled) when no other condition requires it + to be disabled. + """ + model_config = ModelConfig(max_model_len=16) + kv_events_config = KVEventsConfig( + enable_kv_cache_events=True, + publisher="null", + ) + + # Leave disable_hybrid_kv_cache_manager as None (the default) so that + # VllmConfig.__post_init__ resolves it automatically. + vllm_config = VllmConfig( + model_config=model_config, + kv_events_config=kv_events_config, + ) + + assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, ( + "kv_events_config must not force-disable the hybrid KV cache manager." + ) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index b8b387fffd9..22220599f15 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1970,6 +1970,7 @@ def test_null_parent_block_hash(): block_size = 1 num_cached_blocks = 2 num_full_blocks = 4 + kv_cache_group_id = 0 pool = BlockPool( num_gpu_blocks=8, @@ -2002,7 +2003,7 @@ def test_null_parent_block_hash(): num_cached_blocks=num_cached_blocks, num_full_blocks=num_full_blocks, block_size=block_size, - kv_cache_group_id=0, + kv_cache_group_id=kv_cache_group_id, ) events = pool.take_events() @@ -2021,6 +2022,7 @@ def test_null_parent_block_hash(): for h in req.block_hashes[num_cached_blocks:num_full_blocks] ] assert event.block_hashes == expected_new_hashes + assert event.group_idx == kv_cache_group_id # Ensure we didn't accidentally assign a hash to the null block. assert pool.null_block.block_hash is None @@ -2087,6 +2089,153 @@ def test_kv_cache_events_with_lora(blocks_to_cache: int): assert block_stored_event.block_size == block_size +@pytest.mark.parametrize("group_id", [0, 1, 2]) +def test_block_stored_event_group_idx(group_id: int): + """Test BlockStored events emitted by cache_full_blocks carry the correct + group_idx.""" + block_size = 4 + num_tokens = block_size * 2 + + pool = BlockPool( + num_gpu_blocks=5, + enable_caching=True, + hash_block_size=block_size, + enable_kv_cache_events=True, + ) + + req = make_request( + "req_grp_idx", + prompt_token_ids=list(range(num_tokens)), + block_size=block_size, + hash_fn=sha256, + ) + + blocks = pool.get_new_blocks(2) + pool.cache_full_blocks( + request=req, + blocks=blocks, + num_cached_blocks=0, + num_full_blocks=2, + block_size=block_size, + kv_cache_group_id=group_id, + ) + + events = pool.take_events() + assert len(events) == 1 + assert isinstance(events[0], BlockStored) + assert events[0].group_idx == group_id + + +def test_block_stored_event_group_idx_multiple_groups(): + """ + Test BlockStored events for separate HMA groups that each carry the + correct group_idx. + + Simulates the HMA scenario where full-attention blocks (group 0) and + sliding-window blocks (group 1) are cached independently and must be + distinguishable by consumers doing HMA-aware prefix-cache routing. + """ + block_size = 4 + num_tokens = block_size * 2 + + # null block + 4 usable (2 per group) + pool = BlockPool( + num_gpu_blocks=5, + enable_caching=True, + hash_block_size=block_size, + enable_kv_cache_events=True, + ) + + req = make_request( + "req_multi_grp", + prompt_token_ids=list(range(num_tokens)), + block_size=block_size, + hash_fn=sha256, + ) + + # Cache blocks for group 0 (full-attention) + blocks_grp0 = pool.get_new_blocks(2) + pool.cache_full_blocks( + request=req, + blocks=blocks_grp0, + num_cached_blocks=0, + num_full_blocks=2, + block_size=block_size, + kv_cache_group_id=0, + ) + + # Cache blocks for group 1 (sliding-window) + blocks_grp1 = pool.get_new_blocks(2) + pool.cache_full_blocks( + request=req, + blocks=blocks_grp1, + num_cached_blocks=0, + num_full_blocks=2, + block_size=block_size, + kv_cache_group_id=1, + ) + + events = pool.take_events() + assert len(events) == 2 + assert isinstance(events[0], BlockStored) + assert events[0].group_idx == 0 + assert isinstance(events[1], BlockStored) + assert events[1].group_idx == 1 + + +@pytest.mark.parametrize("group_id", [0, 1, 2]) +def test_block_removed_event_group_idx(group_id: int): + """ + Test BlockRemoved events emitted on eviction carry the group_idx extracted + from the evicted block's BlockHashWithGroupId via get_group_id(). + """ + block_size = 4 + num_tokens = block_size * 2 + + # null block + 4 usable; allocate all 4, cache 2, free all, re-allocate + # all 4 so the 2 cached blocks are forced through _maybe_evict_cached_block. + pool = BlockPool( + num_gpu_blocks=5, + enable_caching=True, + hash_block_size=block_size, + enable_kv_cache_events=True, + ) + + req = make_request( + "req_evict_grp", + prompt_token_ids=list(range(num_tokens)), + block_size=block_size, + hash_fn=sha256, + ) + + # Allocate all usable blocks and cache the first two for the target group. + all_blocks = pool.get_new_blocks(4) + pool.cache_full_blocks( + request=req, + blocks=all_blocks, + num_cached_blocks=0, + num_full_blocks=2, + block_size=block_size, + kv_cache_group_id=group_id, + ) + + # Drain the BlockStored events so only eviction events remain later. + pool.take_events() + + # Return all blocks to the free queue so they become eviction candidates. + pool.free_blocks(all_blocks) + + # Re-allocate all blocks; the two with hashes trigger BlockRemoved events. + pool.get_new_blocks(4) + + events = pool.take_events() + removed_events = [e for e in events if isinstance(e, BlockRemoved)] + + assert len(removed_events) == 2 + for event in removed_events: + assert event.group_idx == group_id + + def test_eagle_enabled_removes_last_block(): """Verify Eagle does NOT remove blocks when request length is divisible by block size.""" diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 8fd2309f7ec..f825220800f 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1039,6 +1039,54 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks(): assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens +def test_scheduler_stats_waiting_queues(): + """Test that scheduler stats correctly report waiting and skipped_waiting queues.""" + # Create scheduler with limited capacity so we can have waiting requests + scheduler = create_scheduler(max_num_batched_tokens=100) + + # Create requests: some will be scheduled, some will wait on capacity, + # and some will be blocked by constraints + all_requests = create_requests(num_requests=5, num_tokens=50) + + # Add 3 requests - only 2 can be scheduled (2 * 50 = 100 tokens) + # The 3rd will remain in waiting queue (capacity constraint) + for request in all_requests[:3]: + scheduler.add_request(request) + + # Manually add 2 more to skipped_waiting to simulate constraint-blocked + for request in all_requests[3:]: + request.status = RequestStatus.WAITING_FOR_REMOTE_KVS + scheduler.skipped_waiting.add_request(request) + + # Schedule - this will schedule 2 requests, leaving 1 in waiting + output = scheduler.schedule() + + # Verify: 2 scheduled, 1 still waiting on capacity, 2 blocked by constraints + assert len(output.scheduled_new_reqs) == 2 + assert len(scheduler.waiting) == 1 + assert len(scheduler.skipped_waiting) == 2 + + # Call update_from_output() to get frontend-facing stat + scheduled_req_ids = list(output.num_scheduled_tokens.keys()) + model_runner_output = ModelRunnerOutput( + req_ids=scheduled_req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(scheduled_req_ids)}, + sampled_token_ids=[[1]] * len(scheduled_req_ids), + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + engine_core_outputs = scheduler.update_from_output(output, model_runner_output) + assert engine_core_outputs and len(engine_core_outputs) > 0 + stats = engine_core_outputs[0].scheduler_stats + assert stats is not None + + # Verify stats match queue lengths after scheduling + assert stats.num_running_reqs == 2 # 2 were scheduled + assert stats.num_waiting_reqs == 1 # 1 waiting on capacity + assert stats.num_skipped_waiting_reqs == 2 # 2 blocked by constraints + + def _assert_right_scheduler_output( output: SchedulerOutput, num_requests: int, diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 543dfc8bb31..3ad18cf9a78 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -6,8 +6,10 @@ Test organization: No GPU required: - TestFindBudgetGraph — greedy budget selection logic - TestGetCumulativeStats — hit/miss rate statistics + - TestGetInputModality — modality routing from mm_kwargs keys GPU required: - TestEncoderCudaGraphCaptureReplay — capture, replay, fallback, counters, chunking + - TestEncoderCudaGraphVideoReplay — video modality capture, replay """ from typing import Any @@ -205,11 +207,19 @@ class SimpleMockViTModel(torch.nn.Module): def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: return EncoderCudaGraphConfig( modalities=["image"], - input_key="pixel_values", + input_key_by_modality={ + "image": "pixel_values", + }, buffer_keys=["dummy_buf"], out_hidden_size=_HIDDEN, ) + def get_input_modality( + self, + mm_kwargs: dict[str, Any], + ) -> str: + return "image" + def get_encoder_cudagraph_budget_range( self, vllm_config, @@ -268,6 +278,7 @@ class SimpleMockViTModel(torch.nn.Module): self, token_budget: int, max_batch_size: int, + max_frames_per_batch: int, device: torch.device, dtype: torch.dtype, ) -> EncoderCudaGraphCaptureInputs: @@ -294,6 +305,7 @@ class SimpleMockViTModel(torch.nn.Module): self, mm_kwargs: dict[str, Any], max_batch_size: int, + max_frames_per_batch: int, ) -> EncoderCudaGraphReplayBuffers: grid_thw = mm_kwargs["image_grid_thw"] n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE) @@ -327,11 +339,16 @@ def _make_manager_for_gpu( max_batch_size: int, device: torch.device, dtype: torch.dtype, + *, + max_frames_per_batch: int | None = None, ) -> EncoderCudaGraphManager: """Create EncoderCudaGraphManager bypassing VllmConfig for GPU tests.""" mgr = object.__new__(EncoderCudaGraphManager) mgr.token_budgets = sorted(token_budgets) mgr.max_batch_size = max_batch_size + mgr.max_frames_per_batch = ( + max_frames_per_batch if max_frames_per_batch is not None else max_batch_size * 2 + ) mgr.use_dp = False mgr.budget_graphs = {} mgr.graph_hits = 0 @@ -366,6 +383,18 @@ def _make_mm_kwargs( } +def _make_video_mm_kwargs( + grid_thw_list: list[list[int]], + device: torch.device, + dtype: torch.dtype, +) -> dict[str, Any]: + """Create video mm_kwargs (pixel_values_videos / video_grid_thw) for testing.""" + return { + "pixel_values_videos": _make_pixel_values(grid_thw_list, device, dtype), + "video_grid_thw": grid_thw_list, + } + + # --------------------------------------------------------------------------- # GPU tests — capture, replay, fallback, counters, chunking # --------------------------------------------------------------------------- @@ -449,3 +478,285 @@ class TestEncoderCudaGraphCaptureReplay: assert len(result) == n_images for out in result: assert out.shape == (4, _HIDDEN) + + +# --------------------------------------------------------------------------- +# SimpleMockViTVideoModel — extends SimpleMockViTModel with video support +# --------------------------------------------------------------------------- + + +class SimpleMockViTVideoModel(SimpleMockViTModel): + """ViT mock that supports both image and video modalities. + + Reuses SimpleMockViTModel's NN weights and _forward() logic. + Only the protocol methods that are key-dependent are overridden. + """ + + def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig: + return EncoderCudaGraphConfig( + modalities=["image", "video"], + input_key_by_modality={ + "image": "pixel_values", + "video": "pixel_values_videos", + }, + buffer_keys=["dummy_buf"], + out_hidden_size=_HIDDEN, + ) + + def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str: + return "video" if "video_grid_thw" in mm_kwargs else "image" + + # ------------------------------------------------------------------ + # Private helpers — route to the correct mm_kwargs keys + # ------------------------------------------------------------------ + + def _get_grid_thw(self, mm_kwargs: dict[str, Any]) -> list[list[int]]: + key = ( + "video_grid_thw" + if self.get_input_modality(mm_kwargs) == "video" + else "image_grid_thw" + ) + return mm_kwargs[key] + + def _get_pixel_values(self, mm_kwargs: dict[str, Any]) -> torch.Tensor: + key = ( + "pixel_values_videos" + if self.get_input_modality(mm_kwargs) == "video" + else "pixel_values" + ) + return mm_kwargs[key] + + # ------------------------------------------------------------------ + # Protocol overrides that depend on modality keys + # ------------------------------------------------------------------ + + def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int: + return len(self._get_grid_thw(mm_kwargs)) + + def get_encoder_cudagraph_per_item_output_tokens( + self, mm_kwargs: dict[str, Any] + ) -> list[int]: + m = _SPATIAL_MERGE + return [t * (h // m) * (w // m) for t, h, w in self._get_grid_thw(mm_kwargs)] + + def get_encoder_cudagraph_per_item_input_sizes( + self, mm_kwargs: dict[str, Any] + ) -> list[int]: + return [t * h * w for t, h, w in self._get_grid_thw(mm_kwargs)] + + def select_encoder_cudagraph_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + modality = self.get_input_modality(mm_kwargs) + pv_key = "pixel_values_videos" if modality == "video" else "pixel_values" + grid_key = "video_grid_thw" if modality == "video" else "image_grid_thw" + + grid_thw = self._get_grid_thw(mm_kwargs) + pixel_values = self._get_pixel_values(mm_kwargs) + + if len(indices) == 0: + return {pv_key: pixel_values[:0], grid_key: []} + + patches_per_item = [t * h * w for t, h, w in grid_thw] + cum_patches = [0] + for p in patches_per_item: + cum_patches.append(cum_patches[-1] + p) + + selected_pv = torch.cat( + [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices] + ) + return {pv_key: selected_pv, grid_key: [grid_thw[i] for i in indices]} + + def prepare_encoder_cudagraph_capture_inputs( + self, + token_budget: int, + max_batch_size: int, + max_frames_per_batch: int, + device: torch.device, + dtype: torch.dtype, + ) -> EncoderCudaGraphCaptureInputs: + per_item_output = token_budget // max_batch_size + frames_per_item = max_frames_per_batch // max_batch_size + if frames_per_item > 1: + # Video-format capture: size cu_seqlens for T frames per item. + tokens_per_frame = ( + per_item_output + frames_per_item - 1 + ) // frames_per_item + grid_config = [ + [frames_per_item, _SPATIAL_MERGE, tokens_per_frame * _SPATIAL_MERGE] + for _ in range(max_batch_size) + ] + else: + grid_config = [ + [1, _SPATIAL_MERGE, per_item_output * _SPATIAL_MERGE] + for _ in range(max_batch_size) + ] + total_patches = _count_input_patches(grid_config) + # Use pixel_values (image key) for capture — same patch shape as video. + dummy_pixel_values = torch.randn( + total_patches, _FLAT, device=device, dtype=dtype + ) + n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE) + dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype) + return EncoderCudaGraphCaptureInputs( + mm_kwargs={ + "pixel_values": dummy_pixel_values, + "image_grid_thw": grid_config, + }, + buffers={"dummy_buf": dummy_buf}, + ) + + def prepare_encoder_cudagraph_replay_buffers( + self, + mm_kwargs: dict[str, Any], + max_batch_size: int, + max_frames_per_batch: int, + ) -> EncoderCudaGraphReplayBuffers: + n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE) + p = next(self.parameters()) + dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype) + return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf}) + + def encoder_cudagraph_forward( + self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor] + ) -> torch.Tensor: + return self._forward(self._get_pixel_values(mm_kwargs)) + + def encoder_eager_forward(self, mm_kwargs: dict[str, Any]) -> torch.Tensor: + return self._forward(self._get_pixel_values(mm_kwargs)) + + +# --------------------------------------------------------------------------- +# No-GPU tests — get_input_modality routing +# --------------------------------------------------------------------------- + + +class TestGetInputModality: + """get_input_modality returns correct modality based on mm_kwargs keys.""" + + def test_image_only_model_always_returns_image(self): + model = SimpleMockViTModel() + mm_kwargs = { + "pixel_values": torch.zeros(1, _FLAT), + "image_grid_thw": [[1, 4, 4]], + } + assert model.get_input_modality(mm_kwargs) == "image" + + def test_video_model_returns_image_for_image_kwargs(self): + model = SimpleMockViTVideoModel() + mm_kwargs = { + "pixel_values": torch.zeros(1, _FLAT), + "image_grid_thw": [[1, 4, 4]], + } + assert model.get_input_modality(mm_kwargs) == "image" + + def test_video_model_returns_video_for_video_kwargs(self): + model = SimpleMockViTVideoModel() + mm_kwargs = { + "pixel_values_videos": torch.zeros(8, _FLAT), + "video_grid_thw": [[2, 4, 4]], + } + assert model.get_input_modality(mm_kwargs) == "video" + + def test_video_model_config_has_both_modalities(self): + model = SimpleMockViTVideoModel() + cfg = model.get_encoder_cudagraph_config() + assert "image" in cfg.modalities + assert "video" in cfg.modalities + assert cfg.input_key_by_modality["image"] == "pixel_values" + assert cfg.input_key_by_modality["video"] == "pixel_values_videos" + + +# --------------------------------------------------------------------------- +# GPU tests — video capture, replay, fallback, and mixed image+video +# --------------------------------------------------------------------------- + +_VIDEO_MAX_BATCH = 4 +_VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +class TestEncoderCudaGraphVideoReplay: + def setup_method(self): + self.device = torch.device("cuda:0") + self.dtype = torch.float16 + self.model = SimpleMockViTVideoModel().to(self.device).half() + self.mgr = _make_manager_for_gpu( + self.model, + _BUDGETS, + _VIDEO_MAX_BATCH, + self.device, + self.dtype, + max_frames_per_batch=_VIDEO_MAX_FRAMES, + ) + self.mgr.capture() + + # --- capture --- + + def test_capture_creates_one_graph_per_budget(self): + assert len(self.mgr.budget_graphs) == len(_BUDGETS) + assert set(self.mgr.budget_graphs.keys()) == set(_BUDGETS) + + # --- output shape --- + + def test_video_execute_returns_one_tensor_per_video(self): + # T=2, 4x4 → 2*(4//2)*(4//2) = 8 tokens per video + grid_thw = [[2, 4, 4], [2, 4, 4]] + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert len(result) == 2 + + def test_video_output_tokens_per_item(self): + # T=2,4x4 → 8 tokens; T=1,4x4 → 4 tokens + grid_thw = [[2, 4, 4], [1, 4, 4]] + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert result[0].shape == (8, _HIDDEN) + assert result[1].shape == (4, _HIDDEN) + + # --- budget fallback --- + + def test_video_eager_fallback_when_tokens_exceed_all_budgets(self): + # T=2, 18x18 → 2*(18//2)*(18//2) = 162 tokens > max budget 64 + grid_thw = [[2, 18, 18]] + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype) + result = self.mgr.execute(mm_kwargs) + assert result is not None + assert len(result) == 1 + assert result[0].shape == (162, _HIDDEN) + assert self.mgr.graph_misses == 1 + + # --- counters --- + + def test_video_hit_counter_increments_by_num_videos(self): + grid_thw = [[2, 4, 4], [1, 4, 4]] + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype) + self.mgr.execute(mm_kwargs) + assert self.mgr.graph_hits == 2 + + def test_video_miss_counter_increments_for_oversized_video(self): + grid_thw = [[2, 18, 18]] # 162 tokens > 64 + mm_kwargs = _make_video_mm_kwargs(grid_thw, self.device, self.dtype) + self.mgr.execute(mm_kwargs) + assert self.mgr.graph_misses == 1 + + # --- image and video sharing the same manager --- + + def test_image_and_video_share_manager(self): + """Image and video inputs can both be executed through the same manager.""" + img_grid = [[1, 4, 4], [1, 4, 4]] + img_result = self.mgr.execute( + _make_mm_kwargs(img_grid, self.device, self.dtype) + ) + + vid_grid = [[2, 4, 4]] + vid_result = self.mgr.execute( + _make_video_mm_kwargs(vid_grid, self.device, self.dtype) + ) + + assert len(img_result) == 2 + assert len(vid_result) == 1 + assert img_result[0].shape == (4, _HIDDEN) + assert vid_result[0].shape == (8, _HIDDEN) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index de5b78cbfaf..b83b24a27c4 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -35,10 +35,10 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( using the high-level v1 LLM() API only (no manual batching). Strategy: - - Create two LLM engines with identical config except max_num_seqs: 1 vs N. - - Compute a baseline output for the needle prompt with the bs=1 engine. - - For many trials, generate a batch (size N) where the needle appears at a - random position among random filler prompts using the bs=N engine. + - Create a single LLM engine configured for the larger batch limit (N). + - Compute a baseline output for the needle prompt when it is run alone. + - For many trials, generate a mixed batch (size N) where the needle appears + at a random position among random filler prompts using the same engine. - Track how many trials match vs mismatch, and report totals at the end. The test fails if any mismatches occur, but we still dump pass/fail counts. @@ -85,11 +85,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( needle_prompt = "There once was a " - llm_bs1 = None - llm_bsN = None + llm = None try: - # Engine with bs=1 behavior - llm_bs1 = LLM_with_max_seqs( + llm = LLM_with_max_seqs( model=model, max_num_seqs=max_batch_size, gpu_memory_utilization=gpu_mem_util, @@ -99,21 +97,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( ) # Baseline generation for the needle prompt alone. - baseline_out = llm_bs1.generate([needle_prompt], sampling) + baseline_out = llm.generate([needle_prompt], sampling) assert len(baseline_out) == 1 assert len(baseline_out[0].outputs) >= 1 baseline_text = baseline_out[0].outputs[0].text - # Engine with larger batch limit (e.g., 64) - llm_bsN = LLM_with_max_seqs( - model=model, - max_num_seqs=max_batch_size, - gpu_memory_utilization=gpu_mem_util, - max_model_len=max_model_len, - enforce_eager=enforce_eager, - attention_config=attention_config, - ) - mismatches = 0 for trial in range(num_trials): @@ -128,8 +116,8 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( else: prompts.append(_random_prompt(min_random_prompt, max_random_prompt)) - # Generate with the larger-batch engine - outputs = llm_bsN.generate(prompts, sampling) + # Generate with the same engine but in a larger batch. + outputs = llm.generate(prompts, sampling) # Find the needle output by position needle_output = outputs[needle_pos] assert needle_output.prompt == needle_prompt @@ -155,12 +143,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle( finally: # Ensure engines are shutdown to free GPU/VRAM across test sessions - if llm_bs1 is not None: + if llm is not None: with contextlib.suppress(Exception): - llm_bs1.shutdown() - if llm_bsN is not None: - with contextlib.suppress(Exception): - llm_bsN.shutdown() + llm.shutdown() @skip_unsupported diff --git a/tests/v1/distributed/test_eagle_dp.py b/tests/v1/distributed/test_eagle_dp.py index 8557fa0c61d..019f9c6f213 100644 --- a/tests/v1/distributed/test_eagle_dp.py +++ b/tests/v1/distributed/test_eagle_dp.py @@ -20,16 +20,14 @@ if current_platform.is_rocm(): else: ATTN_BACKENDS = ["FLASH_ATTN"] +# On SM<90 (e.g., L4), batch invariance does not support CUDA graphs. +# See https://github.com/vllm-project/vllm/pull/30018 and +# tests/v1/determinism/utils.py for the documented limitation. +IS_DEVICE_CAPABILITY_BELOW_90 = not current_platform.has_device_capability(90) + @pytest.mark.asyncio @pytest.mark.parametrize("attn_backend", ATTN_BACKENDS) -@pytest.mark.xfail( - not current_platform.is_rocm(), - reason="EAGLE + DP > 1 produces wrong outputs when async spec decode " - "correction is active. Root cause under investigation. " - "See: https://github.com/vllm-project/vllm/issues/31913", - strict=False, -) @pytest.mark.xfail( current_platform.is_rocm(), reason="Test may fail on ROCm until batch invariance is enabled. " @@ -37,7 +35,7 @@ else: strict=False, ) async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str): - if not current_platform.is_rocm(): + if not current_platform.is_rocm() and not current_platform.is_xpu(): # This test checks that running a model with and without eagle # leads to identical tokens. # @@ -57,7 +55,7 @@ async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str): engine_args = AsyncEngineArgs( model=target_model, tokenizer_mode="auto", - enforce_eager=False, + enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90, tensor_parallel_size=int(os.getenv("TP_SIZE", 1)), data_parallel_size=DP_SIZE, data_parallel_backend="mp", # ray takes more time diff --git a/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py b/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py index 5cbdc412323..9e000223c14 100644 --- a/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py @@ -5,9 +5,6 @@ This script contains: 1. test lora with speculative decoding for batch inference """ -import random - -import numpy as np import pytest import torch @@ -15,6 +12,7 @@ from vllm import LLM, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory from vllm.lora.request import LoRARequest from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed LORA_TEST_PROMPT_MAP: dict[str, str] = {} @@ -63,10 +61,7 @@ def test_batch_inference_correctness( with monkeypatch.context() as m: # Disable randomness m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8") - torch.manual_seed(SEED) - np.random.seed(SEED) - random.seed(SEED) - torch.cuda.manual_seed_all(SEED) + set_random_seed(SEED) torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index 69a1c38a453..21a651c62ab 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -514,7 +514,6 @@ async def test_header_dp_rank_argument(): serving_render = OpenAIServingRender( model_config=engine.model_config, renderer=engine.renderer, - io_processor=engine.io_processor, model_registry=models.registry, request_logger=None, chat_template=None, diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index d2a9ceb2dee..ab5946ad3ba 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -936,6 +936,13 @@ async def test_engine_core_client_future_utility_async( client.shutdown() +@pytest.mark.parametrize( + "model_name,num_groups", + [ + ("meta-llama/Llama-3.2-1B-Instruct", 1), + ("google/gemma-3-1b-it", 7), + ], +) @pytest.mark.parametrize( "multiprocessing_mode,publisher_config", [(True, "tcp"), (False, "inproc")], @@ -944,12 +951,14 @@ async def test_engine_core_client_future_utility_async( def test_kv_cache_events( multiprocessing_mode: bool, publisher_config, + model_name: str, + num_groups: int, ): block_size = 16 num_blocks = 2 engine_args = EngineArgs( - model=MODEL_NAME, + model=model_name, enforce_eager=True, enable_prefix_caching=True, block_size=block_size, @@ -985,26 +994,29 @@ def test_kv_cache_events( assert result is not None, "No message received" seq, received = result - assert seq == 0, "Sequence number mismatch" - assert len(received.events) == 1, "We should have exactly one BlockStored event" - event = received.events[0] - assert isinstance(event, BlockStored), "We should have a BlockStored event" - assert len(event.block_hashes) == num_blocks, ( - "We should have a BlockStored event with 2 block_hashes" - ) - assert event.block_size == block_size, ( - "Block size should be the same as the block size" - ) - assert event.parent_block_hash is None, "Parent block hash should be None" - assert event.lora_id is None, "Lora id should be None" - assert event.lora_name is None, "Lora name should be None" - assert len(event.token_ids) == num_blocks * block_size, ( - "Token ids should be the same as the custom tokens" - ) - assert event.token_ids == custom_tokens, ( - "Token ids should be the same as the custom tokens" + assert len(received.events) == num_groups, ( + f"Expected {num_groups} BlockStored event(s), got {len(received.events)}" ) + + for index, event in enumerate(received.events): + assert isinstance(event, BlockStored), "We should have a BlockStored event" + assert len(event.block_hashes) == num_blocks, ( + "We should have a BlockStored event with 2 block_hashes" + ) + assert event.block_size == block_size, ( + "Block size should be the same as the block size" + ) + assert event.parent_block_hash is None, "Parent block hash should be None" + assert event.lora_id is None, "Lora id should be None" + assert event.lora_name is None, "Lora name should be None" + assert len(event.token_ids) == num_blocks * block_size, ( + "Token ids should be the same as the custom tokens" + ) + assert event.token_ids == custom_tokens, ( + "Token ids should be the same as the custom tokens" + ) + assert event.group_idx == index finally: client.shutdown() subscriber.close() diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index ece48e009d2..1919349790f 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -84,6 +84,7 @@ def test_incremental_detokenization( engine_core = MockEngineCore( tokens_list=dummy_test_vectors.generation_tokens, + prompts_list=dummy_test_vectors.prompt_tokens, request_ids=[req.request_id for req in requests], ) @@ -506,6 +507,7 @@ def test_logprobs_processor( engine_core = MockEngineCore( tokens_list=dummy_test_vectors.generation_tokens, + prompts_list=dummy_test_vectors.prompt_tokens, generated_logprobs_raw=None if num_sample_logprobs is None else dummy_test_vectors.generation_logprobs, @@ -691,6 +693,7 @@ def test_stop_token( engine_core = MockEngineCore( tokens_list=[generation_tokens], + prompts_list=dummy_test_vectors.prompt_tokens, generated_logprobs_raw=[generation_logprobs] if do_logprobs else None, prompt_logprobs_raw=None, eos_token_id=sampling_params.eos_token_id, @@ -794,6 +797,7 @@ def test_stop_string( engine_core = MockEngineCore( tokens_list=dummy_test_vectors.generation_tokens, + prompts_list=dummy_test_vectors.prompt_tokens, generated_logprobs_raw=dummy_test_vectors.generation_logprobs if num_sample_logprobs else None, @@ -917,6 +921,7 @@ def test_iteration_stats(dummy_test_vectors): engine_core = MockEngineCore( dummy_test_vectors.generation_tokens, + dummy_test_vectors.prompt_tokens, request_ids=[req.request_id for req in requests], ) @@ -927,7 +932,7 @@ def test_iteration_stats(dummy_test_vectors): inactive_request = requests[num_active] # First iteration has 2 prefills. - outputs = engine_core.get_outputs()[:num_active] + outputs = engine_core.get_outputs(num_active) iteration_stats = IterationStats() output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats) total_prompt_tokens = sum( @@ -941,7 +946,7 @@ def test_iteration_stats(dummy_test_vectors): assert iteration_stats.num_generation_tokens == num_active # Just decodes in this step. - outputs = engine_core.get_outputs()[:num_active] + outputs = engine_core.get_outputs(num_active) iteration_stats = IterationStats() output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats) @@ -951,7 +956,7 @@ def test_iteration_stats(dummy_test_vectors): # Add a new request - prefill and 2 decodes in this step. output_processor.add_request(inactive_request, None) num_active += 1 - outputs = engine_core.get_outputs()[:num_active] + outputs = engine_core.get_outputs(num_active) iteration_stats = IterationStats() output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats) total_prompt_tokens = len(dummy_test_vectors.prompt_tokens[num_active - 1]) @@ -960,7 +965,7 @@ def test_iteration_stats(dummy_test_vectors): assert iteration_stats.num_generation_tokens == num_active # Just decodes in this step. - outputs = engine_core.get_outputs()[:num_active] + outputs = engine_core.get_outputs(num_active) iteration_stats = IterationStats() output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats) @@ -1003,6 +1008,7 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors): engine_core = MockEngineCore( dummy_test_vectors.generation_tokens, + dummy_test_vectors.prompt_tokens, request_ids=[req.request_id for req in requests], ) diff --git a/tests/v1/engine/utils.py b/tests/v1/engine/utils.py index de953a58843..013e73bd8e4 100644 --- a/tests/v1/engine/utils.py +++ b/tests/v1/engine/utils.py @@ -11,6 +11,7 @@ from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from vllm.engine.arg_utils import EngineArgs from vllm.v1.engine import EngineCoreOutput, FinishReason +from vllm.v1.metrics.stats import PrefillStats from vllm.v1.outputs import LogprobsLists, LogprobsTensors GeneralTokenizerType: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast @@ -330,6 +331,7 @@ class MockEngineCore: def __init__( self, tokens_list: list[list[int]], + prompts_list: list[list[int]], # For each request, for each sampled token offset, # a tuple of # (list of topk token ids, list of sample logprob vals, rank) @@ -346,12 +348,13 @@ class MockEngineCore: ) -> None: self.num_requests = len(tokens_list) self.tokens_list = tokens_list - self.current_idx = 0 + self.prompts_list = prompts_list self.generated_logprobs_raw = generated_logprobs_raw self.do_logprobs = generated_logprobs_raw is not None self.prompt_logprobs_raw = prompt_logprobs_raw self.do_prompt_logprobs = prompt_logprobs_raw is not None self.request_finished = [False for _ in range(self.num_requests)] + self.request_token_idx = [0 for _ in range(self.num_requests)] self.eos_token_id = eos_token_id self.stop_token_ids = stop_token_ids self.request_ids = ( @@ -360,14 +363,18 @@ class MockEngineCore: else [f"request-{i}" for i in range(self.num_requests)] ) - def get_outputs(self) -> list[EngineCoreOutput]: + def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]: do_logprobs = self.do_logprobs do_prompt_logprobs = self.do_prompt_logprobs - token_idx = self.current_idx outputs = [] - for req_idx, token_ids in enumerate(self.tokens_list): + for req_idx, (token_ids, prompt_token_ids) in enumerate( + zip(self.tokens_list, self.prompts_list) + ): + if num_active != -1 and req_idx >= num_active: + break if not self.request_finished[req_idx]: + token_idx = self.request_token_idx[req_idx] if do_logprobs: assert self.generated_logprobs_raw is not None (logprobs_token_ids_, logprobs_, sampled_token_ranks_) = ( @@ -381,19 +388,32 @@ class MockEngineCore: else: logprobs = None if do_prompt_logprobs: - if self.current_idx == 0: + if token_idx == 0: assert self.prompt_logprobs_raw is not None prompt_logprobs = self.prompt_logprobs_raw[req_idx] else: prompt_logprobs = None else: prompt_logprobs = None + + # Add prefill_stats on first output (prefill) for this request + if token_idx == 0: + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=len(prompt_token_ids), + num_local_cached_tokens=0, + num_external_cached_tokens=0, + ) + else: + prefill_stats = None + new_token_id = token_ids[token_idx] output = EngineCoreOutput( request_id=self.request_ids[req_idx], new_token_ids=[new_token_id], new_logprobs=logprobs, new_prompt_logprobs_tensors=prompt_logprobs, + prefill_stats=prefill_stats, ) if token_idx == len(token_ids) - 1: output.finish_reason = FinishReason.LENGTH @@ -407,5 +427,6 @@ class MockEngineCore: self.request_finished[req_idx] = True outputs.append(output) - self.current_idx += 1 + self.request_token_idx[req_idx] += 1 + return outputs diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh new file mode 100755 index 00000000000..2e71858983e --- /dev/null +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh @@ -0,0 +1,205 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Integration accuracy test for MultiConnector (NixlConnector + OffloadingConnector). +# +# Launches a P/D setup where both prefill and decode instances use MultiConnector +# wrapping NixlConnector and OffloadingConnector, then runs gsm8k accuracy via +# test_accuracy.py. +# +# By default runs two configurations: +# 1. Normal KV layout (NixlConnector without cross-layer blocks) +# 2. Cross-layer KV layout (NixlConnector with enable_cross_layers_blocks) +# +# Usage: +# bash tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh +# +# Environment variables: +# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) +# GPU_MEMORY_UTILIZATION - GPU memory fraction (default: 0.6) +# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve +# SKIP_CROSS_LAYERS - set to 1 to skip the cross-layer layout test +# SKIP_NORMAL_LAYOUT - set to 1 to skip the normal layout test +set -xe + +# ── Configuration ──────────────────────────────────────────────────────── + +MODEL_NAMES=${MODEL_NAMES:-} +if [[ -n "$MODEL_NAMES" ]]; then + MODELS=("$MODEL_NAMES") +else + MODELS=("Qwen/Qwen3-0.6B") +fi + +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.6} +BLOCK_SIZE=${BLOCK_SIZE:-128} +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} + +GIT_ROOT=$(git rev-parse --show-toplevel) +SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "") + +# ── KV transfer configs ───────────────────────────────────────────────── + +# Normal layout: OffloadingConnector prefers cross-layer but NixlConnector +# does not, so MultiConnector.prefer_cross_layer_blocks = False. +KV_CONFIG_NORMAL='{ + "kv_connector":"MultiConnector", + "kv_role":"kv_both", + "kv_connector_extra_config":{ + "connectors":[ + {"kv_connector":"NixlConnector","kv_role":"kv_both"}, + {"kv_connector":"OffloadingConnector","kv_role":"kv_both", + "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}} + ] + } +}' +# Remove whitespace for CLI safety +KV_CONFIG_NORMAL=$(echo "$KV_CONFIG_NORMAL" | tr -d '[:space:]') + +# Cross-layer layout: both connectors prefer cross-layer blocks. +KV_CONFIG_CROSS_LAYERS='{ + "kv_connector":"MultiConnector", + "kv_role":"kv_both", + "kv_connector_extra_config":{ + "connectors":[ + {"kv_connector":"NixlConnector","kv_role":"kv_both", + "kv_connector_extra_config":{"enable_cross_layers_blocks":"True"}}, + {"kv_connector":"OffloadingConnector","kv_role":"kv_both", + "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}} + ] + } +}' +KV_CONFIG_CROSS_LAYERS=$(echo "$KV_CONFIG_CROSS_LAYERS" | tr -d '[:space:]') + +# ── Helpers ────────────────────────────────────────────────────────────── + +trap 'kill $(jobs -pr) 2>/dev/null' SIGINT SIGTERM EXIT + +wait_for_server() { + local port=$1 + timeout 1200 bash -c " + until curl -s localhost:${port}/v1/completions > /dev/null; do + sleep 1 + done" && return 0 || return 1 +} + +cleanup_instances() { + echo "Cleaning up any running vLLM instances..." + pkill -f "vllm serve" || true + sleep 2 +} + +get_num_gpus() { + if [[ "$SMI_BIN" == *"nvidia"* ]]; then + $SMI_BIN --query-gpu=name --format=csv,noheader | wc -l + elif [[ "$SMI_BIN" == *"rocm"* ]]; then + $SMI_BIN -l | grep -c GPU + else + echo "1" + fi +} + +# ── Run tests for one model with a given KV config ─────────────────────── + +run_tests_for_model() { + local model_name=$1 + local kv_config=$2 + local label=$3 + + echo "================================================================" + echo "Testing model: $model_name ($label)" + echo "KV config: $kv_config" + echo "================================================================" + + local PREFILL_PORT=8100 + local DECODE_PORT=8200 + local PREFILL_GPU=0 + local DECODE_GPU=1 + local PREFILL_SIDE_CHANNEL_PORT=5559 + local DECODE_SIDE_CHANNEL_PORT=5659 + + # ── Start prefill instance ── + echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT" + BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \ + VLLM_KV_CACHE_LAYOUT='HND' \ + UCX_NET_DEVICES=all \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \ + vllm serve $model_name \ + --port $PREFILL_PORT \ + --enforce-eager \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ + --tensor-parallel-size 1 \ + --kv-transfer-config '$kv_config'" + + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" + for arg in "${extra_args[@]}"; do + BASE_CMD="${BASE_CMD} $arg" + done + fi + eval "$BASE_CMD &" + + # ── Start decode instance ── + echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT" + BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \ + VLLM_KV_CACHE_LAYOUT='HND' \ + UCX_NET_DEVICES=all \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \ + vllm serve $model_name \ + --port $DECODE_PORT \ + --enforce-eager \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \ + --tensor-parallel-size 1 \ + --kv-transfer-config '$kv_config'" + + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" + for arg in "${extra_args[@]}"; do + BASE_CMD="${BASE_CMD} $arg" + done + fi + eval "$BASE_CMD &" + + # ── Wait for servers ── + echo "Waiting for prefill instance on port $PREFILL_PORT to start..." + wait_for_server "$PREFILL_PORT" + echo "Waiting for decode instance on port $DECODE_PORT to start..." + wait_for_server "$DECODE_PORT" + + # ── Start proxy ── + PROXY_CMD="python3 ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py --port 8192" + PROXY_CMD+=" --prefiller-hosts localhost" + PROXY_CMD+=" --prefiller-ports $PREFILL_PORT" + PROXY_CMD+=" --decoder-hosts localhost" + PROXY_CMD+=" --decoder-ports $DECODE_PORT" + + echo "Starting proxy server with command: $PROXY_CMD" + $PROXY_CMD & + sleep 5 + + # ── Run accuracy test ── + echo "Running accuracy tests for $model_name ($label)" + TEST_MODEL=$model_name python3 -m pytest -s -x \ + "${GIT_ROOT}"/tests/v1/kv_connector/nixl_integration/test_accuracy.py + + # ── Cleanup ── + cleanup_instances + sleep 3 +} + +# ── Main ───────────────────────────────────────────────────────────────── + +for model in "${MODELS[@]}"; do + if [[ -z "${SKIP_NORMAL_LAYOUT:-}" ]]; then + run_tests_for_model "$model" "$KV_CONFIG_NORMAL" "MultiConnector normal layout" + fi + + if [[ -z "${SKIP_CROSS_LAYERS:-}" ]]; then + run_tests_for_model "$model" "$KV_CONFIG_CROSS_LAYERS" "MultiConnector cross-layer layout" + fi +done + +echo "All MultiConnector accuracy tests passed!" diff --git a/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh new file mode 100755 index 00000000000..a80950b3413 --- /dev/null +++ b/tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh @@ -0,0 +1,174 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector). +# +# Launches a P/D setup where both prefill and decode instances use MultiConnector +# wrapping NixlConnector and OffloadingConnector, then runs scenario-based edge +# case tests including Prometheus metrics validation. +# +# Tests cover: block-size boundaries, decode-side cache-hit scenarios +# (cold / full / partial), direct decode (control), and prefill-side CPU +# offload recovery after GPU eviction. +# +# Usage: +# bash tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh +# +# Environment variables: +# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B) +# KV_CACHE_MEMORY_BYTES - GPU KV cache size in bytes (default: 268435456 = 256 MiB) +# BLOCK_SIZE - KV cache block size (default: 128) +# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve +set -xe + +# ── Configuration ──────────────────────────────────────────────────────── + +MODEL_NAMES=${MODEL_NAMES:-} +if [[ -n "$MODEL_NAMES" ]]; then + MODELS=("$MODEL_NAMES") +else + MODELS=("Qwen/Qwen3-0.6B") +fi + +KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-268435456} # 256 MiB +MAX_MODEL_LEN=${MAX_MODEL_LEN:-2048} +BLOCK_SIZE=${BLOCK_SIZE:-128} +VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-} + +GIT_ROOT=$(git rev-parse --show-toplevel) + +# ── KV transfer config ────────────────────────────────────────────────── + +KV_CONFIG='{ + "kv_connector":"MultiConnector", + "kv_role":"kv_both", + "kv_connector_extra_config":{ + "connectors":[ + {"kv_connector":"NixlConnector","kv_role":"kv_both"}, + {"kv_connector":"OffloadingConnector","kv_role":"kv_both", + "kv_connector_extra_config":{"cpu_bytes_to_use":2147483648}} + ] + } +}' +KV_CONFIG=$(echo "$KV_CONFIG" | tr -d '[:space:]') + +# ── Helpers ────────────────────────────────────────────────────────────── + +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT + +wait_for_server() { + local port=$1 + timeout 1200 bash -c " + until curl -s localhost:${port}/v1/completions > /dev/null; do + sleep 1 + done" && return 0 || return 1 +} + +cleanup_instances() { + echo "Cleaning up any running vLLM instances and proxy..." + pkill -f "vllm serve" || true + pkill -f "toy_proxy_server.py" || true + sleep 2 +} + +# ── Run tests for one model ────────────────────────────────────────────── + +run_tests_for_model() { + local model_name=$1 + + echo "================================================================" + echo "Testing model: $model_name (MultiConnector edge cases)" + echo "================================================================" + + local PREFILL_PORT=8100 + local DECODE_PORT=8200 + local PROXY_PORT=8192 + local PREFILL_GPU=0 + local DECODE_GPU=1 + local PREFILL_SIDE_CHANNEL_PORT=5559 + local DECODE_SIDE_CHANNEL_PORT=5659 + + # ── Start prefill instance ── + echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT" + BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \ + VLLM_KV_CACHE_LAYOUT='HND' \ + UCX_NET_DEVICES=all \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \ + vllm serve \"$model_name\" \ + --port $PREFILL_PORT \ + --enforce-eager \ + --block-size ${BLOCK_SIZE} \ + --max-model-len $MAX_MODEL_LEN \ + --kv-cache-memory-bytes $KV_CACHE_MEMORY_BYTES \ + --tensor-parallel-size 1 \ + --kv-transfer-config '$KV_CONFIG'" + + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" + for arg in "${extra_args[@]}"; do + BASE_CMD="${BASE_CMD} $arg" + done + fi + eval "$BASE_CMD &" + + # ── Start decode instance ── + echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT" + BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \ + VLLM_KV_CACHE_LAYOUT='HND' \ + UCX_NET_DEVICES=all \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \ + vllm serve \"$model_name\" \ + --port $DECODE_PORT \ + --enforce-eager \ + --block-size ${BLOCK_SIZE} \ + --max-model-len $MAX_MODEL_LEN \ + --kv-cache-memory-bytes $KV_CACHE_MEMORY_BYTES \ + --tensor-parallel-size 1 \ + --kv-transfer-config '$KV_CONFIG'" + + if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then + IFS=',' read -r -a extra_args <<< "$VLLM_SERVE_EXTRA_ARGS" + for arg in "${extra_args[@]}"; do + BASE_CMD="${BASE_CMD} $arg" + done + fi + eval "$BASE_CMD &" + + # ── Wait for servers ── + echo "Waiting for prefill instance on port $PREFILL_PORT to start..." + wait_for_server "$PREFILL_PORT" + echo "Waiting for decode instance on port $DECODE_PORT to start..." + wait_for_server "$DECODE_PORT" + + # ── Start proxy ── + echo "Starting proxy server on port $PROXY_PORT" + python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \ + --port "$PROXY_PORT" \ + --prefiller-hosts localhost \ + --prefiller-ports "$PREFILL_PORT" \ + --decoder-hosts localhost \ + --decoder-ports "$DECODE_PORT" & + sleep 5 + + # ── Run edge case tests ── + echo "Running MultiConnector edge case tests for $model_name" + PREFILL_PORT=$PREFILL_PORT \ + DECODE_PORT=$DECODE_PORT \ + PROXY_PORT=$PROXY_PORT \ + BLOCK_SIZE=$BLOCK_SIZE \ + python3 -m pytest -s -x \ + "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py" + + # ── Cleanup ── + cleanup_instances + sleep 3 +} + +# ── Main ───────────────────────────────────────────────────────────────── + +for model in "${MODELS[@]}"; do + run_tests_for_model "$model" +done + +echo "All MultiConnector edge case tests passed!" diff --git a/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py new file mode 100644 index 00000000000..f109190a4a0 --- /dev/null +++ b/tests/v1/kv_connector/nixl_integration/test_multi_connector_edge_cases.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Integration edge-case tests for MultiConnector (NixlConnector + OffloadingConnector). + +Tests cover: + - Output correctness across block-size boundaries (proxy vs direct prefill). + - Decode-side Prometheus metrics validation (local_cache_hit, + external_kv_transfer, local_compute) for cold/warm/partial cache scenarios. + - Prefill-side CPU offload recovery after GPU cache eviction. + +Requires running servers started by run_multi_connector_edge_case_test.sh. +""" + +import os +import time +import urllib.request + +import openai +import regex as re + +# ── Server configuration from environment ───────────────────────────────── + +PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost") +PREFILL_PORT = os.environ["PREFILL_PORT"] +DECODE_HOST = os.getenv("DECODE_HOST", "localhost") +DECODE_PORT = os.environ["DECODE_PORT"] +PROXY_HOST = os.getenv("PROXY_HOST", "localhost") +PROXY_PORT = os.environ["PROXY_PORT"] +BLOCK_SIZE = int(os.getenv("BLOCK_SIZE", "128")) + +# ── OpenAI clients ──────────────────────────────────────────────────────── + +decode_client = openai.OpenAI( + api_key="EMPTY", + base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1", +) +prefill_client = openai.OpenAI( + api_key="EMPTY", + base_url=f"http://{PREFILL_HOST}:{PREFILL_PORT}/v1", +) +proxy_client = openai.OpenAI( + api_key="EMPTY", + base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1", +) + +_MODEL = None + + +def _get_model() -> str: + global _MODEL + if _MODEL is None: + models = decode_client.models.list() + _MODEL = models.data[0].id + return _MODEL + + +def _complete(client: openai.OpenAI, prompt: str, max_tokens: int = 20): + """Send a completion request and return (text, prompt_tokens).""" + resp = client.completions.create( + model=_get_model(), + prompt=prompt, + max_tokens=max_tokens, + temperature=0, + ) + return resp.choices[0].text, resp.usage.prompt_tokens + + +# ── Prometheus metrics helpers ──────────────────────────────────────────── + +_METRIC_RE = re.compile( + r'vllm:prompt_tokens_by_source_total\{.*?source="([^"]+)".*?\}\s+' + r"([\d.eE+\-]+)" +) + + +def _fetch_metrics(host: str, port: str) -> dict[str, float]: + """Scrape prompt_tokens_by_source counters from a vLLM server.""" + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode() + result = { + "local_compute": 0.0, + "local_cache_hit": 0.0, + "external_kv_transfer": 0.0, + } + for m in _METRIC_RE.finditer(body): + source, val = m.group(1), float(m.group(2)) + if source in result: + result[source] += val + return result + + +def _fetch_decode_metrics() -> dict[str, float]: + return _fetch_metrics(DECODE_HOST, DECODE_PORT) + + +def _fetch_prefill_metrics() -> dict[str, float]: + return _fetch_metrics(PREFILL_HOST, PREFILL_PORT) + + +_NIXL_BYTES_RE = re.compile(r"vllm:nixl_bytes_transferred_sum\b.*?\s+([\d.eE+\-]+)") + + +def _fetch_nixl_bytes(host: str, port: str) -> float: + """Scrape total NIXL bytes transferred from a vLLM server.""" + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode() + total = 0.0 + for m in _NIXL_BYTES_RE.finditer(body): + total += float(m.group(1)) + return total + + +_OFFLOAD_BYTES_RE = re.compile( + r'vllm:kv_offload_total_bytes_total\{.*?transfer_type="([^"]+)".*?\}\s+' + r"([\d.eE+\-]+)" +) + + +def _fetch_offload_bytes(host: str, port: str) -> dict[str, float]: + """Scrape kv_offload_total_bytes counters (CPU_to_GPU / GPU_to_CPU).""" + body = urllib.request.urlopen(f"http://{host}:{port}/metrics").read().decode() + result = {"CPU_to_GPU": 0.0, "GPU_to_CPU": 0.0} + for m in _OFFLOAD_BYTES_RE.finditer(body): + transfer_type, val = m.group(1), float(m.group(2)) + if transfer_type in result: + result[transfer_type] += val + return result + + +def _metrics_delta(before: dict, after: dict) -> dict[str, float]: + return {k: after.get(k, 0) - before.get(k, 0) for k in before} + + +# ── Prompts (unique per test to avoid cross-test cache interference) ────── + +SHORT_PROMPT = "Red Hat is " + +MEDIUM_PROMPT = ( + "Red Hat is the best company in the world to work for because it works " + "on open source software, which means that all the contributions are " + "delivered to the community. As a result," +) + + +def _make_prompt(n_tokens: int) -> str: + """Build a prompt of ~n_tokens tokens (1 word ~ 1 token).""" + return "word " * n_tokens + + +BLOCK_BOUNDARY_PROMPT = _make_prompt(BLOCK_SIZE) +ABOVE_BOUNDARY_PROMPT = _make_prompt(BLOCK_SIZE + 2) +MULTI_BLOCK_PROMPT = _make_prompt(BLOCK_SIZE * 4) + +FULL_CACHE_HIT_PROMPT = ( # noqa: E501 + "The history of computing begins with Charles Babbage who designed the " + "Analytical Engine in the 1830s which is considered the first general " + "purpose computer design in history. Ada Lovelace is widely regarded as " + "the first computer programmer for her work on the Analytical Engine. " + "The modern era of computing began with Alan Turing who formalized the " + "concept of computation with his Turing machine in 1936. During World " + "War Two Turing worked at Bletchley Park to break the Enigma cipher. " + "After the war the first electronic computers were built including ENIAC " + "at the University of Pennsylvania and Colossus at Bletchley Park. " + "These early machines filled entire rooms and used vacuum tubes for logic. " + "The invention of the transistor at Bell Labs in 1947 revolutionized " + "computing by making smaller and more reliable machines possible. " + "The integrated circuit followed in the late 1950s combining multiple " + "transistors on a single chip. This led to the microprocessor in the 1970s " + "and eventually to the personal computer revolution of the 1980s." +) + +PARTIAL_CACHE_PREFIX = ( # noqa: E501 + "Machine learning has transformed the field of artificial intelligence " + "by enabling computers to learn patterns from data without being " + "explicitly programmed for every task. The field has evolved dramatically " + "since its inception in the 1950s when Arthur Samuel coined the term while " + "working at IBM. Early approaches focused on symbolic reasoning and expert " + "systems that encoded human knowledge as rules. The statistical revolution " + "of the 1990s shifted the paradigm toward data driven methods. Support " + "vector machines and random forests became popular for classification tasks. " + "The breakthrough of deep learning in 2012 with AlexNet winning ImageNet " + "changed everything. Neural networks with many layers could automatically " + "learn hierarchical feature representations from raw data." +) +PARTIAL_CACHE_EXTENDED = PARTIAL_CACHE_PREFIX + ( + " Transformers have become the dominant architecture for natural language " + "processing tasks including translation, summarization, and generation. " + "The attention mechanism allows models to weigh the importance of different " + "parts of the input sequence. Large language models like GPT and BERT " + "demonstrated that pre-training on massive text corpora followed by fine " + "tuning on specific tasks could achieve state of the art results across " + "a wide range of benchmarks. Scaling laws suggest that larger models " + "trained on more data continue to improve in capability." +) + +# ═══════════════════════════════════════════════════════════════════════════ +# Output correctness across block-size boundaries (decode-side metrics) +# +# Each test sends via proxy, verifies output matches prefill_direct at +# temperature=0, and checks decode-side metrics for NIXL transfer. +# ═══════════════════════════════════════════════════════════════════════════ + + +def test_short_prompt_correctness(): + """Short prompt (< block_size): output matches prefill, NIXL used.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, _ = _complete(proxy_client, SHORT_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, SHORT_PROMPT) + print(f"SHORT PROMPT: {proxy_text=}, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_block_boundary_correctness(): + """Exactly block_size tokens: output matches prefill, NIXL used.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, pt = _complete(proxy_client, BLOCK_BOUNDARY_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, BLOCK_BOUNDARY_PROMPT) + print(f"BLOCK BOUNDARY: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_above_block_boundary_correctness(): + """Just above block_size (partial second block): output matches prefill.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, pt = _complete(proxy_client, ABOVE_BOUNDARY_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, ABOVE_BOUNDARY_PROMPT) + print(f"ABOVE BOUNDARY: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_multi_block_correctness(): + """Multi-block prompt (~4x block_size): output matches prefill.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, pt = _complete(proxy_client, MULTI_BLOCK_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + prefill_text, _ = _complete(prefill_client, MULTI_BLOCK_PROMPT) + print(f"MULTI BLOCK: {pt} prompt tokens, nixl_bytes_delta={n1 - n0}") + assert proxy_text == prefill_text + assert d["external_kv_transfer"] > 0, ( + "NIXL transfer did not occur — decode may have silently fallen back " + "to local compute" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Decode-side KV source validation via Prometheus metrics +# +# Scrape vllm:prompt_tokens_by_source_total from the DECODE server to +# verify which code path (GPU prefix, NIXL, local compute) was exercised. +# ═══════════════════════════════════════════════════════════════════════════ + + +def test_cold_decode_no_cache_hit_metrics(): + """Cold decode: external_kv_transfer==P, local_cache_hit==0, local_compute==0.""" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, P = _complete(proxy_client, MEDIUM_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + print(f"COLD DECODE: {P} prompt tokens, metrics delta: {d}") + print(f" nixl_bytes_delta={n1 - n0}") + assert len(proxy_text) > 0, "proxy returned empty response" + assert d["external_kv_transfer"] == P, ( + f"expected external_kv_transfer={P}, got {d['external_kv_transfer']}" + ) + assert d["local_compute"] == 0, ( + f"expected local_compute=0, got {d['local_compute']}" + ) + assert d["local_cache_hit"] == 0, ( + f"expected local_cache_hit=0, got {d['local_cache_hit']}" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase, got delta={n1 - n0}" + ) + + +def test_full_decode_gpu_cache_hit_metrics(): + """Prime decode, resend via proxy: local_cache_hit==cached blocks.""" + decode_text, _ = _complete(decode_client, FULL_CACHE_HIT_PROMPT) + + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, P = _complete(proxy_client, FULL_CACHE_HIT_PROMPT) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + cached = (P // BLOCK_SIZE) * BLOCK_SIZE + expected_nixl = P - cached + + print(f"FULL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}") + print(f" metrics delta: {d}, nixl_bytes_delta={n1 - n0}") + assert len(proxy_text) > 0, "proxy returned empty response" + assert d["local_cache_hit"] == cached, ( + f"expected local_cache_hit={cached}, got {d['local_cache_hit']}" + ) + assert d["external_kv_transfer"] == expected_nixl, ( + f"expected external_kv_transfer={expected_nixl}, " + f"got {d['external_kv_transfer']}" + ) + assert d["local_compute"] == 0, ( + f"expected local_compute=0, got {d['local_compute']}" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase (partial NIXL for " + f"uncached tail), got delta={n1 - n0}" + ) + + +def test_partial_decode_gpu_cache_hit_metrics(): + """Prime with prefix, extend via proxy: partial local_cache_hit.""" + _, prefix_tokens = _complete(decode_client, PARTIAL_CACHE_PREFIX) + cached = (prefix_tokens // BLOCK_SIZE) * BLOCK_SIZE + assert cached >= BLOCK_SIZE, ( + f"PARTIAL_CACHE_PREFIX too short ({prefix_tokens} tokens) for partial " + f"cache hit test with block_size={BLOCK_SIZE}" + ) + + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + proxy_text, P = _complete(proxy_client, PARTIAL_CACHE_EXTENDED) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + expected_nixl = P - cached + + print(f"PARTIAL CACHE HIT: {P} tokens, cached={cached}, nixl={expected_nixl}") + print(f" metrics delta: {d}, nixl_bytes_delta={n1 - n0}") + assert len(proxy_text) > 0, "proxy returned empty response" + assert d["external_kv_transfer"] == expected_nixl, ( + f"expected external_kv_transfer={expected_nixl}, " + f"got {d['external_kv_transfer']}" + ) + assert d["local_cache_hit"] == cached, ( + f"expected local_cache_hit={cached}, got {d['local_cache_hit']}" + ) + assert d["local_compute"] == 0, ( + f"expected local_compute=0, got {d['local_compute']}" + ) + assert n1 - n0 > 0, ( + f"expected nixl_bytes_transferred to increase (NIXL for uncached " + f"tail), got delta={n1 - n0}" + ) + + +def test_decode_direct_all_local_compute(): + """Direct decode (no proxy): local_compute==P, no transfers.""" + prompt = "The speed of light is approximately" + n0 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + m0 = _fetch_decode_metrics() + text, P = _complete(decode_client, prompt) + time.sleep(1) + m1 = _fetch_decode_metrics() + n1 = _fetch_nixl_bytes(DECODE_HOST, DECODE_PORT) + d = _metrics_delta(m0, m1) + + print(f"DIRECT DECODE: {text!r} ({P} tokens), metrics delta: {d}") + print(f" nixl_bytes_delta={n1 - n0}") + assert len(text.strip()) > 0, "empty output from direct decode" + assert d["local_compute"] == P, ( + f"expected local_compute={P}, got {d['local_compute']}" + ) + assert d["external_kv_transfer"] == 0, ( + f"expected external_kv_transfer=0, got {d['external_kv_transfer']}" + ) + assert n1 - n0 == 0, ( + f"expected no nixl_bytes_transferred for direct decode, got delta={n1 - n0}" + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Prefill-side CPU offload validation via Prometheus metrics +# +# Scrape vllm:prompt_tokens_by_source_total from the PREFILL server. +# Exercises the OffloadingConnector read path: after GPU cache eviction, +# the OffloadingConnector restores KV from CPU (NixlConnector cannot help +# for direct requests without kv_transfer_params). +# ═══════════════════════════════════════════════════════════════════════════ + +EVICTION_PROMPT = ( # noqa: E501 + "Quantum computing leverages quantum mechanical phenomena like " + "superposition and entanglement to perform computations that would be " + "intractable for classical computers. This has implications for " + "cryptography, drug discovery, and optimization problems. Richard Feynman " + "first proposed the idea of quantum computing in 1982 when he observed " + "that simulating quantum systems on classical computers was exponentially " + "hard. Peter Shor developed a quantum algorithm for factoring large " + "numbers in polynomial time which threatens RSA encryption. Grover search " + "algorithm provides a quadratic speedup for unstructured search problems. " + "Companies like IBM Google and Rigetti are building quantum processors " + "with increasing numbers of qubits. Error correction remains a major " + "challenge as quantum states are extremely fragile and prone to decoherence." +) + + +def test_prefill_cpu_offload_after_gpu_eviction(): + """Prefill-side: evict GPU, re-request directly, CPU offload restores KV.""" + text1, P = _complete(prefill_client, EVICTION_PROMPT, max_tokens=30) + + for i in range(100): + _complete(prefill_client, f"Eviction prompt number {i}: " + _make_prompt(200)) + + ob0 = _fetch_offload_bytes(PREFILL_HOST, PREFILL_PORT) + m0 = _fetch_prefill_metrics() + text2, _ = _complete(prefill_client, EVICTION_PROMPT, max_tokens=30) + + cpu_to_gpu_delta = 0.0 + for _ in range(10): + time.sleep(1) + ob1 = _fetch_offload_bytes(PREFILL_HOST, PREFILL_PORT) + cpu_to_gpu_delta = ob1["CPU_to_GPU"] - ob0["CPU_to_GPU"] + if cpu_to_gpu_delta > 0: + break + + m1 = _fetch_prefill_metrics() + d = _metrics_delta(m0, m1) + + print(f"PREFILL CPU OFFLOAD: run1={text1[:60]!r}, run2={text2[:60]!r}") + print(f" prefill metrics delta: {d}") + print(f" cpu_to_gpu bytes delta: {cpu_to_gpu_delta}") + assert text1 == text2, f"inconsistent after eviction: {text1=!r}, {text2=!r}" + assert cpu_to_gpu_delta > 0, ( + f"expected cpu_to_gpu bytes > 0 (OffloadingConnector should restore " + f"KV from CPU to GPU), got {cpu_to_gpu_delta}" + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 291d0574e50..bdc81dc1ac3 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -124,12 +124,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool): return [BlockHash(str(i).encode()) for i in int_hashes] def take_events() -> Iterable[OffloadingEvent]: - yield OffloadingEvent( - keys=to_keys([1, 2, 3]), block_size=16, medium="A", removed=False - ) - yield OffloadingEvent( - keys=to_keys([4, 5, 6]), block_size=32, medium="B", removed=True - ) + yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False) + yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True) runner.manager.take_events.side_effect = take_events events = list(runner.scheduler_connector.take_events()) @@ -137,7 +133,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): event = events[0] assert isinstance(event, BlockStored) assert event.block_hashes == to_hashes([1, 2, 3]) - assert event.block_size == 16 + assert event.block_size == 0 assert event.medium == "A" assert event.token_ids == [] assert event.parent_block_hash is None diff --git a/tests/v1/kv_connector/unit/test_mooncake_connector.py b/tests/v1/kv_connector/unit/test_mooncake_connector.py index f21f8ecdc5c..7b6fe3af0ce 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_connector.py +++ b/tests/v1/kv_connector/unit/test_mooncake_connector.py @@ -609,6 +609,51 @@ def test_register_kv_caches(): assert bl == tensor1[0].nbytes // tensor1.shape[1] +def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes(): + """Mixed MLA+Eagle caches should register by byte length, not shape.""" + + vllm_config = create_vllm_config( + kv_connector="MooncakeConnector", kv_role="kv_consumer" + ) + + with ( + set_current_vllm_config(vllm_config), + patch_worker_dependencies(), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event" + ), + patch( + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread" + ) as mock_thread, + ): + connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER) + worker = connector.connector_worker + mock_thread.return_value.is_alive.return_value = False + + worker.use_mla = True + worker.kv_topo.is_mla = True + + # MLA cache tensor: shape[-2] is the block size. + mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16) + # Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size. + eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16) + kv_caches = {"mla_layer": mla_cache, "eagle_layer": eagle_cache} + + with patch.object( + worker.engine, "batch_register_memory", return_value=0 + ) as mock_batch_register: + connector.register_kv_caches(kv_caches) + + mock_batch_register.assert_called_once() + registered_ptrs, registered_lens = mock_batch_register.call_args[0] + assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()] + assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes] + assert worker.block_len_per_layer == [ + mla_cache.nbytes // mla_cache.shape[0], + eagle_cache.nbytes // eagle_cache.shape[0], + ] + + @pytest.mark.asyncio @patch( "vllm.distributed.kv_transfer.kv_connector.v1.mooncake." diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 671a80137b6..855c3411713 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -21,7 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( MultiKVConnectorStats, MultiKVConnectorWorkerMetadata, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( NixlKVConnectorStats, ) from vllm.v1.kv_cache_interface import KVCacheConfig diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index f39b78dd2b7..d67b14e8dd4 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -24,13 +24,13 @@ from vllm.distributed.kv_transfer.kv_connector.utils import ( TpKVTopology, get_current_attn_backend, ) -from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector +from vllm.distributed.kv_transfer.kv_connector.v1 import nixl +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( MultiKVConnectorStats, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( - KVConnectorRole, +from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( NixlAgentMetadata, NixlConnector, NixlConnectorMetadata, @@ -38,6 +38,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( NixlConnectorWorker, NixlHandshakePayload, NixlKVConnectorStats, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( compute_nixl_compatibility_hash, ) from vllm.distributed.kv_transfer.kv_transfer_state import ( @@ -320,7 +322,7 @@ def test_prompt_less_than_block_size(): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_kv_transfer_handshake(dist_init): @@ -523,6 +525,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): kv_cache_layout="HND", block_size=self.block_size, ssm_sizes=(0, 0), + attn_backend_name=self.backend_name, ), remote_tp_rank=remote_tp_rank, remote_tp_size=remote_tp_size, @@ -533,7 +536,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): class TestNixlHandshake: @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_multi_xfer_one_engine( @@ -620,7 +623,7 @@ class TestNixlHandshake: connector.clear_connector_metadata() @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) @pytest.mark.parametrize( @@ -690,7 +693,7 @@ class TestNixlHandshake: raise TimeoutError("Took too long to complete async handshake.") @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) @pytest.mark.parametrize("local_tp_size", [1, 2]) @@ -702,7 +705,7 @@ class TestNixlHandshake: remote configurations. """ monkeypatch.setattr( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", lambda: local_tp_size, ) @@ -759,7 +762,7 @@ class TestNixlHandshake: check_handshake(6) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_prefill_tp_size_greater_than_decode_tp_size_mla( @@ -862,7 +865,7 @@ class TestNixlHandshake: assert req_id not in conn_p1.connector_worker._reqs_to_process @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_concurrent_load_kv( @@ -927,7 +930,7 @@ class TestNixlHandshake: raise TimeoutError("Took too long to complete async handshake.") @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_fails_on_kv_cache_layout_mismatch( @@ -942,7 +945,7 @@ class TestNixlHandshake: # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 with patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", # noqa: E501 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501 return_value=2, ): # Initialize connector and worker (with fake NIXL wrapper) @@ -972,6 +975,7 @@ class TestNixlHandshake: kv_cache_layout=mismatched_layout, block_size=worker.block_size, ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, ) with pytest.raises(RuntimeError): @@ -980,7 +984,7 @@ class TestNixlHandshake: worker.add_remote_agent(meta, remote_tp_size=1) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( @@ -995,7 +999,7 @@ class TestNixlHandshake: # Mock TP world size to 2 to force heterogeneous TP when # remote_tp_size=1 with patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.get_tensor_model_parallel_world_size", # noqa: E501 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.get_tensor_model_parallel_world_size", # noqa: E501 return_value=2, ): # Initialize connector and worker (with fake NIXL wrapper) @@ -1028,6 +1032,7 @@ class TestNixlHandshake: kv_cache_layout="HND", block_size=worker.block_size, ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, ) # We don't check layout for homogeneous TP and MLA for now, as the @@ -1039,7 +1044,7 @@ class TestNixlHandshake: # we put here is important. First run ray, it will clean up the resources, then # the rest of the tests. @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_kv_connector_stats(default_vllm_config, dist_init): @@ -1224,8 +1229,8 @@ def test_multi_kv_connector_stats_aggregation(): worker_patterns = [(2, 1), (3, 0), (0, 5)] # (Nixl, Foo) worker_outputs: list[ModelRunnerOutput] = [] - for i, (nixl, foo) in enumerate(worker_patterns): - stats = make_multi_stats(nixl, foo) + for i, (nixl_count, foo) in enumerate(worker_patterns): + stats = make_multi_stats(nixl_count, foo) output = ModelRunnerOutput( req_ids=[f"req_{i}"], req_id_to_index={f"req_{i}": 0}, @@ -1253,7 +1258,7 @@ def test_multi_kv_connector_stats_aggregation(): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_scheduler_kv_connector_stats_aggregation(): @@ -1321,7 +1326,7 @@ def test_scheduler_kv_connector_stats_aggregation(): @pytest.mark.parametrize("distributed_executor_backend", ["ray", None]) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend): @@ -1510,13 +1515,14 @@ def test_register_kv_caches( backend_cls = TritonAttentionBackend - nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector" + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker" + nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector" with ( - patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper, - patch(f"{nixl_module}.threading.Event"), - patch(f"{nixl_module}.threading.Thread") as mock_thread, - patch(f"{nixl_module}.get_current_attn_backend") as mock_get_attn_backend, - patch(f"{nixl_module}.get_current_attn_backends") as mock_get_attn_backends, + patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper, + patch(f"{nixl_worker}.threading.Event"), + patch(f"{nixl_worker}.threading.Thread") as mock_thread, + patch(f"{nixl_connector}.get_current_attn_backend") as mock_get_attn_backend, + patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends, ): # Ensure get_attn_backend returns the correct value due to # _cached_get_attn_backend returning the backend from previous @@ -1751,28 +1757,26 @@ def test_kv_buffer_to_nixl_memory_types( vllm_config = create_vllm_config() # Override the default memory types in the config vllm_config.kv_transfer_config.kv_buffer_device = kv_buffer_device - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( _NIXL_SUPPORTED_DEVICE, ) _NIXL_SUPPORTED_DEVICE.update(FakePlatform.get_nixl_supported_devices()) with ( + patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper"), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Event" ), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Event" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.threading.Thread" ), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Thread" - ), - patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.current_platform", FakePlatform, ), patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector._NIXL_SUPPORTED_DEVICE", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils._NIXL_SUPPORTED_DEVICE", _NIXL_SUPPORTED_DEVICE, ), ): # noqa: E501 @@ -1787,7 +1791,7 @@ def test_kv_buffer_to_nixl_memory_types( @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): @@ -1852,7 +1856,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init): @@ -1972,7 +1976,7 @@ class FailingNixlWrapper(FakeNixlWrapper): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FailingNixlWrapper, ) @pytest.mark.parametrize( @@ -2062,10 +2066,10 @@ def test_transfer_failure_logging( slot_mapping={}, ) - # Capture logs from the nixl_connector logger specifically + # Capture logs from the nixl.worker logger specifically # vLLM loggers have propagate=False, so we need to capture directly nixl_logger = logging.getLogger( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector" + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker" ) captured_logs: list[logging.LogRecord] = [] @@ -2127,7 +2131,7 @@ def test_transfer_failure_logging( @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FailingNixlWrapper, ) def test_handshake_failure_returns_finished(default_vllm_config, dist_init): @@ -2178,7 +2182,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init): @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FailingNixlWrapper, ) def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init): @@ -2254,7 +2258,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init) ], ) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_compatibility_hash_validation( @@ -2325,7 +2329,7 @@ def test_compatibility_hash_validation( elif "connector_version" in version_override: stack.enter_context( patch.object( - nixl_connector, + nixl.metadata, "NIXL_CONNECTOR_VERSION", version_override["connector_version"], ) @@ -2347,6 +2351,7 @@ def test_compatibility_hash_validation( kv_cache_layout="HND", block_size=prefill_block_size, ssm_sizes=(0, 0), + attn_backend_name=decode_worker.backend_name, ) handshake_payload = NixlHandshakePayload( compatibility_hash=remote_hash, @@ -2361,7 +2366,7 @@ def test_compatibility_hash_validation( # Patch zmq_ctx to return our mock socket with ( patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"), - patch.object(nixl_connector, "zmq_ctx") as mock_zmq_ctx, + patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx, ): mock_zmq_ctx.return_value.__enter__.return_value = mock_socket @@ -2395,7 +2400,7 @@ def test_compatibility_hash_validation( ], ) @patch( - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper", FakeNixlWrapper, ) def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario): @@ -2460,7 +2465,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) mock_socket.recv.return_value = msg_bytes with ( patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"), - patch.object(nixl_connector, "zmq_ctx") as mock_zmq_ctx, + patch.object(nixl.worker, "zmq_ctx") as mock_zmq_ctx, ): mock_zmq_ctx.return_value.__enter__.return_value = mock_socket diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index adb0acae1cb..30913ff98ee 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -31,10 +31,10 @@ from .utils import ( (False, [0]), ], ) -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform") +@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform") def test_sw_sizes(mock_platform, swa_enabled, expected_sw_sizes): """Test sw_sizes is correctly computed based on SWA enabled/disabled.""" - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) @@ -65,7 +65,7 @@ def test_logical_to_kernel_block_ids_with_hma(): When HMA is enabled, the logical block size may differ from the kernel block size. Each logical block maps to multiple kernel blocks. """ - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -89,6 +89,99 @@ def test_logical_to_kernel_block_ids_with_hma(): ) +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "has_mamba,swa_enabled,mamba_enabled,remote_ratio," + "remote_block_ids,expected_remote_block_ids", + [ + # Non-mamba (FA+SWA): both groups expanded via _logical_to_kernel_block_ids. + # Regression for https://github.com/vllm-project/vllm/pull/39724 + ( + False, + True, + False, + 1, + ([0, 1, 2], [3, 4]), + [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]], + ), + # Mamba (FA+Mamba): FA expanded via _logical_to_remote_kernel_block_ids, + # Mamba passed through unchanged. + # remote_ratio=261 (Nemotron 30B TP=1) != local_ratio=2 so that using + # the wrong conversion method produces different FA results. + ( + True, + False, + True, + 261, + ([0, 1, 2], [10, 11]), + [[0, 1, 261, 262, 522, 523], [10, 11]], + ), + ], + ids=["non_mamba_fa_swa", "mamba_fa_ssm"], +) +def test_read_blocks_for_req_expands_remote_ids( + has_mamba, + swa_enabled, + mamba_enabled, + remote_ratio, + remote_block_ids, + expected_remote_block_ids, +): + """_read_blocks_for_req must expand remote logical block IDs to kernel + block IDs when kernel block size != logical block size. + + Non-mamba path uses _logical_to_kernel_block_ids (all groups expanded). + Mamba path uses _logical_to_remote_kernel_block_ids (FA expanded, Mamba + passed through). + """ + from unittest.mock import MagicMock + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + + worker = object.__new__(NixlConnectorWorker) + worker._has_mamba = has_mamba + worker._physical_blocks_per_logical_kv_block = 2 + worker.kv_cache_config = make_kv_cache_config( + block_size=16, swa_enabled=swa_enabled, mamba_enabled=mamba_enabled + ) + + remote_engine_id = "remote-engine" + if has_mamba: + worker._mamba_phys_ratio = {remote_engine_id: remote_ratio} + + # Mock kv_topo: empty remote ranks skips the transfer machinery entirely, + # isolating the block-ID expansion logic. + worker.kv_topo = MagicMock() + worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = [] + worker.kv_topo.tp_ratio_from_engine_id.return_value = 1 + + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="test-req", + local_block_ids=([0, 1], [2, 3]), + kv_transfer_params={ + "remote_block_ids": remote_block_ids, + "remote_engine_id": remote_engine_id, + "remote_request_id": "prefill-test-req", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 1, + }, + ) + + meta = metadata.reqs_to_recv["test-req"] + worker._read_blocks_for_req("test-req", meta) + + assert meta.remote.block_ids == expected_remote_block_ids, ( + f"Expected {expected_remote_block_ids}, got {meta.remote.block_ids}" + ) + + @pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)]) def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): """Test that a prefill instance returns fewer "remote blocks" for the SWA groups @@ -102,7 +195,7 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size): llm_kwargs = { "model": model_name, "enforce_eager": True, - "gpu_memory_utilization": 0.5, + "gpu_memory_utilization": 0.47, "kv_transfer_config": kv_transfer_config, "max_model_len": 2048, # NOTE: Make sure HMA is enabled @@ -169,7 +262,7 @@ def test_nixl_metadata_hma_block_ids_structure(): Test that NixlConnectorMetadata correctly stores block IDs for multiple KV cache groups when HMA is enabled. """ - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) @@ -211,7 +304,7 @@ def test_nixl_metadata_hma_block_ids_structure(): def test_get_block_descs_ids_hybrid_ssm(): """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM when ratio=1 (no kernel block size mismatch).""" - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -247,7 +340,7 @@ def test_get_block_descs_ids_hybrid_ssm(): def test_get_block_descs_ids_kernel_block_mismatch(): """Test _get_block_descs_ids uses different strides for FA (kernel blocks) vs SSM (logical blocks) when ratio > 1.""" - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( NixlConnectorWorker, ) @@ -284,7 +377,7 @@ def test_get_block_descs_ids_kernel_block_mismatch(): def test_nixl_metadata_hybrid_ssm_block_ids(): """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM groups with different block counts (kernel mismatch active).""" - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( NixlConnectorMetadata, ) @@ -392,7 +485,7 @@ def test_mamba_n1_p_side_truncation(): ], ids=["fa_swa_mamba", "fa_swa_only", "fa_only"], ) -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform") +@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform") def test_has_mamba_init( mock_platform, swa_enabled, @@ -401,7 +494,7 @@ def test_has_mamba_init( expected_is_hma, ): """Test _has_mamba / _is_hma_required derived from kv_cache_groups.""" - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index 283b4f25e6e..44fc6d06d77 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -587,7 +587,7 @@ def test_cannot_recv(): assert_scheduler_empty(scheduler) -@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform") +@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler.current_platform") def test_p_side_chunked_prefill_mamba(mock_platform): """P-side integration: Mamba N-1 truncation + chunked prefill completes. diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 75dc479470e..5f0036807b0 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -476,7 +476,7 @@ def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False): Only sets the two flags needed by the N-1 prefill logic. """ - from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( NixlConnectorScheduler, ) diff --git a/tests/v1/kv_offload/test_cpu_gpu.py b/tests/v1/kv_offload/test_cpu_gpu.py index 2da3a5e56b1..de482aec4a4 100644 --- a/tests/v1/kv_offload/test_cpu_gpu.py +++ b/tests/v1/kv_offload/test_cpu_gpu.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random import time +import uuid import pytest import torch from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec from vllm.v1.kv_offload.spec import ( CanonicalKVCacheRef, @@ -36,6 +38,7 @@ NUM_MAPPINGS = [3] @pytest.mark.parametrize("num_tensors", NUM_TENSORS) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("use_shared_memory", [False, True]) @torch.inference_mode() def test_transfer( default_vllm_config, @@ -48,6 +51,7 @@ def test_transfer( num_tensors: int, seed: int, device: str, + use_shared_memory: bool, ) -> None: set_random_seed(seed) @@ -83,10 +87,24 @@ def test_transfer( tensors=kv_cache_tensors, group_data_refs=kv_cache_groups_data_refs, ) + + mmap_region: SharedOffloadRegion | None = None + if use_shared_memory: + cpu_page_size = gpu_page_size_bytes * num_tensors * block_size_factor + mmap_region = SharedOffloadRegion( + instance_id=str(uuid.uuid4()), + total_size_bytes=num_cpu_blocks * cpu_page_size, + num_blocks=num_cpu_blocks, + rank=0, + num_workers=1, + cpu_page_size=cpu_page_size, + ) + handlers = CpuGpuOffloadingHandlers( kv_caches=kv_caches, block_size_factor=block_size_factor, num_cpu_blocks=num_cpu_blocks, + mmap_region=mmap_region, ) # select block mappings @@ -137,10 +155,8 @@ def test_transfer( if finished: assert finished[0].job_id == 1 assert finished[0].success - assert ( - finished[0].transfer_type == ("GPU", "CPU") - if gpu_to_cpu - else ("CPU", "GPU") + assert finished[0].transfer_type == ( + ("GPU", "CPU") if gpu_to_cpu else ("CPU", "GPU") ) assert finished[0].transfer_size == ( len(gpu_blocks) * handler.group_block_size_in_bytes[0] @@ -161,9 +177,9 @@ def test_transfer( orig_dst_tensors, ): # view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison. - src_view = src_tensor.view(-1, gpu_page_size_bytes) - dst_view = dst_tensor.view(-1, gpu_page_size_bytes) - orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes) + src_view = src_tensor.reshape(-1, gpu_page_size_bytes) + dst_view = dst_tensor.reshape(-1, gpu_page_size_bytes) + orig_dst_view = orig_dst_tensor.reshape(-1, gpu_page_size_bytes) for dst_sub_block in range(num_dst_sub_blocks): src_sub_block = dst_to_src.get(dst_sub_block) if src_sub_block is not None: @@ -171,3 +187,12 @@ def test_transfer( else: expected = orig_dst_view[dst_sub_block] torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu()) + + # Drop loop-variable refs so mmap_obj has no exported buffers at cleanup. + del orig_tensor, tensor, src_tensor, dst_tensor, orig_dst_tensor + del src_view, dst_view, orig_dst_view, expected + + handlers.cpu_to_gpu_handler.shutdown() + handlers.gpu_to_cpu_handler.shutdown() + if mmap_region: + mmap_region.cleanup() diff --git a/tests/v1/kv_offload/test_cpu_manager.py b/tests/v1/kv_offload/test_cpu_manager.py index 85629c96cea..7a2ba837eca 100644 --- a/tests/v1/kv_offload/test_cpu_manager.py +++ b/tests/v1/kv_offload/test_cpu_manager.py @@ -59,7 +59,6 @@ def verify_load_output( def verify_events( events: Iterable[OffloadingEvent], - block_size: int, expected_stores: tuple[set[int], ...] = (), expected_evictions: tuple[set[int], ...] = (), ): @@ -67,7 +66,6 @@ def verify_events( evictions: list[set[OffloadKey]] = [] for event in events: assert event.medium == CPULoadStoreSpec.medium() - assert event.block_size == block_size if event.removed: evictions.append(set(event.keys)) else: @@ -98,9 +96,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): candidate to make room for [3, 4, 5] - After complete_store([2, 3, 4, 5]), block 2 must still be present. """ - block_size = 256 manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy=eviction_policy, enable_events=True, @@ -138,10 +134,9 @@ def test_cpu_manager(): """ Tests CPUOffloadingManager with lru policy. """ - # initialize a CPU backend with a capacity of 4 blocks - block_size = 256 + # initialize a CPU manager with a capacity of 4 blocks cpu_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True + num_blocks=4, cache_policy="lru", enable_events=True ) # prepare store [1, 2] @@ -163,9 +158,7 @@ def test_cpu_manager(): # complete store [1, 2] cpu_manager.complete_store(to_keys([1, 2])) - verify_events( - cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},) - ) + verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] assert cpu_manager.lookup(to_keys([1])) == 1 @@ -184,9 +177,7 @@ def test_cpu_manager(): ) # verify eviction event - verify_events( - cpu_manager.take_events(), block_size=block_size, expected_evictions=({1},) - ) + verify_events(cpu_manager.take_events(), expected_evictions=({1},)) # prepare store with no space assert cpu_manager.prepare_store(to_keys([1, 6])) is None @@ -241,7 +232,6 @@ def test_cpu_manager(): verify_events( cpu_manager.take_events(), - block_size=block_size, expected_stores=({3, 4, 5}, {6, 7, 8}), expected_evictions=({2, 3, 4}, {8}), ) @@ -254,7 +244,6 @@ class TestARCPolicy: self, num_blocks: int = 4, enable_events: bool = True ) -> tuple[CPUOffloadingManager, ARCCachePolicy]: manager = CPUOffloadingManager( - block_size=256, num_blocks=num_blocks, cache_policy="arc", enable_events=enable_events, @@ -289,9 +278,7 @@ class TestARCPolicy: # complete store [1, 2] cpu_manager.complete_store(to_keys([1, 2])) - verify_events( - cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},) - ) + verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] assert cpu_manager.lookup(to_keys([1])) == 1 @@ -547,9 +534,8 @@ def test_filter_reused_manager(): """ Tests FilterReusedOffloadingManager with a CPUOffloadingManager. """ - block_size = 256 lru_manager = CPUOffloadingManager( - block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True + num_blocks=4, cache_policy="lru", enable_events=True ) manager = FilterReusedOffloadingManager( diff --git a/tests/v1/kv_offload/test_shared_offload_region.py b/tests/v1/kv_offload/test_shared_offload_region.py new file mode 100644 index 00000000000..b33a27ca645 --- /dev/null +++ b/tests/v1/kv_offload/test_shared_offload_region.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for SharedOffloadRegion.""" + +import contextlib +import mmap +import os +import threading +import time +import uuid + +import pytest + +from vllm.utils.system_utils import get_mp_context +from vllm.v1.kv_offload.cpu.shared_offload_region import ( + SharedOffloadRegion, + _wait_for_file_size, +) + +PAGE_SIZE = mmap.PAGESIZE + + +# --------------------------------------------------------------------------- +# Helpers / fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _set_spawn_method(monkeypatch): + # On WSL, NVML is not compatible with fork so vLLM auto-overrides the + # multiprocessing start method to 'spawn' with a warning. Set it explicitly + # here so the override is a no-op and the warning is suppressed. + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + +def _make_region( + instance_id: str, + num_blocks: int = 4, + cpu_page_size: int = PAGE_SIZE, + num_workers: int = 1, + rank: int = 0, +) -> SharedOffloadRegion: + total_size_bytes = num_blocks * num_workers * cpu_page_size + assert total_size_bytes % PAGE_SIZE == 0 + return SharedOffloadRegion( + instance_id=instance_id, + total_size_bytes=total_size_bytes, + num_blocks=num_blocks, + rank=rank, + num_workers=num_workers, + cpu_page_size=cpu_page_size, + ) + + +def _cleanup_file(path: str) -> None: + """Best-effort file removal for test teardown.""" + with contextlib.suppress(FileNotFoundError): + os.unlink(path) + + +@contextlib.contextmanager +def _region(instance_id: str, **kwargs): + """Context manager: create one region, clean up on exit.""" + r = _make_region(instance_id, **kwargs) + try: + yield r + finally: + r.cleanup() + _cleanup_file(r.mmap_path) + + +@contextlib.contextmanager +def _multi_region( + instance_id: str, + num_workers: int, + num_blocks: int = 4, + cpu_page_size: int = PAGE_SIZE, +): + """Context manager: create one SharedOffloadRegion per rank, clean up on exit.""" + total = num_blocks * num_workers * cpu_page_size + regions = [ + SharedOffloadRegion( + instance_id=instance_id, + total_size_bytes=total, + num_blocks=num_blocks, + rank=rank, + num_workers=num_workers, + cpu_page_size=cpu_page_size, + ) + for rank in range(num_workers) + ] + try: + yield regions + finally: + for r in regions: + r.cleanup() + _cleanup_file(regions[0].mmap_path) + + +def _race_construct( + instance_id: str, + num_workers: int, + num_blocks: int = 4, + cpu_page_size: int = PAGE_SIZE, +) -> tuple[list[SharedOffloadRegion], list[Exception]]: + """Spawn num_workers threads that all race to construct SharedOffloadRegion.""" + total = num_blocks * num_workers * cpu_page_size + regions: list[SharedOffloadRegion | None] = [None] * num_workers + errors: list[Exception] = [] + barrier = threading.Barrier(num_workers) + + def worker(rank: int) -> None: + barrier.wait() # all threads start at the same instant + try: + regions[rank] = SharedOffloadRegion( + instance_id=instance_id, + total_size_bytes=total, + num_blocks=num_blocks, + rank=rank, + num_workers=num_workers, + cpu_page_size=cpu_page_size, + ) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_workers)] + for t in threads: + t.start() + for t in threads: + t.join() + + return [r for r in regions if r is not None], errors + + +def _mp_race_construct_and_write( + instance_id: str, + total_bytes: int, + num_blocks: int, + rank: int, + num_workers: int, + cpu_page_size: int, + fill_value: int, + done_queue, + cleanup_queue, +) -> None: + """Race to construct a SharedOffloadRegion, write fill_value, then wait + for the parent's cleanup signal before tearing down. The wait gives the + parent a window to read the raw mmap before the creator removes the file.""" + try: + region = SharedOffloadRegion( + instance_id=instance_id, + total_size_bytes=total_bytes, + num_blocks=num_blocks, + rank=rank, + num_workers=num_workers, + cpu_page_size=cpu_page_size, + ) + t = region.create_next_view(cpu_page_size) + t[:, :] = fill_value + done_queue.put({"rank": rank, "error": None}) + cleanup_queue.get() # wait for parent's verification to finish + del t # release view before cleanup to avoid BufferError + region.cleanup() + except Exception as e: + done_queue.put({"rank": rank, "error": repr(e)}) + + +@pytest.fixture +def iid(): + """Fresh instance ID for each test.""" + return str(uuid.uuid4()) + + +# --------------------------------------------------------------------------- +# create_next_view — shape, stride and storage offset +# --------------------------------------------------------------------------- + + +def test_create_next_view_shape_and_stride(iid): + """Returned tensor must have shape (num_blocks, tensor_page_size) and + stride (row_stride, 1) where row_stride = cpu_page_size * num_workers.""" + with _region(iid, num_blocks=4, cpu_page_size=2 * PAGE_SIZE) as r: + t = r.create_next_view(PAGE_SIZE) + assert t.shape == (4, PAGE_SIZE) + # num_workers=1 → row_stride = cpu_page_size + assert t.stride() == (2 * PAGE_SIZE, 1) + del t + + +def test_create_next_view_storage_offset_rank0(iid): + """rank=0 worker's first tensor must start at byte 0 of the mmap.""" + with _region(iid, cpu_page_size=PAGE_SIZE, num_workers=2, rank=0) as r: + t = r.create_next_view(PAGE_SIZE) + assert t.data_ptr() == r._base.data_ptr() # storage_offset == 0 + del t + + +def test_create_next_view_storage_offset_rank1(iid): + """rank=1 worker's first tensor must start cpu_page_size bytes into the mmap.""" + with _multi_region(iid, num_workers=2, num_blocks=4) as (r0, r1): + t1 = r1.create_next_view(PAGE_SIZE) + assert t1.data_ptr() == r1._base.data_ptr() + PAGE_SIZE + del t1 + + +def test_create_next_view_row_stride_with_multiple_workers(iid): + """With num_workers=4, row_stride must be 4 * cpu_page_size.""" + with _region(iid, num_blocks=2, num_workers=4) as r: + t = r.create_next_view(PAGE_SIZE) + assert t.stride(0) == 4 * PAGE_SIZE + del t + + +# --------------------------------------------------------------------------- +# create_next_view — cursor advancement +# --------------------------------------------------------------------------- + + +def test_create_next_view_cursor_advances(iid): + """Each call to create_next_view must advance _worker_offset by tensor_page_size.""" + with _region(iid, cpu_page_size=3 * PAGE_SIZE) as r: + assert r._worker_offset == 0 + r.create_next_view(PAGE_SIZE) + assert r._worker_offset == PAGE_SIZE + r.create_next_view(PAGE_SIZE) + assert r._worker_offset == 2 * PAGE_SIZE + r.create_next_view(PAGE_SIZE) + assert r._worker_offset == 3 * PAGE_SIZE # exactly at area end + + +def test_create_next_view_exact_fill_succeeds(iid): + """Allocations whose total exactly equals cpu_page_size must all succeed.""" + with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r: + r.create_next_view(PAGE_SIZE) # first half + r.create_next_view(PAGE_SIZE) # fills to area end — must not raise + + +# --------------------------------------------------------------------------- +# create_next_view — overflow guard +# --------------------------------------------------------------------------- + + +def test_create_next_view_single_overflow_raises(iid): + """A single allocation larger than cpu_page_size must raise AssertionError.""" + with ( + _region(iid) as r, + pytest.raises(AssertionError, match="exceeds worker area end"), + ): + r.create_next_view(PAGE_SIZE + 1) + + +def test_create_next_view_cumulative_overflow_raises(iid): + """Successive allocations that cumulatively exceed cpu_page_size must raise.""" + with _region(iid, cpu_page_size=2 * PAGE_SIZE) as r: + r.create_next_view(PAGE_SIZE) # ok — half used + r.create_next_view(PAGE_SIZE) # ok — full + with pytest.raises(AssertionError, match="exceeds worker area end"): + r.create_next_view(1) # one byte too many + + +def test_create_next_view_overflow_does_not_mutate_cursor(iid): + """A failed create_next_view must leave _worker_offset unchanged.""" + with _region(iid) as r: + offset_before = r._worker_offset + with pytest.raises(AssertionError): + r.create_next_view(PAGE_SIZE + 1) + assert r._worker_offset == offset_before + + +# --------------------------------------------------------------------------- +# create_next_view — data correctness and layout +# --------------------------------------------------------------------------- + + +def test_create_next_view_write_visible_in_raw_mmap(iid): + """Writes into a create_next_view view must appear at the correct raw mmap offset""" + with _region(iid, num_blocks=4) as r: + t = r.create_next_view(PAGE_SIZE) + t[2, :] = 42 # write to block row 2 + + raw = memoryview(r.mmap_obj) + # num_workers=1 → row_stride = PAGE_SIZE; block 2 starts at byte 2*PAGE_SIZE + chunk = bytes(raw[2 * PAGE_SIZE : 3 * PAGE_SIZE]) + assert all(b == 42 for b in chunk) + del raw, t + + +def test_create_next_view_multi_tensor_layout(iid): + """Two tensors from the same worker land at consecutive byte offsets per row.""" + with _region(iid, num_blocks=2, cpu_page_size=2 * PAGE_SIZE) as r: + ta = r.create_next_view(PAGE_SIZE) + tb = r.create_next_view(PAGE_SIZE) + + ta[:, :] = 1 + tb[:, :] = 2 + + raw = memoryview(r.mmap_obj) + for blk in range(2): + row_offset = blk * 2 * PAGE_SIZE # num_workers=1 + assert all(b == 1 for b in raw[row_offset : row_offset + PAGE_SIZE]) + assert all( + b == 2 for b in raw[row_offset + PAGE_SIZE : row_offset + 2 * PAGE_SIZE] + ) + del raw, ta, tb + + +def test_create_next_view_multiprocess_slots(iid): + """Each worker process calls create_next_view and writes distinct data; + the parent verifies each slot lands at the correct interleaved offset.""" + num_workers = 2 + num_blocks = 4 + total_bytes = num_blocks * num_workers * PAGE_SIZE + + ctx = get_mp_context() + done_queue = ctx.Queue() + cleanup_queue = ctx.Queue() + + # Parent is rank 0 (creator); child is rank 1 (joiner). + region = SharedOffloadRegion( + instance_id=iid, + total_size_bytes=total_bytes, + num_blocks=num_blocks, + rank=0, + num_workers=num_workers, + cpu_page_size=PAGE_SIZE, + ) + try: + child = ctx.Process( + target=_mp_race_construct_and_write, + args=( + iid, + total_bytes, + num_blocks, + 1, + num_workers, + PAGE_SIZE, + 22, + done_queue, + cleanup_queue, + ), + ) + child.start() + + t0 = region.create_next_view(PAGE_SIZE) + t0[:, :] = 11 + + result = done_queue.get(timeout=30) + assert result["error"] is None, result["error"] + + raw = memoryview(region.mmap_obj) + for blk in range(num_blocks): + row_start = blk * num_workers * PAGE_SIZE + w0 = bytes(raw[row_start : row_start + PAGE_SIZE]) + w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE]) + assert all(b == 11 for b in w0), f"block {blk}: rank0 slot wrong" + assert all(b == 22 for b in w1), f"block {blk}: rank1 slot wrong" + + del raw, t0 # release before finally triggers cleanup + cleanup_queue.put(True) + child.join(timeout=10) + assert child.exitcode == 0 + finally: + region.cleanup() + _cleanup_file(region.mmap_path) + + +def test_create_next_view_worker_isolation(iid): + """Writes by worker 0 must not affect worker 1's slot and vice versa.""" + num_workers = 2 + num_blocks = 4 + with _multi_region(iid, num_workers=num_workers, num_blocks=num_blocks) as regions: + t0 = regions[0].create_next_view(PAGE_SIZE) + t1 = regions[1].create_next_view(PAGE_SIZE) + + t0[:, :] = 11 + t1[:, :] = 22 + + raw = memoryview(regions[0].mmap_obj) + for blk in range(num_blocks): + row_start = blk * num_workers * PAGE_SIZE + w0 = bytes(raw[row_start : row_start + PAGE_SIZE]) + w1 = bytes(raw[row_start + PAGE_SIZE : row_start + 2 * PAGE_SIZE]) + assert all(b == 11 for b in w0), f"block {blk}: worker0 slot corrupted" + assert all(b == 22 for b in w1), f"block {blk}: worker1 slot corrupted" + del raw, t0, t1 # release before finally triggers cleanup + + +# --------------------------------------------------------------------------- +# Constructor — creator vs joiner semantics +# --------------------------------------------------------------------------- + + +def test_creator_flag_set_on_first_open(iid): + """The first worker to open the file must have _creator == True.""" + with _region(iid) as r: + assert r._creator is True + + +def test_joiner_flag_not_set(iid): + """A second worker opening the same file must have _creator == False.""" + with _multi_region(iid, num_workers=2) as (r0, r1): + assert r0._creator is True + assert r1._creator is False + + +def test_file_exists_after_construction(iid): + """The mmap file must be present on disk after __init__ completes.""" + with _region(iid) as r: + assert os.path.exists(r.mmap_path) + + +def test_file_has_correct_size(iid): + """The mmap file size on disk must equal total_size_bytes.""" + with _region(iid, num_blocks=4) as r: + assert os.path.getsize(r.mmap_path) == 4 * PAGE_SIZE + + +# --------------------------------------------------------------------------- +# Multi-worker race — concurrent construction +# --------------------------------------------------------------------------- + + +def test_multi_worker_race_exactly_one_creator(iid): + """When N threads race to create the same region, exactly one becomes creator.""" + num_workers = 8 + regions, errors = _race_construct(iid, num_workers=num_workers) + try: + assert not errors, f"Workers raised: {errors}" + assert len(regions) == num_workers, "Some workers failed to construct" + + creators = [r for r in regions if r._creator] + assert len(creators) == 1, f"Expected 1 creator, got {len(creators)}" + assert sum(1 for r in regions if not r._creator) == num_workers - 1, ( + f"Expected {num_workers - 1} non-creators, got " + f"{sum(1 for r in regions if not r._creator)}" + ) + + for r in regions: + assert not r.mmap_obj.closed + assert r.total_size_bytes == 4 * num_workers * PAGE_SIZE + finally: + for r in regions: + r.cleanup() + _cleanup_file(regions[0].mmap_path) + + +def test_multi_worker_race_shared_memory_visible(iid): + """After a concurrent construction race, MAP_SHARED is intact across all workers.""" + num_workers = 4 + regions, errors = _race_construct(iid, num_workers=num_workers) + assert not errors + try: + regions[0].mmap_obj[0:1] = b"\xab" + for r in regions[1:]: + assert memoryview(r.mmap_obj)[0:1] == b"\xab" + finally: + for r in regions: + r.cleanup() + _cleanup_file(regions[0].mmap_path) + + +def test_multiprocess_race_construct_and_write(iid): + """N processes race to construct the same SharedOffloadRegion, each writes + fill_value = rank+1 into their slot; parent verifies interleaved layout.""" + num_workers = 4 + num_blocks = 3 + total_bytes = num_blocks * num_workers * PAGE_SIZE + + ctx = get_mp_context() + done_queue = ctx.Queue() + cleanup_queue = ctx.Queue() + + procs = [ + ctx.Process( + target=_mp_race_construct_and_write, + args=( + iid, + total_bytes, + num_blocks, + rank, + num_workers, + PAGE_SIZE, + rank + 1, + done_queue, + cleanup_queue, + ), + ) + for rank in range(num_workers) + ] + for p in procs: + p.start() + + results = {} + for _ in range(num_workers): + r = done_queue.get(timeout=30) + results[r["rank"]] = r + + for rank, r in results.items(): + assert r["error"] is None, f"rank {rank}: {r['error']}" + + # Read the raw file while all workers still hold it open. + mmap_path = f"/dev/shm/vllm_offload_{iid}.mmap" + with open(mmap_path, "rb") as f: + raw = f.read() + + for blk in range(num_blocks): + for w in range(num_workers): + slot_start = (blk * num_workers + w) * PAGE_SIZE + slot = raw[slot_start : slot_start + PAGE_SIZE] + expected = w + 1 # fill_value = rank + 1 + assert all(b == expected for b in slot), ( + f"block {blk}, worker {w}: expected {expected} but got wrong bytes" + ) + + # Unblock all workers to clean up. + for _ in range(num_workers): + cleanup_queue.put(True) + for p in procs: + p.join(timeout=10) + assert p.exitcode == 0 + + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + + +def test_cleanup_creator_all_effects(iid): + """cleanup() on the creator closes mmap, closes fd, and removes the file.""" + r = _make_region(iid) + path = r.mmap_path + fd = r.fd + mmap_obj = r.mmap_obj + + r.cleanup() + + assert mmap_obj.closed, "mmap should be closed after cleanup" + assert not os.path.exists(path), "creator should remove the file" + with pytest.raises(OSError): + os.fstat(fd) # fd should be closed + + +def test_cleanup_non_creator_all_effects(iid): + """cleanup() on a non-creator closes mmap and fd, but leaves the file on disk.""" + r0 = _make_region(iid) # creator + r1 = _make_region(iid) # joiner + path = r0.mmap_path + fd1 = r1.fd + mmap_obj1 = r1.mmap_obj + try: + r1.cleanup() + + assert mmap_obj1.closed, "mmap should be closed after cleanup" + assert os.path.exists(path), "non-creator must not remove the file" + with pytest.raises(OSError): + os.fstat(fd1) # fd should be closed + finally: + r0.cleanup() + _cleanup_file(path) + + +def test_cleanup_idempotent(iid): + """Calling cleanup() twice must not raise any exception.""" + r = _make_region(iid) + r.cleanup() + r.cleanup() # must be a no-op + + +def test_cleanup_after_create_next_view_releases_mmap(iid): + """cleanup() must close the mmap even after create_next_view was called. + create_next_view returns a view that shares storage with _base; both must be + released before mmap.close() can succeed.""" + r = _make_region(iid) + mmap_obj = r.mmap_obj + + t = r.create_next_view(PAGE_SIZE) + del t + + r.cleanup() + + assert mmap_obj.closed, "mmap should be closed after releasing the tensor" + + +# --------------------------------------------------------------------------- +# _wait_for_file_size +# --------------------------------------------------------------------------- + + +def test_wait_for_file_size_already_large_enough(tmp_path): + """_wait_for_file_size must return immediately when file is already big enough.""" + fd = os.open(str(tmp_path / "ready.mmap"), os.O_CREAT | os.O_RDWR, 0o600) + try: + os.ftruncate(fd, PAGE_SIZE) + start = time.monotonic() + _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0) + assert time.monotonic() - start < 0.5 + finally: + os.close(fd) + + +def test_wait_for_file_size_waits_for_grow(tmp_path): + """_wait_for_file_size must return once a background thread grows the file.""" + fd = os.open(str(tmp_path / "grow.mmap"), os.O_CREAT | os.O_RDWR, 0o600) + try: + + def grow(): + time.sleep(0.05) + os.ftruncate(fd, PAGE_SIZE) + + t = threading.Thread(target=grow) + t.start() + _wait_for_file_size(fd, PAGE_SIZE, timeout=5.0) # must not raise + t.join() + finally: + os.close(fd) + + +def test_wait_for_file_size_timeout(tmp_path): + """_wait_for_file_size must raise TimeoutError when the file never grows.""" + fd = os.open(str(tmp_path / "stuck.mmap"), os.O_CREAT | os.O_RDWR, 0o600) + try: + with pytest.raises(TimeoutError): + _wait_for_file_size(fd, PAGE_SIZE, timeout=0.1) + finally: + os.close(fd) diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index bf29793710a..9ee6a70abe4 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -106,6 +106,7 @@ class MockReasoningConfig: reasoning_start_token_ids = [THINK_START_TOKEN_ID] reasoning_end_token_ids = [THINK_END_TOKEN_ID] + enabled = True def _generate_fake_sampling_metadata( diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index d49874adc99..21f496ea4ae 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -1,7 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm.v1.engine import FinishReason -from vllm.v1.metrics.stats import IterationStats, PromptTokenStats, RequestStateStats +from vllm.v1.metrics.stats import ( + IterationStats, + PrefillStats, + PromptTokenStats, + RequestStateStats, +) def test_iteration_stats_repr(): @@ -21,6 +26,7 @@ def test_prefill_kv_computed_with_cache(): # Case 1: With prefix cache (1200 tokens cached) iteration_stats.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-001", num_prompt_tokens=10000, max_tokens_param=100, req_stats=req_stats, @@ -30,6 +36,7 @@ def test_prefill_kv_computed_with_cache(): finished_req = iteration_stats.finished_requests[0] assert finished_req.num_prompt_tokens == 10000 assert finished_req.num_cached_tokens == 1200 + assert finished_req.request_id == "test-req-001" # Verify calculation: prefill KV = prompt tokens - cached tokens prefill_kv_computed = finished_req.num_prompt_tokens - max( @@ -50,6 +57,7 @@ def test_prefill_kv_computed_no_cache(): # Case 2: No prefix cache iteration_stats.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-002", num_prompt_tokens=2000, max_tokens_param=100, req_stats=req_stats, @@ -59,6 +67,7 @@ def test_prefill_kv_computed_no_cache(): finished_req = iteration_stats.finished_requests[0] assert finished_req.num_prompt_tokens == 2000 assert finished_req.num_cached_tokens == 0 + assert finished_req.request_id == "test-req-002" # Verify calculation: prefill KV = full prompt when no cache prefill_kv_computed = finished_req.num_prompt_tokens - max( @@ -79,6 +88,7 @@ def test_prefill_kv_computed_edge_cases(): # Case 3: Negative num_cached_tokens (shouldn't happen, but handle gracefully) iteration_stats.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-003", num_prompt_tokens=100, max_tokens_param=10, req_stats=req_stats, @@ -91,11 +101,13 @@ def test_prefill_kv_computed_edge_cases(): finished_req.num_cached_tokens, 0 ) assert prefill_kv_computed == 100 # Should treat negative as 0 + assert finished_req.request_id == "test-req-003" # Case 4: All tokens cached (shouldn't happen in practice) iteration_stats2 = IterationStats() iteration_stats2.update_from_finished_request( finish_reason=FinishReason.STOP, + request_id="test-req-004", num_prompt_tokens=100, max_tokens_param=10, req_stats=req_stats, @@ -107,6 +119,7 @@ def test_prefill_kv_computed_edge_cases(): finished_req2.num_cached_tokens, 0 ) assert prefill_kv_computed2 == 0 # All cached, nothing computed + assert finished_req2.request_id == "test-req-004" def test_prompt_token_stats_all_computed(): @@ -114,15 +127,18 @@ def test_prompt_token_stats_all_computed(): stats = PromptTokenStats() # Case 1: No caching (All tokens computed locally) - stats.update_from_output( - num_cached_tokens=0, - num_external_computed_tokens=0, - prompt_len=1000, + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=1000, + num_local_cached_tokens=0, + num_external_cached_tokens=0, ) + stats.update_from_output(prefill_stats) assert stats.computed == 1000 assert stats.local_cache_hit == 0 assert stats.external_kv_transfer == 0 + assert stats.cached_tokens == 0 assert stats.total == 1000 @@ -131,15 +147,19 @@ def test_prompt_token_stats_partial_local_cache(): stats = PromptTokenStats() # Case 2: Partial local cache - stats.update_from_output( - num_cached_tokens=300, - num_external_computed_tokens=0, - prompt_len=1000, + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=1000, + num_local_cached_tokens=300, + num_external_cached_tokens=0, ) + stats.update_from_output(prefill_stats) assert stats.computed == 700 assert stats.local_cache_hit == 300 assert stats.external_kv_transfer == 0 + assert stats.cached_tokens == 300 + assert stats.total == 1000 def test_prompt_token_stats_partial_external_transfer(): @@ -147,15 +167,19 @@ def test_prompt_token_stats_partial_external_transfer(): stats = PromptTokenStats() # Case 3: Partial external transfer - stats.update_from_output( - num_cached_tokens=500, - num_external_computed_tokens=500, - prompt_len=1000, + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=1000, + num_local_cached_tokens=0, + num_external_cached_tokens=500, ) + stats.update_from_output(prefill_stats) assert stats.computed == 500 assert stats.local_cache_hit == 0 assert stats.external_kv_transfer == 500 + assert stats.cached_tokens == 500 + assert stats.total == 1000 def test_prompt_token_stats_mixed_sources(): @@ -163,49 +187,60 @@ def test_prompt_token_stats_mixed_sources(): stats = PromptTokenStats() # Case 4: Mixed sources - stats.update_from_output( - num_cached_tokens=600, - num_external_computed_tokens=200, - prompt_len=1000, + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=1000, + num_local_cached_tokens=400, + num_external_cached_tokens=200, ) + stats.update_from_output(prefill_stats) assert stats.computed == 400 assert stats.local_cache_hit == 400 assert stats.external_kv_transfer == 200 + assert stats.cached_tokens == 600 + assert stats.total == 1000 def test_prompt_token_stats_full_local_cache_recompute(): """Test full local cache triggers last token recomputation. - When all tokens are cached, the scheduler reduces num_cached_tokens by 1 - to force the model to recompute the last token. + When all tokens are cached, the scheduler forces the model to recompute + the last token (num_computed_tokens=1), with the rest from cache. """ stats = PromptTokenStats() - # Case 5: Full local cache (999 cached after reduction, 1 recomputed) - stats.update_from_output( - num_cached_tokens=999, - num_external_computed_tokens=0, - prompt_len=1000, + # Case 5: Full local cache (999 cached, 1 recomputed) + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=1000, + num_local_cached_tokens=999, + num_external_cached_tokens=0, ) + stats.update_from_output(prefill_stats) assert stats.computed == 1 - assert stats.local_cache_hit == 1000 - assert stats.recomputed_tokens == 1 + assert stats.local_cache_hit == 999 + assert stats.external_kv_transfer == 0 + assert stats.cached_tokens == 999 + assert stats.total == 1000 def test_prompt_token_stats_full_external_transfer_recompute(): """Test full external transfer triggers last token recomputation.""" stats = PromptTokenStats() - # Case 6: Full external transfer (999 cached after reduction, 1 recomputed) - stats.update_from_output( - num_cached_tokens=999, - num_external_computed_tokens=1000, - prompt_len=1000, + # Case 6: Full external transfer (999 from external, 1 recomputed) + prefill_stats = PrefillStats() + prefill_stats.set( + num_prompt_tokens=1000, + num_local_cached_tokens=0, + num_external_cached_tokens=999, ) + stats.update_from_output(prefill_stats) assert stats.computed == 1 assert stats.local_cache_hit == 0 - assert stats.external_kv_transfer == 1000 - assert stats.recomputed_tokens == 1 + assert stats.external_kv_transfer == 999 + assert stats.cached_tokens == 999 + assert stats.total == 1000 diff --git a/tests/v1/sample/test_batched_count_greater_than.py b/tests/v1/sample/test_batched_count_greater_than.py new file mode 100644 index 00000000000..c9ace93c696 --- /dev/null +++ b/tests/v1/sample/test_batched_count_greater_than.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test that batched_count_greater_than does not trigger 0/1 specialization +recompiles when batch_size varies.""" + +import torch + +from vllm.platforms import current_platform +from vllm.v1.sample.ops.logprobs import batched_count_greater_than +from vllm.v1.sample.sampler import Sampler + +DEVICE = current_platform.device_type + + +def test_batched_count_greater_than_correctness(): + """Basic correctness: counts elements >= the corresponding value.""" + x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device=DEVICE) + values = torch.tensor([[2.0], [5.0]], device=DEVICE) + result = batched_count_greater_than(x, values) + expected = torch.tensor([2, 2], device=DEVICE) + torch.testing.assert_close(result, expected) + + +def test_gather_logprobs_no_recompile(): + """Sampler.gather_logprobs with batch_size=1 then 2 must not recompile. + + This guards against 0/1 specialization: dynamo normally specializes on + tensor sizes 0 and 1, causing a recompile when the size first exceeds 1. + The mark_unbacked calls in gather_logprobs prevent this. + """ + torch._dynamo.reset() + + compile_count = 0 + orig_backend = current_platform.simple_compile_backend + + def counting_backend(gm, example_inputs): + nonlocal compile_count + compile_count += 1 + if orig_backend == "inductor": + return torch._inductor.compile(gm, example_inputs) + return gm + + # Monkey-patch batched_count_greater_than with our counting backend + # so we can detect recompiles through the production code path. + import vllm.v1.sample.ops.logprobs as logprobs_module + import vllm.v1.sample.sampler as sampler_module + + unwrapped = batched_count_greater_than._torchdynamo_orig_callable + patched = torch.compile(unwrapped, backend=counting_backend) + orig_fn = logprobs_module.batched_count_greater_than + + logprobs_module.batched_count_greater_than = patched + sampler_module.batched_count_greater_than = patched + + try: + vocab_size = 32 + num_logprobs = 3 + + # Call 1: batch_size=1 + logprobs1 = torch.randn(1, vocab_size, device=DEVICE) + token_ids1 = torch.randint( + 0, vocab_size, (1,), device=DEVICE, dtype=torch.int64 + ) + Sampler.gather_logprobs(logprobs1, num_logprobs, token_ids1) + assert compile_count == 1, f"Expected 1 compile, got {compile_count}" + + # Call 2: batch_size=2 — should NOT recompile + logprobs2 = torch.randn(2, vocab_size, device=DEVICE) + token_ids2 = torch.randint( + 0, vocab_size, (2,), device=DEVICE, dtype=torch.int64 + ) + Sampler.gather_logprobs(logprobs2, num_logprobs, token_ids2) + assert compile_count == 1, ( + f"Recompiled on batch_size 1->2 (0/1 specialization). " + f"Expected 1 compile, got {compile_count}" + ) + + # Call 3: batch_size=8 — should NOT recompile + logprobs3 = torch.randn(8, vocab_size, device=DEVICE) + token_ids3 = torch.randint( + 0, vocab_size, (8,), device=DEVICE, dtype=torch.int64 + ) + Sampler.gather_logprobs(logprobs3, num_logprobs, token_ids3) + assert compile_count == 1, ( + f"Recompiled on batch_size change. Expected 1 compile, got {compile_count}" + ) + finally: + # Restore original function + logprobs_module.batched_count_greater_than = orig_fn + sampler_module.batched_count_greater_than = orig_fn + torch._dynamo.reset() diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 511f2668075..23f1f1c1f98 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -127,7 +127,7 @@ def test_flashinfer_sampler(): # ============================================================================= -@pytest.mark.skipif("CPU" in DEVICE_TYPE, reason="CUDA/XPU not available") +@pytest.mark.skipif("cpu" in DEVICE_TYPE, reason="CUDA/XPU not available") class TestTritonTopkTopp: """Tests for the Triton top-k/top-p kernel.""" diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 5d587fa3ec1..188e84abca0 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -755,12 +755,6 @@ def test_load_model( use_distinct_lm_head, monkeypatch, ): - if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm(): - pytest.skip( - "TRITON_ATTN does not support " - "multi-token eagle spec decode on current platform" - ) - if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") diff --git a/tests/v1/spec_decode/test_eagle_step_kernel.py b/tests/v1/spec_decode/test_eagle_step_kernel.py index 275a157d1be..83a29bc0469 100644 --- a/tests/v1/spec_decode/test_eagle_step_kernel.py +++ b/tests/v1/spec_decode/test_eagle_step_kernel.py @@ -15,8 +15,8 @@ DEVICE_TYPE = current_platform.device_type # Skip if no CUDA - Triton kernel requires GPU pytest.importorskip("triton") -if not torch.cuda.is_available(): - pytest.skip("CUDA required for EAGLE kernel tests", allow_module_level=True) +if not current_platform.is_cuda_alike() and not current_platform.is_xpu(): + pytest.skip("CUDA/XPU required for EAGLE kernel tests", allow_module_level=True) def _reference_eagle_step_slot_mapping( diff --git a/tests/v1/spec_decode/test_max_len.py b/tests/v1/spec_decode/test_max_len.py index 42991f9f1ae..1e1c6745191 100644 --- a/tests/v1/spec_decode/test_max_len.py +++ b/tests/v1/spec_decode/test_max_len.py @@ -38,12 +38,6 @@ def test_ngram_max_len(num_speculative_tokens: int): def test_eagle_max_len( monkeypatch: pytest.MonkeyPatch, num_speculative_tokens: int, attn_backend: str ): - if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm(): - pytest.skip( - "TRITON_ATTN does not support " - "multi-token eagle spec decode on current platform" - ) - if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") diff --git a/tests/v1/spec_decode/test_speculators_dflash.py b/tests/v1/spec_decode/test_speculators_dflash.py new file mode 100644 index 00000000000..2ba580695dd --- /dev/null +++ b/tests/v1/spec_decode/test_speculators_dflash.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline +from tests.utils import large_gpu_mark +from vllm import LLM +from vllm.config import SpeculativeConfig +from vllm.distributed import cleanup_dist_env_and_memory + +MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators" + +EXPECTED_GSM8K_ACCURACY = 0.885 +ACCURACY_RTOL = 0.03 +EXPECTED_ACCEPTANCE_LEN = 3.45 +ACCEPTANCE_LEN_RTOL = 0.15 + +# Expected per-position acceptance rates (accepted_at_pos / num_drafts) +# Based on GSM8K evaluation with Qwen3-8B dflash speculators. +EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282] +PER_POS_RTOL = 0.15 + + +def compute_spec_decode_stats( + metrics, +) -> dict: + """Extract all spec-decode metrics and compute derived stats.""" + name2metric = {m.name: m for m in metrics} + + n_drafts = name2metric["vllm:spec_decode_num_drafts"].value + n_draft_tokens = name2metric["vllm:spec_decode_num_draft_tokens"].value + n_accepted = name2metric["vllm:spec_decode_num_accepted_tokens"].value + + per_pos_vec = name2metric["vllm:spec_decode_num_accepted_tokens_per_pos"].values + + acceptance_len = 1 + (n_accepted / n_drafts) if n_drafts > 0 else 1.0 + draft_tokens_per_step = (n_draft_tokens / n_drafts) if n_drafts > 0 else 0 + overall_acceptance_rate = (n_accepted / n_draft_tokens) if n_draft_tokens > 0 else 0 + per_pos_rates = [v / n_drafts for v in per_pos_vec] if n_drafts > 0 else [] + + return { + "num_drafts": n_drafts, + "num_draft_tokens": n_draft_tokens, + "num_accepted_tokens": n_accepted, + "acceptance_len": acceptance_len, + "draft_tokens_per_step": draft_tokens_per_step, + "overall_acceptance_rate": overall_acceptance_rate, + "per_pos_accepted": list(per_pos_vec), + "per_pos_acceptance_rates": per_pos_rates, + } + + +def print_spec_decode_stats(stats: dict) -> None: + """Print all spec-decode metrics and derived values.""" + print("\n===== Spec Decode Metrics =====") + print(f" num_drafts: {stats['num_drafts']}") + print(f" num_draft_tokens: {stats['num_draft_tokens']}") + print(f" num_accepted_tokens: {stats['num_accepted_tokens']}") + print(f" draft_tokens_per_step: {stats['draft_tokens_per_step']:.2f}") + print(f" overall_acceptance_rate: {stats['overall_acceptance_rate']:.4f}") + print(f" acceptance_len (1+acc/drafts): {stats['acceptance_len']:.4f}") + print(" per-position accepted tokens:", stats["per_pos_accepted"]) + print(" per-position acceptance rates:") + for i, rate in enumerate(stats["per_pos_acceptance_rates"]): + print(f" pos {i}: {rate:.4f}") + print("===============================\n") + + +def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch): + """ + Test DFlash speculators model properly initializes speculative decoding. + + Verifies: + 1. Speculative config is automatically initialized from speculators config + 2. Method is detected as 'dflash' + 3. The draft model path is correctly set + 4. Speculative tokens count is valid (num_speculative_tokens=8) + 5. Text generation works with speculative decoding enabled + """ + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + with vllm_runner( + MODEL_PATH, + dtype=torch.bfloat16, + enforce_eager=True, + quantization="fp8", + ) as vllm_model: + vllm_config = vllm_model.llm.llm_engine.vllm_config + + assert isinstance(vllm_config.speculative_config, SpeculativeConfig), ( + "Speculative config should be initialized for speculators model" + ) + + spec_config = vllm_config.speculative_config + assert spec_config.method == "dflash", ( + f"Expected method='dflash', got '{spec_config.method}'" + ) + assert spec_config.num_speculative_tokens > 0, ( + f"Expected positive speculative tokens, " + f"got {spec_config.num_speculative_tokens}" + ) + assert spec_config.model == MODEL_PATH, ( + f"Draft model should be {MODEL_PATH}, got {spec_config.model}" + ) + + vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20) + assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}" + + +@pytest.mark.slow_test +@large_gpu_mark(min_gb=40) +def test_dflash_speculators_correctness(monkeypatch): + """ + E2E correctness test for DFlash via the speculators auto-detect path. + + Evaluates GSM8k accuracy to ensure the speculators-format model produces + correct outputs, and checks that acceptance length does not collapse under + batched inference (lm-eval style). + + Observed per-position acceptance rates on GSM8K (1319 prompts): + pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282, + pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023 + Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128) + """ + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + spec_llm = LLM( + model=MODEL_PATH, + trust_remote_code=True, + max_model_len=4096, + max_num_seqs=128, + gpu_memory_utilization=0.85, + enforce_eager=False, + disable_log_stats=False, + ) + + results = evaluate_gsm8k_offline(spec_llm) + accuracy = results["accuracy"] + accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL) + assert accuracy >= accuracy_threshold, ( + f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}" + ) + + current_metrics = spec_llm.get_metrics() + stats = compute_spec_decode_stats(current_metrics) + print_spec_decode_stats(stats) + + acceptance_len = stats["acceptance_len"] + al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL) + assert acceptance_len >= al_threshold, ( + f"DFlash speculators acceptance length too low: " + f"{acceptance_len:.2f} < {al_threshold:.2f}" + ) + + # Check per-position acceptance rates for the first few positions. + per_pos_rates = stats["per_pos_acceptance_rates"] + for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES): + assert i < len(per_pos_rates), ( + f"Missing per-position acceptance rate for position {i}" + ) + threshold = expected_rate * (1 - PER_POS_RTOL) + assert per_pos_rates[i] >= threshold, ( + f"Per-position acceptance rate at pos {i} too low: " + f"{per_pos_rates[i]:.4f} < {threshold:.4f} " + f"(expected ~{expected_rate:.4f})" + ) + + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() diff --git a/tests/v1/spec_decode/test_tree_attention.py b/tests/v1/spec_decode/test_tree_attention.py index cb487acec0a..1b6fa4f6f48 100644 --- a/tests/v1/spec_decode/test_tree_attention.py +++ b/tests/v1/spec_decode/test_tree_attention.py @@ -14,6 +14,7 @@ from tests.v1.attention.utils import ( ) from vllm.config import ParallelConfig, SpeculativeConfig from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -323,8 +324,7 @@ def forward_attention( def test_tree_attn_correctness( reference_backend: AttentionBackendEnum, ) -> None: - torch.manual_seed(42) - torch.cuda.manual_seed_all(42) + set_random_seed(42) device = "cuda" tree_attn_masks = { diff --git a/tests/v1/test_tensor_ipc_queue.py b/tests/v1/test_tensor_ipc_queue.py index a3fcb97ca17..a70f5d48cc5 100644 --- a/tests/v1/test_tensor_ipc_queue.py +++ b/tests/v1/test_tensor_ipc_queue.py @@ -14,6 +14,7 @@ import pytest import torch import torch.multiprocessing as torch_mp +from vllm.platforms import current_platform from vllm.v1.engine.tensor_ipc import ( TensorIpcData, TensorIpcReceiver, @@ -21,6 +22,8 @@ from vllm.v1.engine.tensor_ipc import ( ) from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +DEVICE_TYPE = current_platform.device_type + @pytest.fixture(scope="module", autouse=True) def setup_multiprocessing(): @@ -53,7 +56,7 @@ def encoder_process( encoder = MsgpackEncoder(oob_tensor_consumer=sender) if torch.cuda.is_available(): - device = "cuda:0" + device = f"{DEVICE_TYPE}:0" tensor = torch.randn( *tensor_data["shape"], dtype=tensor_data["dtype"], device=device ) @@ -384,7 +387,7 @@ def mixed_tensor_encoder_process( # Create only CUDA tensor for IPC (CPU will be serialized) # But actually, let's just send CUDA tensor directly - cuda_tensor = torch.randn(4, 5, device="cuda:0") + cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0") # Manually send via IPC to test the mechanism cuda_tensor_shared = cuda_tensor.share_memory_() @@ -651,7 +654,7 @@ def test_ipc_disabled_mode(): # If CUDA is available, test with CUDA tensor too if torch.cuda.is_available(): - cuda_tensor = torch.randn(4, 5, device="cuda:0") + cuda_tensor = torch.randn(4, 5, device=f"{DEVICE_TYPE}:0") encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor}) assert len(encoded_cuda) > 0 assert tensor_queues[0].empty(), ( diff --git a/tools/generate_cmake_presets.py b/tools/generate_cmake_presets.py index 85847c2c0fe..6bc8443c447 100644 --- a/tools/generate_cmake_presets.py +++ b/tools/generate_cmake_presets.py @@ -128,7 +128,7 @@ def generate_presets(output_path="CMakeUserPresets.json", force_overwrite=False) presets = { "version": 6, - # Keep in sync with CMakeLists.txt and requirements/build.txt + # Keep in sync with CMakeLists.txt and requirements/build/cuda.txt "cmakeMinimumRequired": {"major": 3, "minor": 26, "patch": 1}, "configurePresets": [configure_preset], "buildPresets": [ diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 0e1adda97b6..9d1edee0472 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -5,6 +5,7 @@ set -e # Default values +# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" DEEPGEMM_GIT_REF="477618cd51baffca09c4b0b87e97c03fe827ef03" WHEEL_DIR="" diff --git a/tools/pre_commit/check_torch_cuda.py b/tools/pre_commit/check_torch_cuda.py index 9bc4ab56dfd..08ee419515d 100644 --- a/tools/pre_commit/check_torch_cuda.py +++ b/tools/pre_commit/check_torch_cuda.py @@ -9,6 +9,7 @@ import regex as re # --------------------------------------------------------------------------- # _TORCH_CUDA_PATTERNS = [ r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|set_device|device\()\b", + r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b", r"\bwith\storch\.cuda\.device\b", # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally r"\bcuda_device_count_stateless\(\)\b", @@ -24,6 +25,14 @@ def scan_file(path: str) -> int: for match in re.finditer(pattern, content, re.MULTILINE): # Calculate line number from match position line_num = content[: match.start() + 1].count("\n") + 1 + matched_text = match.group(0) + if "manual_seed" in matched_text: + print( + f"{path}:{line_num}: " + "\033[91merror:\033[0m " + f"Found {matched_text} API call. Use set_random_seed instead." + ) + return 1 print( f"{path}:{line_num}: " "\033[91merror:\033[0m " # red color diff --git a/tools/pre_commit/generate_nightly_torch_test.py b/tools/pre_commit/generate_nightly_torch_test.py index a3d7f7a609b..354da54df8f 100644 --- a/tools/pre_commit/generate_nightly_torch_test.py +++ b/tools/pre_commit/generate_nightly_torch_test.py @@ -3,15 +3,15 @@ """ Generates specialized requirements files for nightly PyTorch testing. -This script reads the main test requirements input file (`requirements/test.in`) +This script reads the main test requirements input file (`requirements/test/cuda.in`) and splits its content into two files: -1. `requirements/nightly_torch_test.txt`: Contains dependencies +1. `requirements/test/nightly-torch.txt`: Contains dependencies except PyTorch-related. 2. `torch_nightly_test.txt`: Contains only PyTorch-related packages. """ -input_file = "requirements/test.in" -output_file = "requirements/nightly_torch_test.txt" +input_file = "requirements/test/cuda.in" +output_file = "requirements/test/nightly-torch.txt" # white list of packages that are not compatible with PyTorch nightly directly # with pip install. Please add your package to this list if it is not compatible diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 1ba1f81564c..41c05efd201 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -36,8 +36,6 @@ SEPARATE_GROUPS = [ EXCLUDE = [ "vllm/model_executor/models", "vllm/model_executor/layers/fla/ops", - # Ignore triton kernels in ops. - "vllm/v1/attention/ops", # TODO: Remove these entries after fixing mypy errors. "vllm/benchmarks", ] diff --git a/use_existing_torch.py b/use_existing_torch.py index 7c58a34d69d..39c327e9670 100644 --- a/use_existing_torch.py +++ b/use_existing_torch.py @@ -30,8 +30,8 @@ def main(argv): args = parser.parse_args(argv) for file in ( - *glob.glob("requirements/*.txt"), - *glob.glob("requirements/*.in"), + *glob.glob("requirements/**/*.txt", recursive=True), + *glob.glob("requirements/**/*.in", recursive=True), "pyproject.toml", ): with open(file) as f: diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index d59b74782be..8c2659b9c7e 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -336,9 +336,13 @@ def _rocm_aiter_fused_topk_fake( router_logits: torch.Tensor, top_k: int, gate_up: bool, -) -> None: - # tuple[torch.Tensor, torch.Tensor]: - pass +) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = x.shape[0] + topk_weights = torch.empty( + (num_tokens, top_k), dtype=torch.float32, device=x.device + ) + topk_indices = torch.empty((num_tokens, top_k), dtype=torch.int32, device=x.device) + return topk_weights, topk_indices # Cache whether aiter supports FP8 MLA parameters @@ -1918,7 +1922,7 @@ class rocm_aiter_ops: @staticmethod def shuffle_weight( - self, tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) + tensor: torch.Tensor, layout: tuple[int, int] = (16, 16) ) -> torch.Tensor: from aiter.ops.shuffle import shuffle_weight diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 0c2a53ec02e..d6780185be9 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -435,6 +435,7 @@ def fused_qk_norm_rope( cos_sin_cache: torch.Tensor, is_neox: bool, position_ids: torch.Tensor, + forced_token_heads_per_warp: int = -1, ) -> None: torch.ops._C.fused_qk_norm_rope( qkv, @@ -448,6 +449,7 @@ def fused_qk_norm_rope( cos_sin_cache, is_neox, position_ids, + forced_token_heads_per_warp, ) @@ -3491,3 +3493,38 @@ if hasattr(torch.ops._C, "hadacore_transform"): @register_fake("_C::hadacore_transform") def _hadacore_transform_fake(x: torch.Tensor, inplace: bool) -> torch.Tensor: return torch.empty_like(x) if not inplace else x + + +if hasattr(torch.ops._C, "minimax_allreduce_rms"): + + @register_fake("_C::minimax_allreduce_rms") + def _minimax_allreduce_rms_fake( + input: torch.Tensor, + norm_weight: torch.Tensor, + workspace: torch.Tensor, + rank: int, + nranks: int, + eps: float, + ) -> torch.Tensor: + return torch.empty_like(input) + + +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + + @register_fake("_C::minimax_allreduce_rms_qk") + def _minimax_allreduce_rms_qk_fake( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + workspace: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + token_num = qkv.shape[0] + return ( + torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device), + torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device), + ) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 646322163ee..a0a321173d1 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -144,6 +144,46 @@ def _xpu_mxfp8_quantize_fake( return x.to(dtype), x_s.to(torch.float8_e8m0fnu) +def _xpu_mxfp4_quantize_impl( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + MXFP4_BLOCK_SIZE = 32 + eps = 1e-10 + assert x.ndim == 2, "input must be 2-D" + assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, ( + f"last dimension {x.shape[-1]} must be divisible by group_size " + f"{MXFP4_BLOCK_SIZE}" + ) + assert x.is_contiguous(), "input groups must be contiguous" + + M, N = x.shape + + # Packed FP4 output: two nibbles per byte + x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8) + x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32) + + torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps) + + x_q = x_q.view(torch.float4_e2m1fn_x2) + x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format) + return x_q, x_s + + +def _xpu_mxfp4_quantize_fake( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + MXFP4_BLOCK_SIZE = 32 + M, N = x.shape + + # Packed FP4 output: two nibbles per byte + x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8) + x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32) + + x_q = x_q.view(torch.float4_e2m1fn_x2) + x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format) + return x_q, x_s + + # Global flag to ensure ops are registered only once _OPS_REGISTERED = False @@ -258,6 +298,9 @@ class xpu_ops: # alibi_slopes = alibi_slopes, # softcap=softcap, return_softmax_lse=return_softmax_lse, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, ) @staticmethod @@ -552,6 +595,12 @@ class xpu_ops: fake_impl=_xpu_mxfp8_quantize_fake, ) + direct_register_custom_op( + op_name="xpu_mxfp4_quantize", + op_func=_xpu_mxfp4_quantize_impl, + fake_impl=_xpu_mxfp4_quantize_fake, + ) + _OPS_REGISTERED = True diff --git a/vllm/benchmarks/datasets/__init__.py b/vllm/benchmarks/datasets/__init__.py new file mode 100644 index 00000000000..5d5e172e7b4 --- /dev/null +++ b/vllm/benchmarks/datasets/__init__.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.benchmarks.datasets.datasets import ( + DEFAULT_NUM_PROMPTS, + AIMODataset, + ASRDataset, + BenchmarkDataset, + BlazeditDataset, + BurstGPTDataset, + ConversationDataset, + CustomDataset, + CustomMMDataset, + HuggingFaceDataset, + InstructCoderDataset, + MLPerfDataset, + MMStarDataset, + MMVUDataset, + MTBenchDataset, + MultiModalConversationDataset, + NextEditPredictionDataset, + PrefixRepetitionRandomDataset, + RandomDataset, + RandomDatasetForReranking, + RandomMultiModalDataset, + SampleRequest, + ShareGPTDataset, + SonnetDataset, + SpecBench, + VisionArenaDataset, + add_dataset_parser, + add_random_dataset_base_args, + add_random_multimodal_dataset_args, + gen_prompt_decode_to_target_len, + get_samples, + is_valid_sequence, + lora_path_on_disk, + lora_tokenizer_cache, + process_image, + process_video, + zeta_prompt, +) +from vllm.benchmarks.datasets.utils import RangeRatio + +__all__ = [ + "DEFAULT_NUM_PROMPTS", + "AIMODataset", + "ASRDataset", + "BenchmarkDataset", + "BlazeditDataset", + "BurstGPTDataset", + "ConversationDataset", + "CustomDataset", + "CustomMMDataset", + "HuggingFaceDataset", + "InstructCoderDataset", + "MLPerfDataset", + "MMStarDataset", + "MMVUDataset", + "MTBenchDataset", + "MultiModalConversationDataset", + "NextEditPredictionDataset", + "PrefixRepetitionRandomDataset", + "RandomDataset", + "RandomDatasetForReranking", + "RandomMultiModalDataset", + "SampleRequest", + "ShareGPTDataset", + "SonnetDataset", + "SpecBench", + "VisionArenaDataset", + "add_dataset_parser", + "add_random_dataset_base_args", + "add_random_multimodal_dataset_args", + "gen_prompt_decode_to_target_len", + "get_samples", + "is_valid_sequence", + "lora_path_on_disk", + "lora_tokenizer_cache", + "process_image", + "process_video", + "RangeRatio", + "zeta_prompt", +] diff --git a/vllm/benchmarks/datasets/create_txt_slices_dataset.py b/vllm/benchmarks/datasets/create_txt_slices_dataset.py new file mode 100644 index 00000000000..3f7c5028a20 --- /dev/null +++ b/vllm/benchmarks/datasets/create_txt_slices_dataset.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Convert a plain-text file (local path or URL) into a JSONL dataset +compatible with ``CustomDataset`` (``--dataset-name custom``), by +randomly slicing the tokenized text into prompts. + +Each line of the output JSONL contains a ``prompt`` (decoded from a random +slice of the tokenized source text) and an ``output_tokens`` count. + +Usage +----- +:: + + python -m vllm.benchmarks.datasets.create_txt_slices_dataset \\ + --input sonnet.txt \\ + --output sonnet_dataset.jsonl \\ + --tokenizer gpt2 \\ + --num-prompts 1000 \\ + --input-len 1024 \\ + --output-len 128 + +The resulting JSONL file can then be used with the serving benchmark:: + + python -m vllm.benchmarks.serve \\ + --dataset-name custom \\ + --dataset-path sonnet_dataset.jsonl \\ + ... +""" + +from __future__ import annotations + +import argparse +import json +import logging +import random +import urllib.request + +import numpy as np +from transformers import AutoTokenizer + +from vllm.benchmarks.datasets.utils import RangeRatio, get_sampling_params + +logger = logging.getLogger(__name__) + + +def load_text(path: str) -> str: + """Load text from a local file or URL.""" + if path.startswith(("http://", "https://")): + with urllib.request.urlopen(path) as response: + return response.read().decode("utf-8") + with open(path, encoding="utf-8") as f: + return f.read() + + +def create_txt_slices_jsonl( + *, + input_path: str, + output_path: str, + tokenizer_name: str, + num_prompts: int, + input_len: int, + output_len: int, + range_ratio: RangeRatio = 0.0, + seed: int = 0, + trust_remote_code: bool = False, +) -> None: + """Read *input_path*, slice it into prompts, and write JSONL to + *output_path*.""" + + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name, trust_remote_code=trust_remote_code + ) + + text = load_text(input_path) + if not text: + raise ValueError("The text file is empty and cannot be sampled from.") + + token_ids = tokenizer(text, add_special_tokens=False).input_ids + if not token_ids: + raise ValueError("Tokenizing the text produced zero tokens; cannot sample.") + + rng_np = np.random.default_rng(seed) + rng_py = random.Random(seed) + + input_lens, output_lens, _ = get_sampling_params( + rng_np, + num_prompts, + range_ratio, + input_len, + output_len, + tokenizer, + ) + + num_available_tokens = len(token_ids) + + records: list[dict[str, object]] = [] + for i in range(num_prompts): + req_input_len = int(input_lens[i]) + req_output_len = int(output_lens[i]) + + # Randomly select a start position and slice with cycling + start_pos = rng_py.randint(0, num_available_tokens - 1) + prompt_token_ids = [ + token_ids[(start_pos + j) % num_available_tokens] + for j in range(req_input_len) + ] + prompt = tokenizer.decode(prompt_token_ids, skip_special_tokens=False) + + records.append({"prompt": prompt, "output_tokens": req_output_len}) + + with open(output_path, "w", encoding="utf-8") as f: + for record in records: + f.write(json.dumps(record, ensure_ascii=False) + "\n") + + logger.info( + "Wrote %d prompts to %s", + len(records), + output_path, + ) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser( + description="Convert a plain-text file into a JSONL dataset " + "for CustomDataset (--dataset-name custom).", + ) + parser.add_argument( + "--input", + required=True, + help="Path or URL to the source text file.", + ) + parser.add_argument( + "--output", + required=True, + help="Path for the output JSONL file.", + ) + parser.add_argument( + "--tokenizer", + required=True, + help="HuggingFace tokenizer name or path.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompt samples to generate (default: 1000).", + ) + parser.add_argument( + "--input-len", + type=int, + default=1024, + help="Target number of input tokens per prompt (default: 1024).", + ) + parser.add_argument( + "--output-len", + type=int, + default=128, + help="Target number of output tokens per prompt (default: 128).", + ) + parser.add_argument( + "--range-ratio", + type=str, + default="0.0", + help="Range ratio for input/output length sampling (default: 0.0). " + "A single float applies to both ISL and OSL. " + 'A JSON dict like \'{"input": 0.3, "output": 0.5}\' sets them ' + "independently. Values must be in [0, 1).", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="Random seed for reproducibility (default: 0).", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code from HuggingFace.", + ) + + args = parser.parse_args(argv) + + logging.basicConfig(level=logging.INFO) + + # Parse --range-ratio: try float first, then JSON dict. + range_ratio: RangeRatio + try: + range_ratio = float(args.range_ratio) + except ValueError: + import json as _json + + range_ratio = _json.loads(args.range_ratio) + + create_txt_slices_jsonl( + input_path=args.input, + output_path=args.output, + tokenizer_name=args.tokenizer, + num_prompts=args.num_prompts, + input_len=args.input_len, + output_len=args.output_len, + range_ratio=range_ratio, + seed=args.seed, + trust_remote_code=args.trust_remote_code, + ) + + +if __name__ == "__main__": + main() diff --git a/vllm/benchmarks/datasets.py b/vllm/benchmarks/datasets/datasets.py similarity index 94% rename from vllm/benchmarks/datasets.py rename to vllm/benchmarks/datasets/datasets.py index dd71762b5ba..986c0c85625 100644 --- a/vllm/benchmarks/datasets.py +++ b/vllm/benchmarks/datasets/datasets.py @@ -22,10 +22,10 @@ import random from abc import ABC, abstractmethod from collections.abc import Callable, Iterator, Mapping from contextlib import suppress -from copy import deepcopy -from dataclasses import dataclass +from dataclasses import dataclass, replace from functools import cache from io import BytesIO +from pathlib import Path from tempfile import NamedTemporaryFile from typing import Any, cast @@ -35,6 +35,11 @@ from huggingface_hub import snapshot_download from PIL import Image from typing_extensions import deprecated +from vllm.benchmarks.datasets.utils import ( + RangeRatio, + _resolve_range_ratios, + get_sampling_params, +) from vllm.inputs import MultiModalDataDict from vllm.lora.request import LoRARequest from vllm.lora.utils import get_adapter_absolute_path @@ -60,10 +65,6 @@ logger = logging.getLogger(__name__) DEFAULT_NUM_PROMPTS = 1000 -# ----------------------------------------------------------------------------- -# Data Classes -# ----------------------------------------------------------------------------- - @dataclass class SampleRequest: @@ -71,9 +72,9 @@ class SampleRequest: Represents a single inference request for benchmarking. """ - prompt: str | list[str] + prompt: str | list[str] | list[dict] prompt_len: int - expected_output_len: int + expected_output_len: int | None multi_modal_data: MultiModalDataDict | dict | list[dict] | None = None lora_request: LoRARequest | None = None request_id: str | None = None @@ -110,7 +111,7 @@ class BenchmarkDataset(ABC): # default seed. self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED self.disable_shuffle = disable_shuffle - self.data = None + self.data: Any | None = None def apply_multimodal_chat_transformation( self, @@ -249,6 +250,7 @@ class BenchmarkDataset(ABC): num_requests: int, request_id_prefix: str = "", no_oversample: bool = False, + **kwargs, ) -> list[SampleRequest]: """ Abstract method to generate sample requests from the dataset. @@ -296,8 +298,10 @@ class BenchmarkDataset(ABC): needed = num_requests - len(requests) additional = [] for i in range(needed): - req = deepcopy(random.choice(requests)) - req.request_id = request_id_prefix + str(len(requests) + i) + req = replace( + random.choice(requests), + request_id=request_id_prefix + str(len(requests) + i), + ) additional.append(req) requests.extend(additional) logger.info("Oversampled requests to reach %d total samples.", num_requests) @@ -533,7 +537,7 @@ class RandomDataset(BenchmarkDataset): request_id_prefix: str = "", no_oversample: bool = False, prefix_len: int = DEFAULT_PREFIX_LEN, - range_ratio: float = DEFAULT_RANGE_RATIO, + range_ratio: RangeRatio = DEFAULT_RANGE_RATIO, input_len: int = DEFAULT_INPUT_LEN, output_len: int = DEFAULT_OUTPUT_LEN, batchsize: int = 1, @@ -542,24 +546,33 @@ class RandomDataset(BenchmarkDataset): lora_assignment: str = "random", **kwargs, ) -> list[SampleRequest]: - # validate total input tokens (prefix + sampled) is at least 1. + resolved_input_rr, _ = _resolve_range_ratios(range_ratio) + num_special = int(tokenizer.num_special_tokens_to_add()) real_input_len = max(0, int(input_len) - num_special) - min_sampled_input = math.floor(real_input_len * (1.0 - float(range_ratio))) + min_sampled_input = math.floor( + real_input_len * (1.0 - float(resolved_input_rr)) + ) min_total_input = int(prefix_len) + min_sampled_input if min_total_input < 1: raise ValueError( "--random-input-len is too small: with tokenizer special " - f"tokens {num_special} and --random-range-ratio {range_ratio}, " + f"tokens {num_special} and " + f"input range ratio {resolved_input_rr}, " "the minimum possible total input tokens (prefix + sampled) is " f"{min_total_input}. Increase --random-input-len and/or " - "--random-prefix-len, or decrease --random-range-ratio so that " - "prefix_len + floor(max(0, random_input_len - num_special)) " - "* (1 - range_ratio) >= 1." + "--random-prefix-len, or decrease the input range ratio " + "so that prefix_len + floor(max(0, random_input_len - " + "num_special)) * (1 - input_range_ratio) >= 1." ) - input_lens, output_lens, offsets = self.get_sampling_params( - num_requests, range_ratio, input_len, output_len, tokenizer + input_lens, output_lens, offsets = get_sampling_params( + self._rng, + num_requests, + range_ratio, + input_len, + output_len, + tokenizer, ) vocab_size = tokenizer.vocab_size @@ -661,55 +674,6 @@ class RandomDataset(BenchmarkDataset): ) return adjusted_tokens - def get_sampling_params( - self, - num_requests: int, - range_ratio: float, - input_len: int, - output_len: int, - tokenizer: TokenizerLike, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """ - Get the sampling parameters for the dataset. - """ - # Enforce range_ratio < 1 - if not (0.0 <= range_ratio < 1.0): - raise ValueError("range_ratio must be in [0, 1).") - num_special_tokens = int(tokenizer.num_special_tokens_to_add()) - real_input_len = max(0, int(input_len) - num_special_tokens) - # Bounds use floor for low and ceil for high - input_low = math.floor(real_input_len * (1 - range_ratio)) - input_high = math.ceil(real_input_len * (1 + range_ratio)) - output_low = math.floor(output_len * (1 - range_ratio)) - output_high = math.ceil(output_len * (1 + range_ratio)) - # Ensure the lower bound for output length is at least 1 to - # prevent sampling 0 tokens. - output_low = max(output_low, 1) - output_high = max(output_high, 1) - - if input_low > input_high: - raise ValueError( - f"Invalid input sampling interval: low={input_low} > high={input_high}" - ) - if output_low > output_high: - raise ValueError( - "Invalid output sampling interval: " - f"low={output_low} > high={output_high}" - ) - - logger.info( - "Sampling input_len from [%s, %s] and output_len from [%s, %s]", - input_low, - input_high, - output_low, - output_high, - ) - - input_lens = self._rng.integers(input_low, input_high + 1, size=num_requests) - output_lens = self._rng.integers(output_low, output_high + 1, size=num_requests) - offsets = self._rng.integers(0, tokenizer.vocab_size, size=num_requests) - return input_lens, output_lens, offsets - def generate_token_sequence( self, *, @@ -776,8 +740,11 @@ class RandomDatasetForReranking(RandomDataset): tokenizer: TokenizerLike, num_requests: int, request_id_prefix: str = "", - range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO, + no_oversample: bool = False, + prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN, + range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO, input_len: int = RandomDataset.DEFAULT_INPUT_LEN, + output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN, batchsize: int = 1, is_reranker: bool = True, **kwargs, @@ -786,8 +753,13 @@ class RandomDatasetForReranking(RandomDataset): query_len_param = (input_len // 2) - n_sep_tokens if is_reranker else input_len - query_lens, _, query_offsets = self.get_sampling_params( - 1, range_ratio, query_len_param, 0, tokenizer + query_lens, _, query_offsets = get_sampling_params( + self._rng, + 1, + range_ratio, + query_len_param, + 0, + tokenizer, ) query_len = int(query_lens[0]) @@ -800,8 +772,13 @@ class RandomDatasetForReranking(RandomDataset): else: doc_len_param = input_len - query_len - n_sep_tokens - doc_lens, _, doc_offsets = self.get_sampling_params( - num_requests, range_ratio, doc_len_param, 0, tokenizer + doc_lens, _, doc_offsets = get_sampling_params( + self._rng, + num_requests, + range_ratio, + doc_len_param, + 0, + tokenizer, ) vocab_size = tokenizer.vocab_size @@ -1175,9 +1152,10 @@ class RandomMultiModalDataset(RandomDataset): request_id_prefix: str = "", no_oversample: bool = False, prefix_len: int = RandomDataset.DEFAULT_PREFIX_LEN, - range_ratio: float = RandomDataset.DEFAULT_RANGE_RATIO, + range_ratio: RangeRatio = RandomDataset.DEFAULT_RANGE_RATIO, input_len: int = RandomDataset.DEFAULT_INPUT_LEN, output_len: int = RandomDataset.DEFAULT_OUTPUT_LEN, + batchsize: int = 1, limit_mm_per_prompt: dict[str, int] = DEFAULT_LIMIT_MM_PER_PROMPT, base_items_per_request: int = DEFAULT_BASE_ITEMS_PER_REQUEST, num_mm_items_range_ratio: float = DEFAULT_NUM_MM_ITEMS_RANGE_RATIO, @@ -1187,9 +1165,18 @@ class RandomMultiModalDataset(RandomDataset): enable_multimodal_chat: bool = DEFAULT_ENABLE_MULTIMODAL_CHAT, **kwargs, ) -> list[SampleRequest]: - # Get the sampling parameters for the dataset - input_lens, output_lens, offsets = self.get_sampling_params( - num_requests, range_ratio, input_len, output_len, tokenizer + if batchsize != 1: + raise NotImplementedError( + "batchsize > 1 is not supported for RandomMultiModalDataset." + ) + + input_lens, output_lens, offsets = get_sampling_params( + self._rng, + num_requests, + range_ratio, + input_len, + output_len, + tokenizer, ) ( @@ -1326,16 +1313,16 @@ class ShareGPTDataset(BenchmarkDataset): self, tokenizer: TokenizerLike, num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, lora_path: str | None = None, max_loras: int | None = None, output_len: int | None = None, enable_multimodal_chat: bool = False, - request_id_prefix: str = "", - no_oversample: bool = False, lora_assignment: str = "random", **kwargs, - ) -> list: - samples: list = [] + ) -> list[SampleRequest]: + samples: list[SampleRequest] = [] ind = 0 for entry in self.data: if len(samples) >= num_requests: @@ -1436,6 +1423,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser): "custom_mm", "prefix_repetition", "spec_bench", + "speed_bench", ], help="Name of the dataset to benchmark on.", ) @@ -1449,8 +1437,8 @@ def add_dataset_parser(parser: FlexibleArgumentParser): type=str, default=None, action=_ValidateDatasetArgs, - help="Path to the sharegpt/sonnet dataset. " - "Or the huggingface dataset ID if using HF dataset.", + help="Path to the sharegpt/sonnet dataset or the HF dataset ID if " + "using HF dataset.", ) parser.add_argument( "--no-oversample", @@ -1620,6 +1608,34 @@ def add_dataset_parser(parser: FlexibleArgumentParser): "repetition dataset.", ) + speed_bench_group = parser.add_argument_group("speed bench dataset options") + speed_bench_group.add_argument( + "--speed-bench-dataset-subset", + type=str, + default="qualitative", + choices={ + "qualitative", + "throughput_1k", + "throughput_2k", + "throughput_8k", + "throughput_16k", + "throughput_32k", + }, + help="Subset of the SPEED-Bench dataset.", + ) + speed_bench_group.add_argument( + "--speed-bench-output-len", + type=int, + default=4096, + help="Num of output tokens per request, used only for speed bench dataset.", + ) + speed_bench_group.add_argument( + "--speed-bench-category", + type=str, + default=None, + help="Category for speed bench dataset. If None, use all categories.", + ) + def add_random_dataset_base_args( parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup, @@ -1648,12 +1664,12 @@ def add_random_dataset_base_args( ) parser_or_group.add_argument( "--random-range-ratio", - type=float, - default=0.0, + type=str, + default="0.0", help="Range ratio for sampling input/output length, " - "used only for random sampling. Must be in the range [0, 1) to define " - "a symmetric sampling range" - "[length * (1 - range_ratio), length * (1 + range_ratio)].", + "used only for random sampling. A single float applies to both " + 'ISL and OSL. A JSON dict like \'{"input": 0.3, "output": 0.5}\' ' + "sets them independently. Values must be in [0, 1).", ) parser_or_group.add_argument( "--random-prefix-len", @@ -1786,10 +1802,25 @@ def add_random_multimodal_dataset_args( ) +def _parse_range_ratio(value: str) -> RangeRatio: + """Parse a ``--random-range-ratio`` CLI string. + + Accepts either a plain float (``"0.3"``) or a JSON dict + (``'{"input": 0.3, "output": 0.5}'``). + """ + try: + return float(value) + except ValueError: + return json.loads(value) + + def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: if not hasattr(args, "request_id_prefix"): args.request_id_prefix = "" + if hasattr(args, "random_range_ratio") and isinstance(args.random_range_ratio, str): + args.random_range_ratio = _parse_range_ratio(args.random_range_ratio) + if args.dataset_name == "custom": dataset = CustomDataset( dataset_path=args.dataset_path, disable_shuffle=args.disable_shuffle @@ -2073,6 +2104,19 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]: request_id_prefix=args.request_id_prefix, no_oversample=args.no_oversample, ), + "speed_bench": lambda: SpeedBench( + dataset_path=args.dataset_path, + dataset_subset=args.speed_bench_dataset_subset, + category=args.speed_bench_category, + disable_shuffle=args.disable_shuffle, + ).sample( + num_requests=args.num_prompts, + tokenizer=tokenizer, + output_len=args.speed_bench_output_len, + enable_multimodal_chat=args.enable_multimodal_chat, + request_id_prefix=args.request_id_prefix, + no_oversample=args.no_oversample, + ), } try: @@ -2120,7 +2164,7 @@ class CustomDataset(BenchmarkDataset): # This will be the standardized format which load_data() # has to convert into depending on the filetype of dataset_path. # sample() will assume this standardized format of self.data - self.data = [] + self.data: list[dict] = [] # Load the JSONL file if self.dataset_path.endswith(".jsonl"): @@ -2149,15 +2193,15 @@ class CustomDataset(BenchmarkDataset): self, tokenizer: TokenizerLike, num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, lora_path: str | None = None, max_loras: int | None = None, output_len: int | None = None, enable_multimodal_chat: bool = False, skip_chat_template: bool = False, - request_id_prefix: str = "", - no_oversample: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: # load all data if needed self.num_available_samples = len(self.data) if num_requests <= 0: @@ -2168,7 +2212,7 @@ class CustomDataset(BenchmarkDataset): num_requests, ) - sampled_requests = [] + sampled_requests: list[SampleRequest] = [] for i, item in enumerate(self.data): if len(sampled_requests) >= num_requests: break @@ -2252,7 +2296,7 @@ class CustomMMDataset(CustomDataset): request_id_prefix: str = "", no_oversample: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: # load all data if needed self.num_available_samples = len(self.data) if num_requests <= 0: @@ -2340,9 +2384,13 @@ class SpecBench(CustomDataset): if not getattr(self, "disable_shuffle", False): random.shuffle(self.data) - def sample(self, **kwargs) -> list: + def sample( + **kwargs, + ) -> list[SampleRequest]: # leverage CustomDataset sample - return super().sample(**kwargs) + return super().sample( + **kwargs, + ) # ----------------------------------------------------------------------------- @@ -2381,14 +2429,14 @@ class SonnetDataset(BenchmarkDataset): self, tokenizer: TokenizerLike, num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, prefix_len: int = DEFAULT_PREFIX_LEN, input_len: int = DEFAULT_INPUT_LEN, output_len: int = DEFAULT_OUTPUT_LEN, return_prompt_formatted: bool = False, - request_id_prefix: str = "", - no_oversample: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: # Calculate average token length for a poem line. tokenized_lines = [tokenizer(line).input_ids for line in self.data] avg_len = sum(len(tokens) for tokens in tokenized_lines) / len(tokenized_lines) @@ -2411,7 +2459,7 @@ class SonnetDataset(BenchmarkDataset): num_prefix_lines = max(round((prefix_len - base_offset) / avg_len), 0) prefix_lines = self.data[:num_prefix_lines] - samples = [] + samples: list[SampleRequest] = [] ind = 0 while len(samples) < num_requests: extra_lines = random.choices( @@ -2482,11 +2530,11 @@ class BurstGPTDataset(BenchmarkDataset): self, tokenizer: TokenizerLike, num_requests: int, - max_loras: int | None = None, - lora_path: str | None = None, request_id_prefix: str = "", no_oversample: bool = False, lora_assignment: str = "random", + max_loras: int | None = None, + lora_path: str | None = None, **kwargs, ) -> list[SampleRequest]: samples = [] @@ -2574,15 +2622,15 @@ class ConversationDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, - enable_multimodal_chat: bool = False, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, + enable_multimodal_chat: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: # Filter examples with at least 2 conversations filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2) - sampled_requests = [] + sampled_requests: list[SampleRequest] = [] ind = 0 dynamic_output = output_len is None @@ -2634,15 +2682,15 @@ class MultiModalConversationDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, - enable_multimodal_chat: bool = False, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, + enable_multimodal_chat: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: # Filter examples with at least 2 conversations filtered_data = self.data.filter(lambda x: len(x["conversations"]) >= 2) - sampled_requests = [] + sampled_requests: list[SampleRequest] = [] ind = 0 dynamic_output = output_len is None @@ -2703,12 +2751,12 @@ class VisionArenaDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, - enable_multimodal_chat: bool = False, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, + enable_multimodal_chat: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name) if parser_fn is None: raise ValueError(f"Unsupported dataset path: {self.hf_name}") @@ -2753,9 +2801,11 @@ class MMVUDataset(HuggingFaceDataset): DEFAULT_OUTPUT_LEN = 128 SUPPORTED_DATASET_PATHS = { - "yale-nlp/MMVU": lambda x: x["question"] - + " " - + (" ".join(f"{k}.{v}" for k, v in x["choices"].items())), + "yale-nlp/MMVU": lambda x: ( + x["question"] + + " " + + (" ".join(f"{k}.{v}" for k, v in x["choices"].items())) + ), } def __init__(self, **kwargs) -> None: @@ -2770,12 +2820,12 @@ class MMVUDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, - enable_multimodal_chat: bool = False, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, + enable_multimodal_chat: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: parser_fn = self.SUPPORTED_DATASET_PATHS.get(self.hf_name) if parser_fn is None: raise ValueError(f"Unsupported dataset path: {self.hf_name}") @@ -2838,15 +2888,15 @@ class InstructCoderDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, output_len: int | None = None, enable_multimodal_chat: bool = False, skip_chat_template: bool = False, - request_id_prefix: str = "", - no_oversample: bool = False, **kwargs, ) -> list[SampleRequest]: output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN - sampled_requests = [] + sampled_requests: list[SampleRequest] = [] for i, prompt in enumerate(self.sample_prompts(n=num_requests)): # apply template if not skip_chat_template: @@ -2903,15 +2953,15 @@ class MTBenchDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, output_len: int | None = None, enable_multimodal_chat: bool = False, skip_chat_template: bool = False, - request_id_prefix: str = "", - no_oversample: bool = False, **kwargs, - ) -> list: + ) -> list[SampleRequest]: output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN - sampled_requests = [] + sampled_requests: list[SampleRequest] = [] for i, item in enumerate(self.data): if len(sampled_requests) >= num_requests: @@ -2976,7 +3026,7 @@ class BlazeditDataset(HuggingFaceDataset): min_distance: float = 0.0, max_distance: float = 1.0, **kwargs, - ) -> list: + ) -> list[SampleRequest]: output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN sampled_requests = [] @@ -3050,12 +3100,12 @@ class AIMODataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, **kwargs, - ) -> list: - sampled_requests = [] + ) -> list[SampleRequest]: + sampled_requests: list[SampleRequest] = [] ind = 0 dynamic_output = output_len is None @@ -3228,18 +3278,18 @@ class ASRDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, **kwargs, - ) -> list: + ) -> list[SampleRequest]: output_len = output_len if output_len is not None else self.DEFAULT_OUTPUT_LEN if "openai" in getattr(tokenizer, "name_or_path", ""): prompt = "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>" else: prompt = "" prompt_len = len(tokenizer(prompt).input_ids) - sampled_requests = [] + sampled_requests: list[SampleRequest] = [] ind = 0 skipped = 0 asr_min_audio_len_sec = kwargs.get("asr_min_audio_len_sec") @@ -3326,9 +3376,9 @@ class MLPerfDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, **kwargs, ) -> list[SampleRequest]: # Force dynamic output length based on reference completion. @@ -3405,12 +3455,12 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset): self, tokenizer: TokenizerLike, num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, prefix_len: int = DEFAULT_PREFIX_LEN, suffix_len: int = DEFAULT_SUFFIX_LEN, num_prefixes: int = DEFAULT_NUM_PREFIXES, output_len: int = DEFAULT_OUTPUT_LEN, - request_id_prefix: str = "", - no_oversample: bool = False, **kwargs, ) -> list[SampleRequest]: vocab_size = tokenizer.vocab_size @@ -3421,7 +3471,7 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset): f"to num_prefixes ({num_prefixes})" ) - def _generate_exact_length_tokens(target_length: int) -> list[int]: + def _generate_exact_length_tokens(target_length: int) -> tuple[list[int], int]: """Generate tokens that decode and re-encode to exactly target_length.""" # Generate random tokens @@ -3491,10 +3541,10 @@ class MMStarDataset(HuggingFaceDataset): self, tokenizer: TokenizerLike, num_requests: int, - output_len: int | None = None, - enable_multimodal_chat: bool = False, request_id_prefix: str = "", no_oversample: bool = False, + output_len: int | None = None, + enable_multimodal_chat: bool = False, **kwargs, ) -> list[SampleRequest]: # If --hf-output-len is not set, use the default output length. @@ -3516,6 +3566,7 @@ class MMStarDataset(HuggingFaceDataset): # if enable_multimodal_chat is False). prompt_len = len(tokenizer(question_text).input_ids) + prompt: str | list[dict] if enable_multimodal_chat: # If multimodal content should be embedded in the chat message, # convert to [{"role":"user","content":[...]}] @@ -3543,3 +3594,48 @@ class MMStarDataset(HuggingFaceDataset): sampled_requests, num_requests, request_id_prefix, no_oversample ) return sampled_requests + + +# ----------------------------------------------------------------------------- +# Speed Bench Dataset Implementation +# ----------------------------------------------------------------------------- + + +class SpeedBench(CustomDataset): + """ + Implements the SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench + Download the dataset using: + curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 - + """ # noqa: E501 + + def __init__(self, **kwargs) -> None: + self.dataset_subset = kwargs.pop("dataset_subset", "qualitative") + self.category = kwargs.pop("category", None) + super().__init__(**kwargs) + self.load_data() + + def load_data(self) -> None: + if self.dataset_path is None: + raise ValueError("dataset_path must be provided for loading data.") + + self.data = [] + + # Load the JSONL file + jsonl_data = pd.read_json( + path_or_buf=Path(self.dataset_path) / f"{self.dataset_subset}.jsonl", + lines=True, + ) + + # check if the JSONL file has a 'turns' column + if "messages" not in jsonl_data.columns: + raise ValueError("JSONL file must contain a 'messages' column.") + + for _, row in jsonl_data.iterrows(): + # sample only from a specific category if specified + if (not self.category) or (self.category == row["category"]): + prompt = row["messages"][0]["content"] + self.data.append({"prompt": prompt}) + + random.seed(self.random_seed) + if not getattr(self, "disable_shuffle", False): + random.shuffle(self.data) diff --git a/vllm/benchmarks/datasets/utils.py b/vllm/benchmarks/datasets/utils.py new file mode 100644 index 00000000000..bc5a4340dd6 --- /dev/null +++ b/vllm/benchmarks/datasets/utils.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shared utilities for benchmark dataset sampling. +""" + +import logging +import math + +import numpy as np + +from vllm.tokenizers import TokenizerLike + +logger = logging.getLogger(__name__) + +# Type alias: a single float applies to both ISL and OSL; a dict allows +# specifying them independently via ``{"input": …, "output": …}``. +RangeRatio = float | dict[str, float] + + +def _resolve_range_ratios( + range_ratio: RangeRatio, +) -> tuple[float, float]: + """Return ``(input_range_ratio, output_range_ratio)`` from *range_ratio*. + + *range_ratio* is either a single float (used for both input and output) + or a dict with ``"input"`` and ``"output"`` keys. + """ + if isinstance(range_ratio, dict): + try: + return float(range_ratio["input"]), float(range_ratio["output"]) + except KeyError as exc: + raise ValueError( + "When range_ratio is a dict it must contain 'input' and " + f"'output' keys, got: {sorted(range_ratio)}" + ) from exc + ratio = float(range_ratio) + return ratio, ratio + + +def get_sampling_params( + rng: np.random.Generator, + num_requests: int, + range_ratio: RangeRatio, + input_len: int, + output_len: int, + tokenizer: TokenizerLike, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Sample per-request input/output token lengths and vocab offsets. + + Lengths are drawn uniformly from integer ranges around the configured + means, controlled by *range_ratio*. It may be a single ``float`` + (applied to both input and output) or a ``dict`` with ``"input"`` and + ``"output"`` keys for independent control. + + Tokenizer special tokens are subtracted from ``input_len`` before + computing the sampling interval. + + Returns: + (input_lens, output_lens, offsets) – three 1-D ``np.ndarray`` of + shape ``(num_requests,)``. + """ + input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio) + + if not (0.0 <= input_range_ratio < 1.0): + raise ValueError("input_range_ratio must be in [0, 1).") + if not (0.0 <= output_range_ratio < 1.0): + raise ValueError("output_range_ratio must be in [0, 1).") + num_special_tokens = int(tokenizer.num_special_tokens_to_add()) + real_input_len = max(0, int(input_len) - num_special_tokens) + input_low = math.floor(real_input_len * (1 - input_range_ratio)) + input_high = math.ceil(real_input_len * (1 + input_range_ratio)) + output_low = math.floor(output_len * (1 - output_range_ratio)) + output_high = math.ceil(output_len * (1 + output_range_ratio)) + # Ensure the lower bound for output length is at least 1 to + # prevent sampling 0 tokens. + output_low = max(output_low, 1) + output_high = max(output_high, 1) + + if input_low > input_high: + raise ValueError( + f"Invalid input sampling interval: low={input_low} > high={input_high}" + ) + if output_low > output_high: + raise ValueError( + f"Invalid output sampling interval: low={output_low} > high={output_high}" + ) + + logger.info( + "Sampling input_len from [%s, %s] and output_len from [%s, %s]", + input_low, + input_high, + output_low, + output_high, + ) + + input_lens = rng.integers(input_low, input_high + 1, size=num_requests) + output_lens = rng.integers(output_low, output_high + 1, size=num_requests) + offsets = rng.integers(0, tokenizer.vocab_size, size=num_requests) + return input_lens, output_lens, offsets diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index b0ef67889d1..61af098f80d 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -237,6 +237,8 @@ async def async_request_openai_completions( generated_text += text or "" elif usage := data.get("usage"): output.output_tokens = usage.get("completion_tokens") + if (pt := usage.get("prompt_tokens")) is not None: + output.prompt_len = pt if first_chunk_received: output.success = True else: @@ -358,6 +360,8 @@ async def async_request_openai_chat_completions( generated_text += content or "" elif usage := data.get("usage"): output.output_tokens = usage.get("completion_tokens") + if (pt := usage.get("prompt_tokens")) is not None: + output.prompt_len = pt most_recent_timestamp = timestamp diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 53ae6ca6a80..980ed1d412f 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -439,7 +439,7 @@ def calculate_metrics( ).input_ids ) actual_output_lens.append(output_len) - total_input += input_requests[i].prompt_len + total_input += outputs[i].prompt_len tpot = 0 if output_len > 1: latency_minus_ttft = outputs[i].latency - outputs[i].ttft diff --git a/vllm/benchmarks/startup.py b/vllm/benchmarks/startup.py index 375b8f9fac3..095fdb07327 100644 --- a/vllm/benchmarks/startup.py +++ b/vllm/benchmarks/startup.py @@ -16,7 +16,7 @@ import shutil import tempfile import time from contextlib import contextmanager -from typing import Any +from typing import Any, NamedTuple import numpy as np from tqdm import tqdm @@ -27,6 +27,82 @@ from vllm.benchmarks.lib.utils import ( ) from vllm.engine.arg_utils import EngineArgs +PERCENTAGES = [10, 25, 50, 75, 90, 99] + + +class MetricDesc(NamedTuple): + """Descriptor for a metric to collect from each iteration.""" + + iter_key: str # key in the iteration result dict + suffix: str # result key suffix, e.g. "startup", "compilation" + display_name: str + + +class MetricStats(NamedTuple): + """Aggregated statistics for a single benchmark metric.""" + + key: str # e.g. "cold_startup", "warm_encoder_compilation" + display_name: str + values: list[float] + avg: float + percentiles: dict[int, float] + + +_BASE_METRICS = [ + MetricDesc("total_startup_time", "startup", "Startup time"), + MetricDesc("compilation_time", "compilation", "Compilation time"), +] +_ENCODER_METRIC = MetricDesc( + "encoder_compilation_time", + "encoder_compilation", + "Encoder compilation time", +) + + +def _compute_metric( + phase: str, + desc: MetricDesc, + iterations: list[dict[str, float]], +) -> MetricStats: + values = [m[desc.iter_key] for m in iterations] + arr = np.array(values) + return MetricStats( + key=f"{phase}_{desc.suffix}", + display_name=desc.display_name, + values=values, + avg=float(np.mean(arr)), + percentiles=dict(zip(PERCENTAGES, np.percentile(arr, PERCENTAGES).tolist())), + ) + + +def _collect_phase_metrics( + phase: str, + iterations: list[dict[str, float]], + has_encoder: bool, +) -> list[MetricStats]: + metrics = [_compute_metric(phase, desc, iterations) for desc in _BASE_METRICS] + if has_encoder: + metrics.append(_compute_metric(phase, _ENCODER_METRIC, iterations)) + return metrics + + +def _print_phase(phase_name: str, metrics: list[MetricStats]) -> None: + print(f"\n{phase_name}:") + for m in metrics: + print(f"Avg {m.display_name.lower()}: {m.avg:.2f} seconds") + for m in metrics: + print(f"{m.display_name} percentiles:") + for pct, val in m.percentiles.items(): + print(f" {pct}%: {val:.2f} seconds") + + +def _metric_to_json(m: MetricStats) -> dict[str, Any]: + return { + f"avg_{m.key}_time": m.avg, + f"{m.key}_times": m.values, + f"{m.key}_percentiles": m.percentiles, + } + @contextmanager def cold_startup(): @@ -72,6 +148,7 @@ def run_startup_in_subprocess(engine_args, result_queue): # Extract compilation time if available compilation_time = 0.0 + encoder_compilation_time = 0.0 if hasattr(llm.llm_engine, "vllm_config"): vllm_config = llm.llm_engine.vllm_config if ( @@ -79,11 +156,15 @@ def run_startup_in_subprocess(engine_args, result_queue): and vllm_config.compilation_config is not None ): compilation_time = vllm_config.compilation_config.compilation_time + encoder_compilation_time = ( + vllm_config.compilation_config.encoder_compilation_time + ) result_queue.put( { "total_startup_time": total_startup_time, "compilation_time": compilation_time, + "encoder_compilation_time": encoder_compilation_time, } ) @@ -93,65 +174,20 @@ def run_startup_in_subprocess(engine_args, result_queue): def save_to_pytorch_benchmark_format( - args: argparse.Namespace, results: dict[str, Any] + args: argparse.Namespace, metrics: list[MetricStats] ) -> None: base_name = os.path.splitext(args.output_json)[0] - - cold_startup_records = convert_to_pytorch_benchmark_format( - args=args, - metrics={ - "avg_cold_startup_time": [results["avg_cold_startup_time"]], - }, - extra_info={ - "cold_startup_times": results["cold_startup_times"], - "cold_startup_percentiles": results["cold_startup_percentiles"], - }, - ) - if cold_startup_records: - write_to_json(f"{base_name}.cold_startup.pytorch.json", cold_startup_records) - - cold_compilation_records = convert_to_pytorch_benchmark_format( - args=args, - metrics={ - "avg_cold_compilation_time": [results["avg_cold_compilation_time"]], - }, - extra_info={ - "cold_compilation_times": results["cold_compilation_times"], - "cold_compilation_percentiles": results["cold_compilation_percentiles"], - }, - ) - if cold_compilation_records: - write_to_json( - f"{base_name}.cold_compilation.pytorch.json", cold_compilation_records - ) - - warm_startup_records = convert_to_pytorch_benchmark_format( - args=args, - metrics={ - "avg_warm_startup_time": [results["avg_warm_startup_time"]], - }, - extra_info={ - "warm_startup_times": results["warm_startup_times"], - "warm_startup_percentiles": results["warm_startup_percentiles"], - }, - ) - if warm_startup_records: - write_to_json(f"{base_name}.warm_startup.pytorch.json", warm_startup_records) - - warm_compilation_records = convert_to_pytorch_benchmark_format( - args=args, - metrics={ - "avg_warm_compilation_time": [results["avg_warm_compilation_time"]], - }, - extra_info={ - "warm_compilation_times": results["warm_compilation_times"], - "warm_compilation_percentiles": results["warm_compilation_percentiles"], - }, - ) - if warm_compilation_records: - write_to_json( - f"{base_name}.warm_compilation.pytorch.json", warm_compilation_records + for m in metrics: + records = convert_to_pytorch_benchmark_format( + args=args, + metrics={f"avg_{m.key}_time": [m.avg]}, + extra_info={ + f"{m.key}_times": m.values, + f"{m.key}_percentiles": m.percentiles, + }, ) + if records: + write_to_json(f"{base_name}.{m.key}.pytorch.json", records) def add_cli_args(parser: argparse.ArgumentParser): @@ -224,97 +260,46 @@ def main(args: argparse.Namespace): os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" print("Setting VLLM_ENABLE_V1_MULTIPROCESSING=0 to collect startup metrics.\n") + # Collect cold startup iterations print("Measuring cold startup time...\n") - cold_startup_times = [] - cold_compilation_times = [] + cold_iterations = [] for i in tqdm(range(args.num_iters_cold), desc="Cold startup iterations"): with cold_startup(): - metrics = create_llm_and_measure_startup() - cold_startup_times.append(metrics["total_startup_time"]) - cold_compilation_times.append(metrics["compilation_time"]) + cold_iterations.append(create_llm_and_measure_startup()) # Warmup for warm startup print("\nWarming up for warm startup measurement...\n") for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"): create_llm_and_measure_startup() + # Collect warm startup iterations print("\nMeasuring warm startup time...\n") - warm_startup_times = [] - warm_compilation_times = [] + warm_iterations = [] for i in tqdm(range(args.num_iters_warm), desc="Warm startup iterations"): - metrics = create_llm_and_measure_startup() - warm_startup_times.append(metrics["total_startup_time"]) - warm_compilation_times.append(metrics["compilation_time"]) + warm_iterations.append(create_llm_and_measure_startup()) - # Calculate statistics - cold_startup_array = np.array(cold_startup_times) - cold_compilation_array = np.array(cold_compilation_times) - warm_startup_array = np.array(warm_startup_times) - warm_compilation_array = np.array(warm_compilation_times) + # Determine if encoder compilation occurred in any iteration + has_encoder = any( + m["encoder_compilation_time"] > 0 for m in cold_iterations + warm_iterations + ) - avg_cold_startup = np.mean(cold_startup_array) - avg_cold_compilation = np.mean(cold_compilation_array) - avg_warm_startup = np.mean(warm_startup_array) - avg_warm_compilation = np.mean(warm_compilation_array) - - percentages = [10, 25, 50, 75, 90, 99] - cold_startup_percentiles = np.percentile(cold_startup_array, percentages) - cold_compilation_percentiles = np.percentile(cold_compilation_array, percentages) - warm_startup_percentiles = np.percentile(warm_startup_array, percentages) - warm_compilation_percentiles = np.percentile(warm_compilation_array, percentages) + cold_metrics = _collect_phase_metrics("cold", cold_iterations, has_encoder) + warm_metrics = _collect_phase_metrics("warm", warm_iterations, has_encoder) + all_metrics = cold_metrics + warm_metrics + # Print results print("\n" + "=" * 60) print("STARTUP TIME BENCHMARK RESULTS") print("=" * 60) - - # Cold startup statistics - print("\nCOLD STARTUP:") - print(f"Avg total startup time: {avg_cold_startup:.2f} seconds") - print(f"Avg compilation time: {avg_cold_compilation:.2f} seconds") - print("Startup time percentiles:") - for percentage, percentile in zip(percentages, cold_startup_percentiles): - print(f" {percentage}%: {percentile:.2f} seconds") - print("Compilation time percentiles:") - for percentage, percentile in zip(percentages, cold_compilation_percentiles): - print(f" {percentage}%: {percentile:.2f} seconds") - - # Warm startup statistics - print("\nWARM STARTUP:") - print(f"Avg total startup time: {avg_warm_startup:.2f} seconds") - print(f"Avg compilation time: {avg_warm_compilation:.2f} seconds") - print("Startup time percentiles:") - for percentage, percentile in zip(percentages, warm_startup_percentiles): - print(f" {percentage}%: {percentile:.2f} seconds") - print("Compilation time percentiles:") - for percentage, percentile in zip(percentages, warm_compilation_percentiles): - print(f" {percentage}%: {percentile:.2f} seconds") - + _print_phase("COLD STARTUP", cold_metrics) + _print_phase("WARM STARTUP", warm_metrics) print("=" * 60) # Output JSON results if specified if args.output_json: - results = { - "avg_cold_startup_time": float(avg_cold_startup), - "avg_cold_compilation_time": float(avg_cold_compilation), - "cold_startup_times": cold_startup_times, - "cold_compilation_times": cold_compilation_times, - "cold_startup_percentiles": dict( - zip(percentages, cold_startup_percentiles.tolist()) - ), - "cold_compilation_percentiles": dict( - zip(percentages, cold_compilation_percentiles.tolist()) - ), - "avg_warm_startup_time": float(avg_warm_startup), - "avg_warm_compilation_time": float(avg_warm_compilation), - "warm_startup_times": warm_startup_times, - "warm_compilation_times": warm_compilation_times, - "warm_startup_percentiles": dict( - zip(percentages, warm_startup_percentiles.tolist()) - ), - "warm_compilation_percentiles": dict( - zip(percentages, warm_compilation_percentiles.tolist()) - ), - } + results: dict[str, Any] = {} + for m in all_metrics: + results.update(_metric_to_json(m)) with open(args.output_json, "w") as f: json.dump(results, f, indent=4) - save_to_pytorch_benchmark_format(args, results) + save_to_pytorch_benchmark_format(args, all_metrics) diff --git a/vllm/collect_env.py b/vllm/collect_env.py index 0cf5681bcf5..1b94adba87e 100644 --- a/vllm/collect_env.py +++ b/vllm/collect_env.py @@ -46,6 +46,17 @@ SystemEnv = namedtuple( "nvidia_driver_version", "nvidia_gpu_models", "cudnn_version", + "xpu_available", + "xpu_runtime_version", + "intel_graphics_compiler_version", + "intel_gpu_models", + "oneapi_compiler_version", + "level_zero_loader_version", + "level_zero_driver_version", + "oneccl_version", + "libigdgmm_version", + "vllm_xpu_kernels_version", + "sycl_version", "pip_version", # 'pip' or 'pip3' "pip_packages", "conda_packages", @@ -277,6 +288,134 @@ def get_rocm_version(run_lambda): ) +def get_xpu_available(): + if TORCH_AVAILABLE and hasattr(torch, "xpu") and torch.xpu.is_available(): + return True + return False + + +def get_xpu_runtime_version(): + if TORCH_AVAILABLE and hasattr(torch.version, "xpu"): + return torch.version.xpu + return None + + +def get_pkg_version(run_lambda, pkg): + assert get_platform() == "linux" + + if pkg == "vllm_xpu_kernels": + rc, out, _ = run_lambda("pip show vllm-xpu-kernels") + if rc == 0: + match = re.search(r"Version: (.*)", out) + return match.group(1).strip() if match else None + return None + + pkg_map = { + "igc": ["intel-igc-core", "libigc2", "libigc1"], + "level_zero_loader": ["level-zero", "libze1"], + "level_zero_driver": ["libze-intel-gpu1", "intel-level-zero-gpu"], + "oneccl": ["intel-oneapi-ccl", "oneccl"], + "libigdgmm": ["libigdgmm12", "libigdgmm"], + } + + pkg_candidates = pkg_map.get(pkg, []) + if not pkg_candidates: + return None + + mgr_name = None + for mgr in ["dpkg", "dnf", "yum", "zypper"]: + rc, _, _ = run_lambda(f"which {mgr}") + if rc == 0: + mgr_name = mgr + break + + if not mgr_name: + return None + + ret = "" + index = -1 + + for pkg_name in pkg_candidates: + if not pkg_name: + continue + + cmd = "" + if mgr_name in ["dnf", "yum"]: + index = 1 + cmd = f"{mgr_name} list | grep -w {pkg_name}" + elif mgr_name == "zypper": + index = 2 + cmd = f"{mgr_name} info {pkg_name} | grep Version" + elif mgr_name == "dpkg": + index = 2 + cmd = f"{mgr_name} -l | grep -w {pkg_name}" + + if cmd: + out = run_and_read_all(run_lambda, cmd) + if out: + ret = out.splitlines()[0] + break + + if not ret or index == -1: + return None + + lst = re.sub(" +", " ", ret).strip().split(" ") + if len(lst) > index: + return lst[index] + + return None + + +def get_intel_graphics_compiler_version(run_lambda): + """Return Intel Graphics Compiler (IGC) version.""" + return get_pkg_version(run_lambda, "igc") + + +def get_level_zero_loader_version(run_lambda): + """Return Level Zero loader runtime version.""" + return get_pkg_version(run_lambda, "level_zero_loader") + + +def get_level_zero_driver_version(run_lambda): + """Return Level Zero driver version.""" + return get_pkg_version(run_lambda, "level_zero_driver") + + +def get_oneapi_ccl_version(run_lambda): + """Return oneAPI Collective Communications Library (oneCCL) version.""" + return get_pkg_version(run_lambda, "oneccl") + + +def get_libigdgmm_version(run_lambda): + return get_pkg_version(run_lambda, "libigdgmm") + + +def get_vllm_xpu_kernels_version(run_lambda): + return get_pkg_version(run_lambda, "vllm_xpu_kernels") + + +def get_intel_gpu_models(): + if TORCH_AVAILABLE and hasattr(torch, "xpu") and torch.xpu.is_available(): + device_count = torch.xpu.device_count() + return "\n".join( + "GPU {}: {}".format(i, torch.xpu.get_device_name(i)) + for i in range(device_count) + ) + return None + + +def get_oneapi_compiler_version(run_lambda): + """Return Intel oneAPI DPC++/C++ Compiler version via icpx.""" + return run_and_parse_first_match( + run_lambda, "icpx --version", r"oneAPI DPC\+\+/C\+\+ Compiler (\S+)" + ) + + +def get_sycl_version(run_lambda): + """Return SYCL/DPC++ compiler build version.""" + return run_and_parse_first_match(run_lambda, "icpx --version", r"\((\d[\d.]+)\)") + + def get_vllm_version(): from vllm import __version__, __version_tuple__ @@ -298,11 +437,12 @@ def get_vllm_version(): def summarize_vllm_build_flags(): - # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc. - return "CUDA Archs: {}; ROCm: {}".format( + flags = "CUDA Archs: {}; ROCm: {}; XPU: {}".format( os.environ.get("TORCH_CUDA_ARCH_LIST", "Not Set"), "Enabled" if os.environ.get("ROCM_HOME") else "Disabled", + "Enabled" if get_xpu_available() else "Disabled", ) + return flags def get_gpu_topo(run_lambda): @@ -574,6 +714,13 @@ def get_env_vars(): "OMP_", "MKL_", "NVIDIA", + "ZE_", + "ONEAPI_", + "SYCL_", + "NEOReadDebugKeys", + "IGC_", + "CCL_", + "I_MPI_", ) for k, v in os.environ.items(): if any(term in k.lower() for term in secret_terms): @@ -637,6 +784,17 @@ def get_env_info(): nvidia_gpu_models=get_gpu_info(run_lambda), nvidia_driver_version=get_nvidia_driver_version(run_lambda), cudnn_version=get_cudnn_version(run_lambda), + xpu_available=str(get_xpu_available()), + xpu_runtime_version=get_xpu_runtime_version(), + intel_graphics_compiler_version=get_intel_graphics_compiler_version(run_lambda), + intel_gpu_models=get_intel_gpu_models(), + oneapi_compiler_version=get_oneapi_compiler_version(run_lambda), + level_zero_loader_version=get_level_zero_loader_version(run_lambda), + level_zero_driver_version=get_level_zero_driver_version(run_lambda), + oneccl_version=get_oneapi_ccl_version(run_lambda), + libigdgmm_version=get_libigdgmm_version(run_lambda), + vllm_xpu_kernels_version=get_vllm_xpu_kernels_version(run_lambda), + sycl_version=get_sycl_version(run_lambda), hip_compiled_version=hip_compiled_version, hip_runtime_version=hip_runtime_version, miopen_runtime_version=miopen_runtime_version, @@ -676,26 +834,15 @@ PyTorch version : {torch_version} Is debug build : {is_debug_build} CUDA used to build PyTorch : {cuda_compiled_version} ROCM used to build PyTorch : {hip_compiled_version} +XPU used to build PyTorch : {xpu_runtime_version} ============================== Python Environment ============================== Python version : {python_version} Python platform : {python_platform} - -============================== - CUDA / GPU Info -============================== -Is CUDA available : {is_cuda_available} -CUDA runtime version : {cuda_runtime_version} -CUDA_MODULE_LOADING set to : {cuda_module_loading} -GPU models and configuration : {nvidia_gpu_models} -Nvidia driver version : {nvidia_driver_version} -cuDNN version : {cudnn_version} -HIP runtime version : {hip_runtime_version} -MIOpen runtime version : {miopen_runtime_version} -Is XNNPACK available : {is_xnnpack_available} - + +{gpu_info} ============================== CPU Info ============================== @@ -790,6 +937,35 @@ def pretty_str(envinfo): if envinfo.cuda_compiled_version is None: mutable_dict["cuda_compiled_version"] = "None" + # If the machine doesn't have XPU, report XPU fields as 'No XPU' + dynamic_xpu_fields = [ + "intel_graphics_compiler_version", + "intel_gpu_models", + "level_zero_loader_version", + "level_zero_driver_version", + "oneccl_version", + "libigdgmm_version", + "vllm_xpu_kernels_version", + ] + all_xpu_fields = dynamic_xpu_fields + [ + "oneapi_compiler_version", + "sycl_version", + ] + all_dynamic_xpu_fields_missing = all( + mutable_dict[field] is None for field in dynamic_xpu_fields + ) + xpu_available = mutable_dict.get("xpu_available") == "True" + if not xpu_available and all_dynamic_xpu_fields_missing: + for field in all_xpu_fields: + mutable_dict[field] = "No XPU" + if envinfo.xpu_runtime_version is None or envinfo.xpu_runtime_version == "N/A": + mutable_dict["xpu_runtime_version"] = "N/A" + + # If intel_gpu_models is multiline, start on the next line + mutable_dict["intel_gpu_models"] = maybe_start_on_next_line( + mutable_dict.get("intel_gpu_models") + ) + # Replace True with Yes, False with No mutable_dict = replace_bools(mutable_dict) @@ -811,6 +987,62 @@ def pretty_str(envinfo): mutable_dict["conda_packages"], "[conda] " ) mutable_dict["cpu_info"] = envinfo.cpu_info + + CUDA_FMT = """ +============================== + CUDA / GPU Info +============================== +Is CUDA available : {is_cuda_available} +CUDA runtime version : {cuda_runtime_version} +CUDA_MODULE_LOADING set to : {cuda_module_loading} +GPU models and configuration : {nvidia_gpu_models} +Nvidia driver version : {nvidia_driver_version} +cuDNN version : {cudnn_version} +HIP runtime version : {hip_runtime_version} +MIOpen runtime version : {miopen_runtime_version} +Is XNNPACK available : {is_xnnpack_available} +""".strip() + + XPU_FMT = """ +============================== + Intel XPU / GPU Info +============================== +Is XPU available : {xpu_available} +XPU runtime version : {xpu_runtime_version} +Intel GPU models : {intel_gpu_models} + +--Compile time-- +oneAPI compiler version : {oneapi_compiler_version} +SYCL compiler build : {sycl_version} +oneCCL version : {oneccl_version} + +--Runtime-- +Intel Graphics Compiler (IGC): {intel_graphics_compiler_version} +Intel GMM (libigdgmm) : {libigdgmm_version} +Level Zero loader version : {level_zero_loader_version} +Level Zero driver version : {level_zero_driver_version} +vLLM XPU kernels version : {vllm_xpu_kernels_version} +""".strip() + + invalid_vers = {"N/A", "Could not collect", "None"} + sections = [] + + if ( + mutable_dict.get("is_cuda_available") in ("True", "Yes") + or mutable_dict.get("cuda_compiled_version") not in invalid_vers + ): + sections.append(CUDA_FMT) + + if ( + mutable_dict.get("xpu_available") in ("True", "Yes") + or mutable_dict.get("xpu_runtime_version") not in invalid_vers + ): + sections.append(XPU_FMT) + + mutable_dict["gpu_info"] = ( + ("\n\n".join(sections) + "\n").format(**mutable_dict) if sections else "" + ) + return env_info_fmt.format(**mutable_dict) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 63dc8874069..7373ad75117 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -265,6 +265,7 @@ class CompilerManager: compile_range: Range, graph_index: int = 0, num_graphs: int = 1, + is_encoder: bool = False, ) -> Any: if graph_index == 0: # before compiling the first graph, record the start time @@ -282,7 +283,10 @@ class CompilerManager: # after loading the last graph for this shape, record the time. # there can be multiple graphs due to piecewise compilation. elapsed = time.perf_counter() - compilation_start_time - compilation_config.compilation_time += elapsed + if is_encoder: + compilation_config.encoder_compilation_time += elapsed + else: + compilation_config.compilation_time += elapsed logger.info_once( "Directly load the compiled graph(s) for compile range %s " "from the cache, took %.3f s", @@ -387,7 +391,10 @@ class CompilerManager: # after compiling the last graph, record the end time if graph_index == num_graphs - 1: elapsed = time.perf_counter() - compilation_start_time - compilation_config.compilation_time += elapsed + if is_encoder: + compilation_config.encoder_compilation_time += elapsed + else: + compilation_config.compilation_time += elapsed logger.info_once( "Compiling a graph for compile range %s takes %.2f s", str(compile_range), @@ -516,16 +523,31 @@ def _decompose_size_nodes(graph: fx.GraphModule) -> None: ) # Replace size node in each user's args. - # Dynamo always passes size as a direct arg: view(clone, size) - # → view(clone, d0, d1, ...) for user in list(node.users): - new_args = [] - for arg in user.args: - if arg is node: - new_args.extend(dims) - else: - new_args.append(arg) - user.args = tuple(new_args) + if ( + user.op == "call_function" + and user.target is operator.getitem + and len(user.args) == 2 + and user.args[0] is node + ): + # getitem(size, idx) → replace with dims[idx] directly. + idx = user.args[1] + assert isinstance(idx, int), ( + f"Expected literal int index for getitem on size(), " + f"got {type(idx).__name__}: {idx}" + ) + user.replace_all_uses_with(dims[idx]) + graph.graph.erase_node(user) + else: + # User consumes the full size tuple (e.g. view(clone, size)) + # → view(clone, d0, d1, ...) + new_args = [] + for arg in user.args: + if arg is node: + new_args.extend(dims) + else: + new_args.append(arg) + user.args = tuple(new_args) graph.graph.erase_node(node) @@ -1115,7 +1137,10 @@ class VllmBackend: logger.info_once( "Dynamo bytecode transform time: %.2f s", dynamo_time, scope="local" ) - self.compilation_config.compilation_time += dynamo_time + if self.is_encoder: + self.compilation_config.encoder_compilation_time += dynamo_time + else: + self.compilation_config.compilation_time += dynamo_time # Record Dynamo time in tracing if available start_time = int(torch_compile_start_time * 1e9) diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index 9c55a42a492..79daf00de66 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -507,6 +507,16 @@ def _support_torch_compile( hash_key, ) + # Hash-level dir; shared across ranks on the same node. + self.compilation_config.local_cache_dir = cache_dir + inductor_cache = os.path.join(cache_dir, "inductor_cache") + os.makedirs(inductor_cache, exist_ok=True) + # Process-wide: post-load execution, CUDA-graph capture, and later + # autotune/recompile all need to write under {hash}/inductor_cache/. + # Unconditional because torch's cache_dir() may have pre-filled the + # /tmp default during import, making setdefault a no-op. + os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache + rank = self.vllm_config.parallel_config.rank dp_rank = self.vllm_config.parallel_config.data_parallel_index cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}") diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index a712c013ce9..3a961cf5348 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -301,7 +301,7 @@ class ActivationQuantFusionPass(VllmPatternMatcherPass): pattern_silu_mul_nvfp4 = SiluMulNvfp4QuantPattern() pattern_silu_mul_nvfp4.register(self.patterns) - if current_platform.is_cuda_alike(): + if current_platform.is_cuda(): for quant_key in [kFp8Dynamic128Sym, kFp8Dynamic64Sym]: for is_scale_transposed in [False, True]: for is_e8m0 in [True, False]: diff --git a/vllm/compilation/passes/fusion/attn_quant_fusion.py b/vllm/compilation/passes/fusion/attn_quant_fusion.py index 98d2be387e1..3e2ed2bc707 100644 --- a/vllm/compilation/passes/fusion/attn_quant_fusion.py +++ b/vllm/compilation/passes/fusion/attn_quant_fusion.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform from vllm.utils.math_utils import round_up +from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement from .matcher_utils import MatcherQuantFP8 @@ -53,21 +54,43 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): @property def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - output_attn: torch.Tensor, - scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> torch.Tensor: + # When _USE_LAYERNAME is enabled (torch >= 2.11), layer_name is + # passed as an explicit pattern input so the pattern matcher + # treats it as a wildcard matching hoisted LayerName placeholders. + # Otherwise it stays as a closure constant (original behavior). + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _pattern_with_ln( # type: ignore[misc] + q, k, v, output_attn, scale, kv_cache_dummy_dep, layer_name + ): + at1 = auto_functionalized( + ATTN_OP, + query=q, + key=k, + value=v, + output=output_attn, + layer_name=layer_name, + output_scale=None, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + attn_out_view = RESHAPE_OP( + at1[1], [q.shape[0], self._num_heads * self._head_size] + ) + return self._quant_matcher(attn_out_view, scale)[0] + + return _pattern_with_ln + + def _pattern(q, k, v, output_attn, scale, kv_cache_dummy_dep): at1 = auto_functionalized( ATTN_OP, query=q, key=k, value=v, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=None, output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -81,14 +104,34 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): @property def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - output_attn: torch.Tensor, - scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> torch.Tensor: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _replacement_with_ln( # type: ignore[misc] + q, k, v, output_attn, scale, kv_cache_dummy_dep, layer_name + ): + output_attn = torch.empty( + [q.shape[0], self._num_heads, self._head_size], + dtype=FP8_DTYPE, + device=q.device, + ) + at1 = auto_functionalized( + ATTN_OP, + query=q, + key=k, + value=v, + output=output_attn, + layer_name=layer_name, + output_scale=scale, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size]) + + return _replacement_with_ln + + def _replacement(q, k, v, output_attn, scale, kv_cache_dummy_dep): output_attn = torch.empty( [q.shape[0], self._num_heads, self._head_size], dtype=FP8_DTYPE, @@ -100,7 +143,7 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): key=k, value=v, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=scale, output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -113,7 +156,7 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): dtype = self._dtype num_heads = self._num_heads head_size = self._head_size - return [ + inputs: list = [ self.empty(5, num_heads, head_size, dtype=dtype), # q self.empty(5, num_heads, head_size, dtype=dtype), # k self.empty(5, num_heads, head_size, dtype=dtype), # v @@ -121,6 +164,9 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): self.empty_fp32(1, 1), # scale self.empty(0, dtype=dtype), # kv_cache_dummy_dep ] + if _USE_LAYERNAME: + inputs.append(_encode_layer_name(self._layer_name)) + return inputs class AttnNvfp4QuantPattern( @@ -144,23 +190,64 @@ class AttnNvfp4QuantPattern( @property def pattern(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _pattern_with_ln( # type: ignore[misc] + q, + k, + v, + output_attn, + output_quant, + output_scale, + input_scale, + kv_cache_dummy_dep, + layer_name, + ): + at1 = auto_functionalized( + ATTN_OP, + query=q, + key=k, + value=v, + output=output_attn, + layer_name=layer_name, + output_scale=None, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + attn_out_view = RESHAPE_OP( + at1[1], [q.shape[0], self._num_heads * self._head_size] + ) + at2 = auto_functionalized( + self._QUANT_OP, + input=attn_out_view, + input_scale=input_scale, + is_sf_swizzled_layout=True, + output=output_quant, + output_scale=output_scale, + ) + return at2[1], torch.ops.aten.view.dtype(at2[2], FP8_DTYPE) + + return _pattern_with_ln + def _pattern( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - output_attn: torch.Tensor, - output_quant: torch.Tensor, - output_scale: torch.Tensor, - input_scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + q, + k, + v, + output_attn, + output_quant, + output_scale, + input_scale, + kv_cache_dummy_dep, + ): at1 = auto_functionalized( ATTN_OP, query=q, key=k, value=v, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=None, output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -176,42 +263,80 @@ class AttnNvfp4QuantPattern( output=output_quant, output_scale=output_scale, ) - output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE) - return at2[1], output_scale_view + return at2[1], torch.ops.aten.view.dtype(at2[2], FP8_DTYPE) return _pattern @property def replacement(self) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _replacement_with_ln( # type: ignore[misc] + q, + k, + v, + output_attn, + _output_quant, + output_scale, + input_scale, + kv_cache_dummy_dep, + layer_name, + ): + output_attn = torch.empty( + [q.shape[0], self._num_heads, self._head_size // 2], + dtype=FP4_DTYPE, + device=q.device, + ) + osv = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE) + at2 = auto_functionalized( + ATTN_OP, + query=q, + key=k, + value=v, + output=output_attn, + layer_name=layer_name, + output_scale=input_scale, + output_block_scale=osv, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + return RESHAPE_OP( + at2[1], [-1, self._num_heads * self._head_size // 2] + ), at2[2] + + return _replacement_with_ln + def _replacement( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - output_attn: torch.Tensor, - _output_quant: torch.Tensor, - output_scale: torch.Tensor, - input_scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + q, + k, + v, + output_attn, + _output_quant, + output_scale, + input_scale, + kv_cache_dummy_dep, + ): output_attn = torch.empty( [q.shape[0], self._num_heads, self._head_size // 2], dtype=FP4_DTYPE, device=q.device, ) - output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE) + osv = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE) at2 = auto_functionalized( ATTN_OP, query=q, key=k, value=v, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=input_scale, - output_block_scale=output_scale_view, + output_block_scale=osv, kv_cache_dummy_dep=kv_cache_dummy_dep, ) - output = RESHAPE_OP(at2[1], [-1, self._num_heads * self._head_size // 2]) - return output, at2[2] + return RESHAPE_OP( + at2[1], [-1, self._num_heads * self._head_size // 2] + ), at2[2] return _replacement @@ -219,18 +344,19 @@ class AttnNvfp4QuantPattern( dtype = self._dtype num_heads = self._num_heads head_size = self._head_size - return [ + inputs: list = [ self.empty_bf16(5, num_heads, head_size), # q self.empty_bf16(5, num_heads, head_size), # k self.empty_bf16(5, num_heads, head_size), # v self.empty_bf16(5, num_heads, head_size), # output_attn - self.empty(5, num_heads * head_size // 2, dtype=FP4_DTYPE), # output_quant - self.empty_i32( - 128, round_up(num_heads * head_size // 16, 4) - ), # output_scale + self.empty(5, num_heads * head_size // 2, dtype=FP4_DTYPE), + self.empty_i32(128, round_up(num_heads * head_size // 16, 4)), self.empty_fp32(1, 1), # input_scale self.empty(0, dtype=dtype), # kv_cache_dummy_dep ] + if _USE_LAYERNAME: + inputs.append(_encode_layer_name(self._layer_name)) + return inputs class AttnQuantFusionPass(VllmFusionPatternMatcherPass): @@ -259,13 +385,19 @@ class AttnQuantFusionPass(VllmFusionPatternMatcherPass): "so no fusion patterns were registered." ) + # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all + # layers produce the same pattern — register once then break. for layer in layers: if layer.impl.fused_output_quant_supported(_FP8_QUANT_KEY): self.register(AttnFp8StaticQuantPattern(layer, dtype)) + if _USE_LAYERNAME: + break if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): for layer in layers: if layer.impl.fused_output_quant_supported(kNvfp4Dynamic): self.register(AttnNvfp4QuantPattern(layer, dtype)) + if _USE_LAYERNAME: + break self.dump_patterns(config, self.pm_pass) diff --git a/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py b/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py new file mode 100644 index 00000000000..7445028da63 --- /dev/null +++ b/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Fusion pass: replace MiniMax QK allreduce + RMS norm with the Lamport +fused kernel (minimax_allreduce_rms_qk) for decode-size batches. + +Pattern (inlined forward_qk in compiled graph): + q, k, v = qkv.split([q_size, kv_size, kv_size], -1) + q_fp32 = q.to(float32); k_fp32 = k.to(float32) + q_var = q_fp32.pow(2).mean(-1, keepdim=True) + k_var = k_fp32.pow(2).mean(-1, keepdim=True) + qk_var = cat([q_var, k_var], -1) + qk_var = allreduce(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, -1) + q_out = (q_fp32 * rsqrt(q_var + eps) * q_weight).to(orig_dtype) + k_out = (k_fp32 * rsqrt(k_var + eps) * k_weight).to(orig_dtype) + return q_out, k_out, v + +Replacement (pure, no in-place on qkv/q/k): + q_out, k_out = minimax_qk_norm_fused(qkv, q_weight, k_weight, workspace, ...) + v = qkv.split([q_size, kv_size, kv_size], -1)[2] + return q_out, k_out, v + +is_applicable_for_range: only fires for compile_range.end <= max_decode_tokens +so that large prefill batches fall through to the original forward_qk (= main). +""" + +import torch +import torch._inductor.pattern_matcher as pm +import torch.fx as fx +from torch._inductor.pattern_matcher import PatternMatcherPass + +from vllm.config import VllmConfig +from vllm.config.utils import Range +from vllm.distributed import tensor_model_parallel_all_reduce +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.utils.torch_utils import direct_register_custom_op + +from ..inductor_pass import enable_fake_mode +from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass + +logger = init_logger(__name__) + +MAX_TOKEN_NUM = 2048 + +_MINIMAX_QK_NORM_FUSED_OP = None +if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"): + + def _minimax_qk_norm_fused( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + max_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.layers.mamba.lamport_workspace import ( + get_allreduce_workspace, + ) + + workspace = get_allreduce_workspace( + rank=rank, + world_size=nranks, + max_tokens=max_tokens, + process_group=get_tp_group().cpu_group, + ) + return torch.ops._C.minimax_allreduce_rms_qk( + qkv, + norm_weight_q, + norm_weight_k, + workspace, + q_size, + kv_size, + rank, + nranks, + eps, + ) + + def _minimax_qk_norm_fused_fake( + qkv: torch.Tensor, + norm_weight_q: torch.Tensor, + norm_weight_k: torch.Tensor, + q_size: int, + kv_size: int, + rank: int, + nranks: int, + eps: float, + max_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + T = qkv.shape[0] + return ( + torch.empty([T, q_size], dtype=qkv.dtype, device=qkv.device), + torch.empty([T, kv_size], dtype=qkv.dtype, device=qkv.device), + ) + + direct_register_custom_op( + op_name="minimax_qk_norm_fused", + op_func=_minimax_qk_norm_fused, + fake_impl=_minimax_qk_norm_fused_fake, + mutates_args=[], + ) + _MINIMAX_QK_NORM_FUSED_OP = torch.ops.vllm.minimax_qk_norm_fused.default + + +class MiniMaxQKNormPattern: + """ + Match the forward_qk allreduce+rms pattern and replace with Lamport kernel. + """ + + def __init__( + self, + q_size: int, + kv_size: int, + eps: float, + tp_world: int, + tp_rank: int, + max_tokens: int, + dtype: torch.dtype, + device: str | None, + ) -> None: + self.q_size = q_size + self.kv_size = kv_size + self.eps = eps + self.tp_world = tp_world + self.tp_rank = tp_rank + self.max_tokens = max_tokens + self.dtype = dtype + self.device = device + + def get_inputs(self) -> list[torch.Tensor]: + T = 4 + qkv = torch.empty( + [T, self.q_size + 2 * self.kv_size], + device=self.device, + dtype=self.dtype, + ) + q_weight = torch.empty([self.q_size], device=self.device, dtype=self.dtype) + k_weight = torch.empty([self.kv_size], device=self.device, dtype=self.dtype) + return [qkv, q_weight, k_weight] + + def register(self, pm_pass: PatternMatcherPass) -> None: + q_size = self.q_size + kv_size = self.kv_size + eps = self.eps + tp_world = self.tp_world + max_tokens = self.max_tokens + tp_rank = self.tp_rank + dtype = self.dtype + + def pattern( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + q_fp32 = q.to(torch.float32) + k_fp32 = k.to(torch.float32) + q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True) + k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True) + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, dim=-1) + q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype) + k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype) + return q_out, k_out, v + + def replacement( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + assert _MINIMAX_QK_NORM_FUSED_OP is not None + q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused( + qkv, + q_weight, + k_weight, + q_size, + kv_size, + tp_rank, + tp_world, + eps, + max_tokens, + ) + _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + return q_out, k_out, v + + pm.register_replacement( + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass + ) + + # Second pattern: three separate split_with_sizes nodes (one per output), + # each with _users=1. This occurs when the QKV projection uses a + # functional GEMM kernel (e.g. cutlass_scaled_mm via auto_functionalized), + # which causes inductor to generate one split per consumer. + def pattern_split3( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0] + k = qkv.split([q_size, kv_size, kv_size], dim=-1)[1] + v = qkv.split([q_size, kv_size, kv_size], dim=-1)[2] + q_fp32 = q.to(torch.float32) + k_fp32 = k.to(torch.float32) + q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True) + k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True) + qk_var = torch.cat([q_var, k_var], dim=-1) + qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world + q_var, k_var = qk_var.chunk(2, dim=-1) + q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype) + k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype) + return q_out, k_out, v + + pm.register_replacement( + pattern_split3, replacement, self.get_inputs(), pm.fwd_only, pm_pass + ) + + +class MiniMaxQKNormPass(VllmPatternMatcherPass): + """ + Replace forward_qk allreduce+norm with the Lamport fused kernel. + Only applied for decode-size compile ranges (small token counts). + """ + + def __init__(self, config: VllmConfig) -> None: + super().__init__(config) + self.disabled = True + + if _MINIMAX_QK_NORM_FUSED_OP is None: + logger.warning_once( + "minimax_allreduce_rms_qk op not found, MiniMaxQKNormPass disabled." + ) + return + + tp_world = get_tensor_model_parallel_world_size() + if tp_world <= 1: + logger.warning_once("MiniMaxQKNormPass disabled: tp_size <= 1.") + return + + if config.model_config is None: + logger.warning_once("MiniMaxQKNormPass disabled: no model_config.") + return + + hf_cfg = config.model_config.hf_config + + model_name = getattr(hf_cfg, "architectures", "")[0] + if model_name != "MiniMaxM2ForCausalLM": + return + + num_attention_heads = getattr(hf_cfg, "num_attention_heads", 0) + num_key_value_heads = getattr(hf_cfg, "num_key_value_heads", 0) + hidden_size = getattr(hf_cfg, "hidden_size", 0) + head_dim = getattr(hf_cfg, "head_dim", 0) + eps: float = getattr(hf_cfg, "rms_norm_eps", 1e-6) + + if ( + num_attention_heads != 48 + or num_key_value_heads != 8 + or hidden_size != 3072 + or head_dim != 128 + ): + logger.warning_once( + "MiniMaxQKNormPass disabled: cannot infer model info from hf_config." + ) + return + + num_heads_per_rank = num_attention_heads // tp_world + num_kv_heads_per_rank = max(1, num_key_value_heads // tp_world) + q_size = num_heads_per_rank * head_dim + kv_size = num_kv_heads_per_rank * head_dim + + self.max_token_num = min( + MAX_TOKEN_NUM, config.scheduler_config.max_num_batched_tokens + ) + + tp_rank = get_tensor_model_parallel_rank() + # Allocate Lamport workspace first. + from vllm.distributed.parallel_state import get_tp_group + from vllm.model_executor.layers.mamba.lamport_workspace import ( + get_allreduce_workspace, + ) + + get_allreduce_workspace( + rank=tp_rank, + world_size=tp_world, + max_tokens=self.max_token_num, + process_group=get_tp_group().cpu_group, + ) + + self.patterns: PatternMatcherPass = PatternMatcherPass( + pass_name="minimax_qk_norm_pass" + ) + self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank) + self.dump_patterns(config, self.patterns) + self.disabled = False + + @enable_fake_mode + def _register_patterns( + self, + q_size: int, + kv_size: int, + eps: float, + tp_world: int, + tp_rank: int, + ) -> None: + MiniMaxQKNormPattern( + q_size=q_size, + kv_size=kv_size, + eps=eps, + tp_world=tp_world, + tp_rank=tp_rank, + max_tokens=self.max_token_num, + dtype=self.model_dtype, + device=self.device, + ).register(self.patterns) + + def is_applicable_for_range(self, compile_range: Range) -> bool: + if self.disabled: + return False + + return bool(compile_range.end <= self.max_token_num) + + @VllmInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + if self.disabled: + return + self.matched_count = self.patterns.apply(graph) + logger.debug("MiniMaxQKNormPass replaced %s patterns", self.matched_count) + + def uuid(self) -> str: + return VllmInductorPass.hash_source(self, MiniMaxQKNormPattern) diff --git a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py index 5a9ef46a0fc..e36400ec8ec 100644 --- a/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py +++ b/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py @@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kNvfp4Dynamic, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement from .matcher_utils import MatcherQuantFP8 @@ -49,21 +50,43 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): @property def pattern(self) -> Callable[..., torch.Tensor]: - def _pattern( - q: torch.Tensor, - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - output_attn: torch.Tensor, - scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> torch.Tensor: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _pattern_with_ln( # type: ignore[misc] + q, + kv_c_normed, + k_pe, + output_attn, + scale, + kv_cache_dummy_dep, + layer_name, + ): + at1 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=layer_name, + output_scale=None, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + # MLA output is already 2D (T, N*V), no reshape needed + return self._quant_matcher(at1[1], scale)[0] + + return _pattern_with_ln + + def _pattern(q, kv_c_normed, k_pe, output_attn, scale, kv_cache_dummy_dep): at1 = auto_functionalized( MLA_ATTN_OP, q=q, kv_c_normed=kv_c_normed, k_pe=k_pe, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=None, output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -75,14 +98,41 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): @property def replacement(self) -> Callable[..., torch.Tensor]: - def _replacement( - q: torch.Tensor, - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - output_attn: torch.Tensor, - scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> torch.Tensor: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _replacement_with_ln( # type: ignore[misc] + q, + kv_c_normed, + k_pe, + output_attn, + scale, + kv_cache_dummy_dep, + layer_name, + ): + # MLA output in quant_dtype + output_attn = torch.empty( + [q.shape[0], self._output_dim], + dtype=FP8_DTYPE, + device=q.device, + ) + at1 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=layer_name, + output_scale=scale, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + return at1[1] + + return _replacement_with_ln + + def _replacement(q, kv_c_normed, k_pe, output_attn, scale, kv_cache_dummy_dep): # MLA output in quant_dtype output_attn = torch.empty( [q.shape[0], self._output_dim], @@ -95,7 +145,7 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): kv_c_normed=kv_c_normed, k_pe=k_pe, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=scale, output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -105,7 +155,7 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): return _replacement def get_inputs(self) -> list[torch.Tensor]: - return [ + inputs: list = [ self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype), self.empty(5, self._kv_lora_rank, dtype=self._dtype), self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype), @@ -113,6 +163,9 @@ class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]): self.empty_fp32(1, 1), self.empty(0, dtype=self._dtype), ] + if _USE_LAYERNAME: + inputs.append(_encode_layer_name(self._layer_name)) + return inputs class MLAAttnNvfp4QuantPattern( @@ -141,21 +194,56 @@ class MLAAttnNvfp4QuantPattern( def pattern( self, ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _pattern_with_ln( # type: ignore[misc] + q, + kv_c_normed, + k_pe, + output_attn, + input_scale, + kv_cache_dummy_dep, + layer_name, + ): + at1 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=layer_name, + output_scale=None, + output_block_scale=None, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + output_quant, output_scale = create_fp4_output_tensors( + at1[1].shape[0], at1[1].shape[1], at1[1].device, True + ) + at2 = auto_functionalized( + self._QUANT_OP, + input=at1[1], + input_scale=input_scale, + is_sf_swizzled_layout=True, + output=output_quant, + output_scale=output_scale, + ) + output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE) + return at2[1], output_scale_view + + return _pattern_with_ln + def _pattern( - q: torch.Tensor, - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - output_attn: torch.Tensor, - input_scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + q, kv_c_normed, k_pe, output_attn, input_scale, kv_cache_dummy_dep + ): at1 = auto_functionalized( MLA_ATTN_OP, q=q, kv_c_normed=kv_c_normed, k_pe=k_pe, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=None, output_block_scale=None, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -182,14 +270,47 @@ class MLAAttnNvfp4QuantPattern( def replacement( self, ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]: + _ln = _encode_layer_name(self._layer_name) + + if _USE_LAYERNAME: + + def _replacement_with_ln( # type: ignore[misc] + q, + kv_c_normed, + k_pe, + output_attn, + input_scale, + kv_cache_dummy_dep, + layer_name, + ): + # MLA output in quant_dtype (FP4 packed as uint8) + output_attn = torch.empty( + [q.shape[0], self._output_dim // 2], + dtype=FP4_DTYPE, + device=q.device, + ) + output_scale = create_fp4_output_tensors( + q.shape[0], self._output_dim, q.device, True + )[1] + output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE) + at2 = auto_functionalized( + MLA_ATTN_OP, + q=q, + kv_c_normed=kv_c_normed, + k_pe=k_pe, + output=output_attn, + layer_name=layer_name, + output_scale=input_scale, + output_block_scale=output_scale_view, + kv_cache_dummy_dep=kv_cache_dummy_dep, + ) + return at2[1], at2[2] + + return _replacement_with_ln + def _replacement( - q: torch.Tensor, - kv_c_normed: torch.Tensor, - k_pe: torch.Tensor, - output_attn: torch.Tensor, - input_scale: torch.Tensor, - kv_cache_dummy_dep: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + q, kv_c_normed, k_pe, output_attn, input_scale, kv_cache_dummy_dep + ): # MLA output in quant_dtype (FP4 packed as uint8) output_attn = torch.empty( [q.shape[0], self._output_dim // 2], @@ -207,7 +328,7 @@ class MLAAttnNvfp4QuantPattern( kv_c_normed=kv_c_normed, k_pe=k_pe, output=output_attn, - layer_name=self._layer_name, + layer_name=_ln, output_scale=input_scale, output_block_scale=output_scale_view, kv_cache_dummy_dep=kv_cache_dummy_dep, @@ -217,7 +338,7 @@ class MLAAttnNvfp4QuantPattern( return _replacement def get_inputs(self) -> list[torch.Tensor]: - return [ + inputs: list = [ self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype), self.empty(5, self._kv_lora_rank, dtype=self._dtype), self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype), @@ -225,6 +346,9 @@ class MLAAttnNvfp4QuantPattern( self.empty_fp32(1, 1), self.empty(0, dtype=self._dtype), ] + if _USE_LAYERNAME: + inputs.append(_encode_layer_name(self._layer_name)) + return inputs class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass): @@ -250,13 +374,19 @@ class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass): "so no fusion patterns were registered." ) + # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all + # layers produce the same pattern — register once then break. for layer in layers: if layer.impl.fused_output_quant_supported(kFp8StaticTensorSym): self.register(MLAAttnFp8StaticQuantPattern(layer, dtype)) + if _USE_LAYERNAME: + break if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): for layer in layers: if layer.impl.fused_output_quant_supported(kNvfp4Dynamic): self.register(MLAAttnNvfp4QuantPattern(layer, dtype)) + if _USE_LAYERNAME: + break self.dump_patterns(config, self.pm_pass) diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py index 245119fa5e0..b7e747a784e 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py @@ -164,6 +164,7 @@ class QkNormRopePattern: cos_sin_cache=cos_sin_cache, is_neox=self.is_neox, position_ids=positions.view(-1), + forced_token_heads_per_warp=-1, ) result_qkv = result[1] diff --git a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py index 830a9640780..bc6754188aa 100644 --- a/vllm/compilation/passes/fusion/rope_kvcache_fusion.py +++ b/vllm/compilation/passes/fusion/rope_kvcache_fusion.py @@ -15,7 +15,13 @@ from vllm.model_executor.layers.attention.attention import ( Attention, get_attention_context, ) -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + _USE_LAYERNAME, + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) from ..inductor_pass import enable_fake_mode from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass @@ -37,7 +43,7 @@ def fused_rope_and_unified_kv_cache_update_impl( positions: torch.Tensor, cos_sin_cache: torch.Tensor, is_neox: bool, - layer_name: str = "", + layer_name: LayerNameType, ) -> torch.Tensor: """ This impl fetches the KV cache and slot mapping from the forward context, @@ -46,6 +52,7 @@ def fused_rope_and_unified_kv_cache_update_impl( that is passed to unified_attention to signal a side effect and the data dependency between them to ensure torch.compile preserves ordering. """ + layer_name = _resolve_layer_name(layer_name) _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name) if layer_slot_mapping is not None: attn_layer.impl.do_rope_and_kv_cache_update( @@ -70,7 +77,7 @@ def fused_rope_and_unified_kv_cache_update_fake( positions: torch.Tensor, cos_sin_cache: torch.Tensor, is_neox: bool, - layer_name: str = "", + layer_name: LayerNameType, ) -> torch.Tensor: return torch.empty(0, device=query.device, dtype=query.dtype) @@ -120,38 +127,30 @@ class RopeReshapeKVCachePattern: num_kv_heads=self.num_kv_heads, ) - def get_inputs(self) -> list[torch.Tensor]: + def get_inputs(self) -> list: # Sample inputs to help pattern tracing T = 5 L = 4096 qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size) positions = empty_i64(T) cos_sin_cache = empty_bf16(L, self.head_size) - return [ - qkv, - positions, - cos_sin_cache, - ] + inputs: list = [qkv, positions, cos_sin_cache] + if _USE_LAYERNAME: + inputs.append(_encode_layer_name(self.layer_name)) + return inputs - def register(self, pm_pass: PatternMatcherPass) -> None: - def pattern( - qkv: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + def _mk_pattern_with_layer_name_input(self, _ln): + """Pattern/replacement with layer_name as an explicit input.""" + + def pattern(qkv, positions, cos_sin_cache, layer_name): q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) q, k = self.rope_matcher(positions, q, k, cos_sin_cache) q = q.view(-1, self.num_heads, self.head_size) k = k.view(-1, self.num_kv_heads, self.head_size) v = v.view(-1, self.num_kv_heads, self.head_size_v) - dummy = torch.ops.vllm.unified_kv_cache_update(k, v, self.layer_name) - return dummy, q, k, v + return torch.ops.vllm.unified_kv_cache_update(k, v, layer_name), q, k, v - def replacement( - qkv: torch.Tensor, - positions: torch.Tensor, - cos_sin_cache: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + def replacement(qkv, positions, cos_sin_cache, layer_name): q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) q = q.view(-1, self.num_heads, self.head_size) k = k.view(-1, self.num_kv_heads, self.head_size) @@ -164,10 +163,50 @@ class RopeReshapeKVCachePattern: positions=positions, cos_sin_cache=cos_sin_cache, is_neox=self.is_neox, - layer_name=self.layer_name, + layer_name=layer_name, ) return results[0], results[1], results[2], v + return pattern, replacement + + def _mk_pattern_with_layer_name_closure(self, _ln): + """Pattern/replacement with layer_name as a closure constant.""" + + def pattern(qkv, positions, cos_sin_cache): + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + q, k = self.rope_matcher(positions, q, k, cos_sin_cache) + q = q.view(-1, self.num_heads, self.head_size) + k = k.view(-1, self.num_kv_heads, self.head_size) + v = v.view(-1, self.num_kv_heads, self.head_size_v) + return torch.ops.vllm.unified_kv_cache_update(k, v, _ln), q, k, v + + def replacement(qkv, positions, cos_sin_cache): + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + q = q.view(-1, self.num_heads, self.head_size) + k = k.view(-1, self.num_kv_heads, self.head_size) + v = v.view(-1, self.num_kv_heads, self.head_size_v) + results = auto_functionalized( + self.FUSED_OP, + query=q, + key=k, + value=v, + positions=positions, + cos_sin_cache=cos_sin_cache, + is_neox=self.is_neox, + layer_name=_ln, + ) + return results[0], results[1], results[2], v + + return pattern, replacement + + def register(self, pm_pass: PatternMatcherPass) -> None: + _ln = _encode_layer_name(self.layer_name) + + if _USE_LAYERNAME: + pattern, replacement = self._mk_pattern_with_layer_name_input(_ln) + else: + pattern, replacement = self._mk_pattern_with_layer_name_closure(_ln) + # NOTE: use view_to_reshape to unify view/reshape to simplify # pattern and increase matching opportunities def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule: @@ -176,7 +215,11 @@ class RopeReshapeKVCachePattern: return gm pm.register_replacement( - pattern, replacement, self.get_inputs(), fwd_and_view_to_reshape, pm_pass + pattern, + replacement, + self.get_inputs(), + fwd_and_view_to_reshape, + pm_pass, ) @@ -205,6 +248,8 @@ class RopeKVCacheFusionPass(VllmPatternMatcherPass): self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num attn_layers = get_layers_from_vllm_config(config, Attention) + # When _USE_LAYERNAME is enabled, layer_name is a wildcard so all + # layers produce the same pattern — register once then break. for _, layer in attn_layers.items(): if layer.impl.fused_rope_kvcache_supported(): for is_neox in [True, False]: @@ -212,6 +257,8 @@ class RopeKVCacheFusionPass(VllmPatternMatcherPass): layer=layer, is_neox=is_neox, ).register(self.patterns) + if _USE_LAYERNAME: + break self.dump_patterns(config, self.patterns) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index b4823a0afde..91e10145607 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -38,6 +38,7 @@ if current_platform.is_cuda_alike(): if current_platform.is_cuda(): from .fusion.allreduce_rms_fusion import AllReduceFusionPass from .fusion.collective_fusion import AsyncTPPass + from .fusion.minimax_qk_norm_fusion import MiniMaxQKNormPass from .inductor_pass import ( CustomGraphPass, @@ -137,6 +138,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc] if self.pass_config.fuse_allreduce_rms: self.passes += [AllReduceFusionPass(config)] + if self.pass_config.fuse_minimax_qk_norm: + self.passes += [MiniMaxQKNormPass(config)] + if self.pass_config.fuse_norm_quant: self.passes += [RMSNormQuantFusionPass(config)] if rocm_aiter_ops.is_enabled(): diff --git a/vllm/compilation/passes/utility/fix_functionalization.py b/vllm/compilation/passes/utility/fix_functionalization.py index 1b656d0c890..15eb23e6f94 100644 --- a/vllm/compilation/passes/utility/fix_functionalization.py +++ b/vllm/compilation/passes/utility/fix_functionalization.py @@ -168,6 +168,7 @@ class FixFunctionalizationPass(VllmInductorPass): "cos_sin_cache", "is_neox", "position_ids", + "forced_token_heads_per_warp", ) self.defunctionalize(graph, node, mutated_args=mutated_args, args=args) elif ( diff --git a/vllm/compilation/piecewise_backend.py b/vllm/compilation/piecewise_backend.py index 7474d0bf841..02a4dad5460 100644 --- a/vllm/compilation/piecewise_backend.py +++ b/vllm/compilation/piecewise_backend.py @@ -270,6 +270,7 @@ class PiecewiseBackend: compile_range=range_entry.compile_range, graph_index=self.piecewise_compile_index, num_graphs=self.total_piecewise_compiles, + is_encoder=self.vllm_backend.is_encoder, ) range_entry.compiled = True @@ -353,12 +354,22 @@ class PiecewiseBackend: return None def __call__(self, *args: Any) -> Any: - runtime_shape = args[self.sym_shape_indices[0]] - range_entry = self._find_range_for_shape(runtime_shape) + if self.sym_shape_indices: + runtime_shape = args[self.sym_shape_indices[0]] + range_entry = self._find_range_for_shape(runtime_shape) + assert range_entry is not None, ( + f"Shape: {runtime_shape} out of considered ranges: " + f"{self.compile_ranges}" + ) + else: + # All inputs have static shapes; use the only compiled range_entry + compiled_entries = [re for re in self.range_entries.values() if re.compiled] + assert len(compiled_entries) == 1, ( + f"Expected exactly one compiled range_entry for static shape " + f"compilation, but found {len(compiled_entries)}" + ) + range_entry = compiled_entries[0] - assert range_entry is not None, ( - f"Shape: {runtime_shape} out of considered ranges: {self.compile_ranges}" - ) assert range_entry.compiled, ( "All ranges should be compiled or loaded up front in " "PiecewiseBackend.__init__. " diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index d5a3e9bfd96..758605d25c6 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -16,6 +16,7 @@ from vllm.config.kv_events import KVEventsConfig from vllm.config.kv_transfer import KVTransferConfig from vllm.config.load import LoadConfig from vllm.config.lora import LoRAConfig +from vllm.config.mamba import MambaConfig from vllm.config.model import ( ModelConfig, iter_architecture_defaults, @@ -83,6 +84,8 @@ __all__ = [ "LoadConfig", # From vllm.config.lora "LoRAConfig", + # From vllm.config.mamba + "MambaConfig", # From vllm.config.model "ModelConfig", "iter_architecture_defaults", diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 1da647a6d6f..561367173d5 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -27,6 +27,11 @@ class AttentionConfig: flash_attn_max_num_splits_for_cuda_graph: int = 32 """Flash Attention max number splits for cuda graph decode.""" + tq_max_kv_splits_for_cuda_graph: int = 32 + """TurboQuant max NUM_KV_SPLITS for cuda graph decode. + Fixes the split count so grid dimensions are constant across captures, + and buffers can be pre-allocated to avoid inflating the memory estimate.""" + use_cudnn_prefill: bool = False """Whether to use cudnn prefill.""" diff --git a/vllm/config/cache.py b/vllm/config/cache.py index cd1554590ea..47a655f22d5 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -24,6 +24,10 @@ CacheDType = Literal[ "fp8_e5m2", "fp8_inc", "fp8_ds_mla", + "turboquant_k8v4", + "turboquant_4bit_nc", + "turboquant_k3v4_nc", + "turboquant_3bit_nc", "int8_per_token_head", "fp8_per_token_head", ] @@ -123,14 +127,6 @@ class CacheConfig: - "align": only cache the mamba state of the last token of each scheduler step and when the token is at position i * block_size. """ - enable_mamba_cache_stochastic_rounding: bool = False - """Enable stochastic rounding when writing SSM state to fp16 cache. - Uses random bits to unbias the rounding error, which can improve - numerical stability for long sequences.""" - mamba_cache_philox_rounds: int = 0 - """Number of Philox PRNG rounds for stochastic rounding random number - generation. 0 uses the Triton default. Higher values improve randomness - quality at the cost of compute.""" # Will be set after profiling. num_gpu_blocks: int | None = field(default=None, init=False) @@ -258,29 +254,3 @@ class CacheConfig: str(cache_dtype), ) return cache_dtype - - def __post_init__(self): - if self.enable_mamba_cache_stochastic_rounding: - from vllm.platforms import current_platform - - if not current_platform.is_cuda(): - raise ValueError( - "Stochastic rounding for Mamba cache is only supported " - "on NVIDIA CUDA platforms. Please do not specify " - "`--enable-mamba-cache-stochastic-rounding`." - ) - if not current_platform.is_device_capability_family(100): - raise ValueError( - "Stochastic rounding for Mamba cache requires compute " - "capability 10.0 (data center Blackwell). The `cvt.rs` PTX " - "instruction is not supported on your GPU. Please do not specify " - "`--enable-mamba-cache-stochastic-rounding`." - ) - if self.mamba_ssm_cache_dtype != "float16": - raise ValueError( - "Stochastic rounding for Mamba cache requires " - "the SSM cache to be float16. Please set it explicitly, " - "by specifying `--mamba-ssm-cache-dtype float16`, or disable " - "stochastic rounding by not specifying " - "`--enable-mamba-cache-stochastic-rounding`." - ) diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index ef2a4bf5a4f..6aca5c9825f 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -26,6 +26,8 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backend import AttentionCGSupport + from vllm.v1.kv_cache_interface import KVCacheConfig else: VllmConfig = object @@ -132,6 +134,8 @@ class PassConfig: """Enable async TP.""" fuse_allreduce_rms: bool = None # type: ignore[assignment] """Enable flashinfer allreduce fusion.""" + fuse_minimax_qk_norm: bool = None # type: ignore[assignment] + """Enable fused allreduce+RMSNorm for MiniMax QK norm.""" enable_qk_norm_rope_fusion: bool = False """Enable fused Q/K RMSNorm + RoPE pass.""" @@ -515,13 +519,21 @@ class CompilationConfig: User-provided values override auto-inference. Example: [2048, 4096, 8192, 13824]""" - encoder_cudagraph_max_images_per_batch: int = 0 - """Maximum number of images per batch for encoder CUDA graph capture. + encoder_cudagraph_max_vision_items_per_batch: int = 0 + """Maximum number of images/videos per batch for encoder CUDA graph capture. Determines the fixed batch size used during graph capture. If 0 (default), auto-inferred as max_budget // min_budget from the model's budget range. User-provided positive value overrides auto-inference.""" + encoder_cudagraph_max_frames_per_batch: int = 0 + """Maximum total video frames per batch for encoder CUDA graph capture. + Controls the cu_seqlens buffer size (one entry per attention sequence, + i.e. one per video frame). If 0 (default), auto-inferred per budget + level as token_budget (tight bound: packing guarantees + sum(T_i) <= token_budget). Positive value overrides auto-inference + and applies to all budget levels.""" + # Inductor capture compile_sizes: list[int | str] | None = None """Sizes to compile for inductor. In addition @@ -698,6 +710,8 @@ class CompilationConfig: """files that are traced for compilation""" compilation_time: float = field(default=0.0, init=False) """time taken for compilation""" + encoder_compilation_time: float = field(default=0.0, init=False) + """time taken for multimodal encoder compilation""" static_forward_context: dict[str, Any] = field(default_factory=dict, init=False) """Per-model forward context @@ -744,6 +758,7 @@ class CompilationConfig: "local_cache_dir", "traced_files", "compilation_time", + "encoder_compilation_time", "static_forward_context", "pass_config", # handled separately below "dynamic_shapes_config", # handled separately below @@ -763,6 +778,7 @@ class CompilationConfig: "enabled_custom_ops": True, "disabled_custom_ops": True, "compilation_time": True, + "encoder_compilation_time": True, "traced_files": True, "inductor_compile_config": { "post_grad_custom_post_pass": True, @@ -960,10 +976,18 @@ class CompilationConfig: # Validate encoder CUDA graph configuration if ( self.cudagraph_mm_encoder - and self.encoder_cudagraph_max_images_per_batch < 0 + and self.encoder_cudagraph_max_vision_items_per_batch < 0 ): raise ValueError( - "encoder_cudagraph_max_images_per_batch must be " + "encoder_cudagraph_max_vision_items_per_batch must be " + "non-negative (0 = auto-infer)" + ) + if ( + self.cudagraph_mm_encoder + and self.encoder_cudagraph_max_frames_per_batch < 0 + ): + raise ValueError( + "encoder_cudagraph_max_frames_per_batch must be " "non-negative (0 = auto-infer)" ) @@ -1241,6 +1265,152 @@ class CompilationConfig: assert "none" in self.custom_ops return f"+{op}" in self.custom_ops + def resolve_cudagraph_mode_and_sizes( + self, + min_cg_support: "AttentionCGSupport", + min_cg_attn_backend: str | None, + uniform_decode_query_len: int = 1, + tensor_parallel_size: int = 1, + kv_cache_config: "KVCacheConfig | None" = None, + max_num_reqs: int | None = None, + is_profiling: bool = False, + ) -> CUDAGraphMode: + from vllm.v1.attention.backend import AttentionCGSupport + + cudagraph_mode = self.cudagraph_mode + if cudagraph_mode is None or cudagraph_mode == CUDAGraphMode.NONE: + self.cudagraph_mode = CUDAGraphMode.NONE + return CUDAGraphMode.NONE + + # Check cudagraph for mixed batch is supported + if ( + cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL + and min_cg_support != AttentionCGSupport.ALWAYS + ): + msg = ( + f"CUDAGraphMode.{cudagraph_mode.name} is not supported " + f"with {min_cg_attn_backend} backend (support: " + f"{min_cg_support})" + ) + if min_cg_support == AttentionCGSupport.NEVER: + # if not supported any full cudagraphs, just raise it. + msg += ( + "; please try cudagraph_mode=PIECEWISE, and " + "make sure compilation mode is VLLM_COMPILE" + ) + raise ValueError(msg) + + # attempt to resolve the full cudagraph related mode + if self.splitting_ops_contain_attention(): + msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE" + cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + else: + msg += "; setting cudagraph_mode=FULL_DECODE_ONLY" + cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + logger.warning(msg) + + # check that if we are doing decode full-cudagraphs it is supported + if ( + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and min_cg_support == AttentionCGSupport.NEVER + ): + msg = ( + f"CUDAGraphMode.{cudagraph_mode.name} is not supported " + f"with {min_cg_attn_backend} backend (support: " + f"{min_cg_support})" + ) + if self.mode == CompilationMode.VLLM_COMPILE and ( + self.splitting_ops_contain_attention() + or self.use_inductor_graph_partition + ): + msg += ( + "; setting cudagraph_mode=PIECEWISE because " + "attention is compiled piecewise" + ) + cudagraph_mode = CUDAGraphMode.PIECEWISE + else: + msg += ( + "; setting cudagraph_mode=NONE because " + "attention is not compiled piecewise" + ) + cudagraph_mode = CUDAGraphMode.NONE + logger.warning(msg) + + # check that if we are doing spec-decode + decode full-cudagraphs it is + # supported + if ( + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and uniform_decode_query_len > 1 + and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value + ): + msg = ( + f"CUDAGraphMode.{cudagraph_mode.name} is not supported" + f" with spec-decode for attention backend " + f"{min_cg_attn_backend} (support: {min_cg_support})" + ) + if self.splitting_ops_contain_attention(): + msg += "; setting cudagraph_mode=PIECEWISE" + cudagraph_mode = CUDAGraphMode.PIECEWISE + else: + msg += "; setting cudagraph_mode=NONE" + cudagraph_mode = CUDAGraphMode.NONE + logger.warning(msg) + + # double check that we can support full cudagraph if they are requested + # even after automatic downgrades + if ( + cudagraph_mode.has_full_cudagraphs() + and min_cg_support == AttentionCGSupport.NEVER + ): + raise ValueError( + f"CUDAGraphMode.{cudagraph_mode.name} is not " + f"supported with {min_cg_attn_backend} backend (" + f"support:{min_cg_support}) " + "; please try cudagraph_mode=PIECEWISE, " + "and make sure compilation mode is VLLM_COMPILE" + ) + + # Adjust cudagraph sizes to be a multiple of uniform_decode_query_len + # to avoid: https://github.com/vllm-project/vllm/issues/28207 and temp-fix: + # https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536 + # Will be removed in the near future when we have separate cudagraph capture + # sizes for decode and mixed prefill-decode. + if ( + cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + and uniform_decode_query_len > 1 + ): + self.adjust_cudagraph_sizes_for_spec_decode( + uniform_decode_query_len, + tensor_parallel_size, + ) + + # For Mamba models with FULL decode cudagraphs, each decode + # sequence needs one Mamba cache block. The decode cudagraph + # dispatcher already caps batch sizes at max_num_seqs, so we just + # need to verify that enough blocks exist. Raising here instead + # of silently capping cudagraph_capture_sizes avoids unintended + # restrictions on PIECEWISE (prefill) cudagraphs. + # See: https://github.com/vllm-project/vllm/issues/34094 + if ( + kv_cache_config is not None + and max_num_reqs is not None + and cudagraph_mode.has_full_cudagraphs() + and not is_profiling + and kv_cache_config.has_mamba_layers + and max_num_reqs > kv_cache_config.num_blocks + ): + raise ValueError( + f"max_num_seqs ({max_num_reqs}) exceeds available Mamba cache " + f"blocks ({kv_cache_config.num_blocks}). Each decode sequence " + "requires one Mamba cache block, so CUDA graph capture cannot " + "proceed. Please lower max_num_seqs to at most " + f"{kv_cache_config.num_blocks} or increase " + "gpu_memory_utilization." + ) + + self.cudagraph_mode = cudagraph_mode + return cudagraph_mode + def adjust_cudagraph_sizes_for_spec_decode( self, uniform_decode_query_len: int, tensor_parallel_size: int ): diff --git a/vllm/config/lora.py b/vllm/config/lora.py index bfef0efa3df..bf47887c7e0 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -7,8 +7,10 @@ import torch from pydantic import ConfigDict, Field, model_validator from typing_extensions import Self +from vllm import envs from vllm.config.utils import config from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.hashing import safe_hash if TYPE_CHECKING: @@ -105,7 +107,14 @@ class LoRAConfig: f"max_cpu_loras ({self.max_cpu_loras}) must be >= " f"max_loras ({self.max_loras})." ) - + if envs.VLLM_LORA_ENABLE_DUAL_STREAM and not current_platform.is_cuda_alike(): + raise ValueError("Dual CUDA streams are only supported on CUDA platforms.") + if envs.VLLM_LORA_ENABLE_DUAL_STREAM and self.fully_sharded_loras: + logger.warning_once( + "fully_sharded_loras isn't compatible with " + "VLLM_LORA_ENABLE_DUAL_STREAM, set VLLM_LORA_ENABLE_DUAL_STREAM=False" + ) + envs.VLLM_LORA_ENABLE_DUAL_STREAM = False return self def verify_with_model_config(self, model_config: ModelConfig): diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py new file mode 100644 index 00000000000..996478c3676 --- /dev/null +++ b/vllm/config/mamba.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import Enum, EnumMeta +from typing import Any + +from pydantic import field_validator + +from vllm.config.utils import config + + +class _MambaBackendEnumMeta(EnumMeta): + """Metaclass for MambaBackendEnum to provide better error messages.""" + + def __getitem__(cls, name: str): + try: + return super().__getitem__(name) + except KeyError: + valid = ", ".join(cls.__members__.keys()) + raise ValueError( + f"Unknown Mamba SSU backend: '{name}'. Valid options are: {valid}" + ) from None + + +class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): + """Enumeration of supported Mamba SSU (selective state update) backends.""" + + TRITON = "triton" + FLASHINFER = "flashinfer" + + +@config +class MambaConfig: + """Configuration for Mamba SSM backends.""" + + backend: MambaBackendEnum = MambaBackendEnum.TRITON + """Mamba SSU backend to use.""" + + enable_stochastic_rounding: bool = False + """Enable stochastic rounding when writing SSM state to fp16 cache. + Uses random bits to unbias the rounding error, which can improve + numerical stability for long sequences.""" + stochastic_rounding_philox_rounds: int = 0 + """Number of Philox PRNG rounds for stochastic rounding random number + generation. 0 uses the Triton default. Higher values improve randomness + quality at the cost of compute.""" + + @field_validator("backend", mode="before") + @classmethod + def validate_backend_before(cls, value: Any) -> Any: + """Enable parsing of the `backend` enum type from string.""" + if isinstance(value, str): + return MambaBackendEnum[value.upper()] + return value + + def __post_init__(self): + if self.enable_stochastic_rounding: + from vllm.platforms import current_platform + + if not current_platform.is_cuda(): + raise ValueError( + "Stochastic rounding for Mamba cache is only supported " + "on NVIDIA CUDA platforms. Please do not specify " + "`--enable-mamba-cache-stochastic-rounding`." + ) + if ( + self.backend == MambaBackendEnum.TRITON + and not current_platform.is_device_capability_family(100) + ): + raise ValueError( + "Stochastic rounding for Mamba cache with triton backend requires " + "compute capability 10.0 (data center Blackwell). The `cvt.rs` " + "PTX instruction is not supported on your GPU. Please do not " + "specify `--enable-mamba-cache-stochastic-rounding`, " + "or set `--mamba-backend flashinfer`." + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index 1cce7f9d94c..2b767b21a7c 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -951,6 +951,7 @@ class ModelConfig: # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", + "gpt_oss_mxfp4", "cpu_awq", "gguf", ] @@ -966,7 +967,7 @@ class ModelConfig: for name in quantization_methods: method = me_quant.get_quantization_config(name) quantization_override = method.override_quantization_method( - quant_cfg, self.quantization + quant_cfg, self.quantization, hf_config=self.hf_config ) if quantization_override is not None: # Raise error if the override is not custom (custom would diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 0b5c97ba063..a42b8422ef3 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -622,6 +622,18 @@ class ParallelConfig: and self.data_parallel_size > 1 ) + @property + def use_batched_dp_moe(self) -> bool: + return ( + self.all2all_backend + in ( + "deepep_low_latency", + "nixl_ep", + ) + and self.enable_expert_parallel + and self.data_parallel_size > 1 + ) + @property def node_rank_within_dp(self) -> int: return self.node_rank % self.nnodes_within_dp diff --git a/vllm/config/pooler.py b/vllm/config/pooler.py index 24368c3494e..f8eefb7c2ba 100644 --- a/vllm/config/pooler.py +++ b/vllm/config/pooler.py @@ -77,10 +77,31 @@ class PoolerConfig: Defaults to None (i.e. set to max_model_len). """ - ## for classification models + ## for classification models — affine score calibration + logit_mean: float | None = None + """ + If provided, subtract this value from classification logits before + activation. Used for affine score calibration (Platt scaling): + activation((logit - logit_mean) / logit_sigma). Defaults to None. + """ + + logit_sigma: float | None = None + """ + If provided, divide the classification logits by this value after + mean subtraction. Used for affine score calibration (Platt scaling): + activation((logit - logit_mean) / logit_sigma). Defaults to None. + """ + + # Deprecated aliases — will be removed in v0.21 logit_bias: float | None = None """ - If provided, apply classification logit biases. Defaults to None. + Deprecated: Use logit_mean instead. Will be removed in v0.21. + """ + + logit_scale: float | None = None + """ + Deprecated: Use logit_sigma instead (note: logit_sigma = 1/logit_scale). + Will be removed in v0.21. """ ## for reward models @@ -98,6 +119,39 @@ class PoolerConfig: """ def __post_init__(self) -> None: + # Handle deprecated logit_bias → logit_mean + if self.logit_bias is not None: + if self.logit_mean is not None: + raise ValueError( + "Cannot set both `logit_bias` and `logit_mean`. " + "`logit_bias` is deprecated, use `logit_mean` instead." + ) + logger.warning( + "`logit_bias` is deprecated and will be removed in v0.21. " + "Use `logit_mean` instead." + ) + self.logit_mean = self.logit_bias + self.logit_bias = None + + # Handle deprecated logit_scale → logit_sigma + if self.logit_scale is not None: + if self.logit_sigma is not None: + raise ValueError( + "Cannot set both `logit_scale` and `logit_sigma`. " + "`logit_scale` is deprecated, use `logit_sigma` instead." + ) + logger.warning( + "`logit_scale` is deprecated and will be removed in v0.21. " + "Use `logit_sigma` instead (logit_sigma = 1/logit_scale)." + ) + if self.logit_scale == 0: + raise ValueError("logit_scale cannot be 0 (division by zero)") + self.logit_sigma = 1.0 / self.logit_scale + self.logit_scale = None + + if self.logit_sigma is not None and self.logit_sigma == 0: + raise ValueError("logit_sigma cannot be 0 (division by zero)") + if pooling_type := self.pooling_type: if self.seq_pooling_type is not None: raise ValueError( diff --git a/vllm/config/reasoning.py b/vllm/config/reasoning.py index be1e2b6da58..ff5546e05eb 100644 --- a/vllm/config/reasoning.py +++ b/vllm/config/reasoning.py @@ -5,6 +5,7 @@ from dataclasses import field from vllm.config.model import ModelConfig from vllm.config.utils import config +from vllm.reasoning import ReasoningParserManager from vllm.tokenizers import cached_tokenizer_from_config @@ -18,11 +19,11 @@ class ReasoningConfig: `initialize_token_ids` and are not intended to be set directly. """ - # NOTE: These parameters are temporary, the intent is to derive them - # automatically from the reasoning parser in a future version. - reasoning_start_str: str = "" + reasoning_parser: str = "" + """The name of the ReasoningParser to use for this model.""" + reasoning_start_str: str = "" """String that indicates the start of reasoning.""" - reasoning_end_str: str = "" + reasoning_end_str: str = "" """String that indicates the end of reasoning content.""" _reasoning_start_token_ids: list[int] | None = field( @@ -36,6 +37,16 @@ class ReasoningConfig: """Private backing field for `reasoning_end_token_ids`. Set by `initialize_token_ids`. Not intended to be configured directly.""" + _enabled: bool = field(default=False, init=False, repr=False) + """Private field indicating whether reasoning token IDs have been initialized. + Set to True by `initialize_token_ids` once token IDs are initialized.""" + + @property + def enabled(self) -> bool: + """Returns True if reasoning is enabled (i.e. if token IDs have been + initialized), False otherwise.""" + return self._enabled + @property def reasoning_start_token_ids(self) -> list[int] | None: """Token IDs derived from `reasoning_start_str`. Set automatically by @@ -54,15 +65,36 @@ class ReasoningConfig: self._reasoning_start_token_ids is not None and self._reasoning_end_token_ids is not None ): - return + self._enabled = True + return # Already initialized tokenizer = cached_tokenizer_from_config(model_config=model_config) + reasoning_start_str = self.reasoning_start_str + reasoning_end_str = self.reasoning_end_str + if self.reasoning_parser is not None and ( + not reasoning_start_str or not reasoning_end_str + ): + parser_cls = ReasoningParserManager.get_reasoning_parser( + self.reasoning_parser + ) + reasoning_parser = parser_cls(tokenizer) + start_token = reasoning_parser.reasoning_start_str + if start_token and not reasoning_start_str: + reasoning_start_str = start_token + end_token = reasoning_parser.reasoning_end_str + if end_token and not reasoning_end_str: + reasoning_end_str = end_token + + if not reasoning_start_str or not reasoning_end_str: + # If we don't have valid strings to tokenize, + # we can't initialize the token IDs. + return self._reasoning_start_token_ids = tokenizer.encode( - self.reasoning_start_str, add_special_tokens=False + reasoning_start_str, add_special_tokens=False ) self._reasoning_end_token_ids = tokenizer.encode( - self.reasoning_end_str, add_special_tokens=False + reasoning_end_str, add_special_tokens=False ) if not self._reasoning_start_token_ids or not self._reasoning_end_token_ids: @@ -72,3 +104,4 @@ class ReasoningConfig: f"reasoning_end_str='{self.reasoning_end_str}'. " "Ensure the strings are valid tokens in the model's vocabulary." ) + self._enabled = True diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index 3cd99bb082e..b9a48144ded 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -40,6 +40,7 @@ class SchedulerConfig: """ DEFAULT_MAX_NUM_BATCHED_TOKENS: ClassVar[int] = 2048 + DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP: ClassVar[int] = 256 DEFAULT_MAX_NUM_SEQS: ClassVar[int] = 128 runner_type: RunnerType = "generate" diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 0e74501dd9a..bbe923f68f1 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -40,6 +40,7 @@ MTPModelTypes = Literal[ "ernie_mtp", "nemotron_h_mtp", "exaone_moe_mtp", + "exaone4_5_mtp", "qwen3_next_mtp", "qwen3_5_mtp", "longcat_flash_mtp", @@ -299,6 +300,10 @@ class SpeculativeConfig: {"n_predict": n_predict, "architectures": ["ErnieMTPModel"]} ) + if hf_config.architectures[0] == "NemotronH_Super_Omni_Reasoning_V3": + # Promote VLM's text_config so MTP detection below fires correctly + hf_config = hf_config.text_config + if ( hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"} and hasattr(hf_config, "num_nextn_predict_layers") @@ -327,7 +332,13 @@ class SpeculativeConfig: hf_config.update( {"n_predict": n_predict, "architectures": ["ExaoneMoeMTP"]} ) - + if "exaone4_5" in hf_config.model_type: + hf_config.model_type = "exaone4_5_mtp" + if hf_config.model_type == "exaone4_5_mtp": + n_predict = getattr(hf_config, "num_nextn_predict_layers", None) + hf_config.update( + {"n_predict": n_predict, "architectures": ["Exaone4_5_MTP"]} + ) if hf_config.model_type in ("qwen3_5", "qwen3_5_moe"): is_moe = hf_config.model_type == "qwen3_5_moe" hf_config.model_type = "qwen3_5_mtp" @@ -818,6 +829,7 @@ class SpeculativeConfig: "kimi_k2", "kimi_k25", "minimax_m2", + "gemma4", ] if ( self.method in ("eagle3", "extract_hidden_states", "dflash") diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6229b44d52a..47bc3547ce8 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -37,6 +37,7 @@ from .kv_events import KVEventsConfig from .kv_transfer import KVTransferConfig from .load import LoadConfig from .lora import LoRAConfig +from .mamba import MambaConfig from .model import ModelConfig from .observability import ObservabilityConfig from .offload import OffloadConfig @@ -275,6 +276,8 @@ class VllmConfig: """Model weight offloading configuration.""" attention_config: AttentionConfig = Field(default_factory=AttentionConfig) """Attention configuration.""" + mamba_config: MambaConfig = Field(default_factory=MambaConfig) + """Mamba configuration.""" kernel_config: KernelConfig = Field(default_factory=KernelConfig) """Kernel configuration.""" lora_config: LoRAConfig | None = None @@ -559,6 +562,16 @@ class VllmConfig: if architectures is not None: hf_config = copy.deepcopy(hf_config) hf_config.architectures = architectures + elif hf_config.architectures is None: + from transformers.models.auto.modeling_auto import ( + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + ) + + if hf_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES: + hf_config = copy.deepcopy(hf_config) + hf_config.architectures = [ + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[hf_config.model_type] + ] model_config = copy.deepcopy(self.model_config) @@ -707,6 +720,18 @@ class VllmConfig: if self.lora_config is not None: self.lora_config.verify_with_model_config(self.model_config) + if ( + self.mamba_config.enable_stochastic_rounding + and self.cache_config.mamba_ssm_cache_dtype != "float16" + ): + raise ValueError( + "Stochastic rounding for Mamba cache requires " + "the SSM cache to be float16. Please set it explicitly, " + "by specifying `--mamba-ssm-cache-dtype float16`, or disable " + "stochastic rounding by not specifying " + "`--enable-mamba-cache-stochastic-rounding`." + ) + if self.quant_config is None and self.model_config is not None: self.quant_config = VllmConfig._get_quantization_config( self.model_config, self.load_config @@ -764,6 +789,16 @@ class VllmConfig: elif self.scheduler_config.async_scheduling is None: # Enable async scheduling unless there is an incompatible option. if ( + self.model_config is not None + and self.model_config.runner_type == "pooling" + ): + # The current implementation of asynchronous scheduling negatively + # impacts performance of pooling models, so we disable by default. + logger.debug( + "Disabling asynchronous scheduling by default for pooling model." + ) + self.scheduler_config.async_scheduling = False + elif ( self.speculative_config is not None and self.speculative_config.method not in get_args(EagleModelTypes) and self.speculative_config.method not in get_args(NgramGPUTypes) @@ -1210,6 +1245,12 @@ class VllmConfig: if self.reasoning_config is not None and self.model_config is not None: self.reasoning_config.initialize_token_ids(self.model_config) + if not self.reasoning_config.enabled: + logger.warning_once( + "Auto-initialization of reasoning token IDs failed. " + "Please check whether your reasoning parser has implemented " + "the `reasoning_start_str` and `reasoning_end_str`." + ) # Hybrid KV cache manager (HMA) runtime rules: # - Explicit enable (--no-disable-kv-cache-manager): error if runtime @@ -1223,9 +1264,6 @@ class VllmConfig: if not current_platform.support_hybrid_kv_cache(): # Hybrid KV cache manager is not supported on non-GPU platforms. need_disable_hybrid_kv_cache_manager = True - if self.kv_events_config is not None: - # Hybrid KV cache manager is not compatible with KV events. - need_disable_hybrid_kv_cache_manager = True if ( self.model_config is not None and self.model_config.attention_chunk_size is not None @@ -1621,6 +1659,22 @@ class VllmConfig: compile_range_end, ) + if compilation_config.pass_config.fuse_minimax_qk_norm: + from vllm.compilation.passes.fusion.minimax_qk_norm_fusion import ( + MAX_TOKEN_NUM, + ) + + max_token_num = min( + MAX_TOKEN_NUM, self.scheduler_config.max_num_batched_tokens + ) + if compile_range_end is not None and max_token_num < compile_range_end: + computed_compile_ranges_endpoints.append(max_token_num) + else: + logger.debug( + "Max num batched tokens below MiniMax QK norm fusion threshold, " + "MiniMax QK norm fusion enabled for all num_tokens." + ) + if compilation_config.compile_ranges_endpoints is not None: for x in compilation_config.compile_ranges_endpoints: assert isinstance(x, int) diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index 21ec7a36e98..d3e304f8b60 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -67,6 +67,8 @@ class BlockStored(KVCacheEvent): KV cache consumers to reconstruct block hashes. """ + group_idx: int | None = None + def __hash__(self) -> int: return hash( ( @@ -77,6 +79,7 @@ class BlockStored(KVCacheEvent): self.lora_id, self.medium, tuple(self.extra_keys) if self.extra_keys else None, + self.group_idx, ) ) @@ -84,9 +87,16 @@ class BlockStored(KVCacheEvent): class BlockRemoved(KVCacheEvent): block_hashes: list[ExternalBlockHash] medium: str | None + group_idx: int | None = None def __hash__(self) -> int: - return hash((tuple(self.block_hashes), self.medium)) + return hash( + ( + tuple(self.block_hashes), + self.medium, + self.group_idx, + ) + ) class AllBlocksCleared(KVCacheEvent): diff --git a/vllm/distributed/kv_transfer/kv_connector/factory.py b/vllm/distributed/kv_transfer/kv_connector/factory.py index 9f8379fecd3..f691b9f18e9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/factory.py +++ b/vllm/distributed/kv_transfer/kv_connector/factory.py @@ -178,7 +178,7 @@ KVConnectorFactory.register_connector( KVConnectorFactory.register_connector( "NixlConnector", - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector", + "vllm.distributed.kv_transfer.kv_connector.v1.nixl", "NixlConnector", ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py index 07e05cc8f89..3d3a093820e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py @@ -7,6 +7,7 @@ from .multi_process_adapter import ( LMCacheMPSchedulerAdapter, LMCacheMPWorkerAdapter, LoadStoreOp, + ParallelStrategy, ) __all__ = [ @@ -15,4 +16,5 @@ __all__ = [ "LMCacheMPSchedulerAdapter", "LMCacheMPWorkerAdapter", "LoadStoreOp", + "ParallelStrategy", ] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py index eff580df902..2e75519df12 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py @@ -79,6 +79,39 @@ def get_lmcache_chunk_size( return chunk_size +@dataclass +class ParallelStrategy: + use_mla: bool + """Whether to use the MLA.""" + + kv_world_size: int + """ + The kv world size, kv_world_size may not be equal to the actual_world_size, + in the case of mla, it will 'exclude' the effect of TP, the value is + calculated by `extract_world_size_and_kv_rank` in `lmcache_mp_connector.py`. + """ + + kv_worker_id: int + """ + The kv worker id of the sub-process, kv_worker_id may not be equal to the + actual_worker_id, in the case of mla, it will 'exclude' the effect of TP, + the value is calculated by `extract_world_size_and_kv_rank` in + `lmcache_mp_connector.py`. + """ + + actual_world_size: int + """The actual world size.""" + + actual_worker_id: int + """The actual worker id of the sub-process.""" + + tp_size: int + """The tensor parallel size.""" + + pp_size: int + """The pipeline parallel size.""" + + @dataclass class LoadStoreOp: block_ids: list[int] @@ -111,10 +144,8 @@ class LMCacheMPSchedulerAdapter: server_url: str, context: zmq.Context, model_name: str, - world_size: int, - kv_rank: int, vllm_block_size: int, - tp_size: int = 1, + parallel_strategy: ParallelStrategy, ): """ Args: @@ -122,11 +153,10 @@ class LMCacheMPSchedulerAdapter: context: The ZMQ context model_name: The model name used for LMCache keys - world_size: The world size used for LMCache keys - kv_rank: The kv rank used for LMCache keys vllm_block_size: The block size used in vLLM - tp_size: Tensor-parallel size for MLA - multi-reader locking (default 1). + parallel_strategy: + The parallel strategy, which includes `use_mla`, + `world_size`, `worker_id` and so on """ self.mq_client = MessageQueueClient(server_url, context) @@ -134,9 +164,7 @@ class LMCacheMPSchedulerAdapter: self.lookup_futures: dict[str, MessagingFuture[LookupResult]] = {} self.model_name = model_name - self.world_size = world_size - self.worker_id = kv_rank - self.tp_size = tp_size + self.parallel_strategy = parallel_strategy # Read chunk size from lmcache self.chunk_size = get_lmcache_chunk_size(self.mq_client) @@ -145,6 +173,21 @@ class LMCacheMPSchedulerAdapter: ) self.blocks_in_chunk = self.chunk_size // vllm_block_size + @property + def world_size(self) -> int: + """The world size.""" + return self.parallel_strategy.kv_world_size + + @property + def worker_id(self) -> int: + """The worker id.""" + return self.parallel_strategy.kv_worker_id + + @property + def tp_size(self) -> int: + """The tensor parallel size.""" + return self.parallel_strategy.tp_size + @_lmcache_nvtx_annotate def maybe_submit_lookup_request( self, @@ -308,9 +351,8 @@ class LMCacheMPWorkerAdapter: server_url: str, context: zmq.Context, model_name: str, - world_size: int, - kv_rank: int, vllm_block_size: int, + parallel_strategy: ParallelStrategy, ): self.mq_client = MessageQueueClient(server_url, context) @@ -336,8 +378,7 @@ class LMCacheMPWorkerAdapter: self.previously_finished: set[str] = set() self.model_name = model_name - self.world_size = world_size - self.worker_id = kv_rank + self.parallel_strategy = parallel_strategy # Read chunk size from lmcache chunk_size = get_lmcache_chunk_size(self.mq_client) @@ -346,6 +387,29 @@ class LMCacheMPWorkerAdapter: ) self.blocks_in_chunk = chunk_size // vllm_block_size + @property + def world_size(self) -> int: + """The world size.""" + return self.parallel_strategy.kv_world_size + + @property + def worker_id(self) -> int: + """The worker id.""" + return self.parallel_strategy.kv_worker_id + + @property + def use_mla(self) -> bool: + """Whether to use MLA.""" + return self.parallel_strategy.use_mla + + @property + def is_first_rank_of_pp_group(self) -> bool: + """Is the first rank of the pipeline parallel group.""" + return ( + self.parallel_strategy.actual_worker_id % self.parallel_strategy.tp_size + == 0 + ) + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """ Register the kv caches with LMCache server diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py index 5f14c733a8b..c6d46b49af5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/lmcache_mp_connector.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import enum -import inspect from collections.abc import Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal @@ -28,12 +27,25 @@ try: LMCacheMPSchedulerAdapter, LMCacheMPWorkerAdapter, LoadStoreOp, + ParallelStrategy, ) + + try: + from lmcache.v1.multiprocess.custom_types import RequestAllocationRecord + except ImportError: + from lmcache.v1.multiprocess.custom_types import ( + BlockAllocationRecord as RequestAllocationRecord, + ) except ImportError: + from lmcache.v1.multiprocess.custom_types import ( + BlockAllocationRecord as RequestAllocationRecord, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration import ( LMCacheMPSchedulerAdapter, LMCacheMPWorkerAdapter, LoadStoreOp, + ParallelStrategy, ) if TYPE_CHECKING: @@ -53,12 +65,6 @@ if TYPE_CHECKING: logger = lmcache_init_logger(__name__) -def _adapter_accepts_tp_size() -> bool: - """Check if the imported adapter accepts tp_size.""" - sig = inspect.signature(LMCacheMPSchedulerAdapter.__init__) - return "tp_size" in sig.parameters - - # Helper functions def reformat_block_ids(block_ids: tuple[list[int], ...] | None) -> list[int]: if block_ids is None: @@ -94,8 +100,8 @@ def extract_world_size_and_kv_rank( # vLLM constructs TP groups first, and then construct other # parallel groups on top of TP groups. # for example, TP=4, PP=2, - # TP group: [0, 1, 2, 3], [4, 5, 6, 7] - # PP group: [0, 4], [1, 5], [2, 6], [3, 7] + # PP group: [0, 1, 2, 3], [4, 5, 6, 7] + # TP group: [0, 4], [1, 5], [2, 6], [3, 7] # So we can "exclude" the effect of TP by rank // tp_size. return world_size // tp_size, rank // tp_size @@ -112,24 +118,24 @@ def create_scheduler_adapter( vllm_config.parallel_config.rank, vllm_config, ) - tp_size = vllm_config.parallel_config.tensor_parallel_size - - # Pass tp_size only when the adapter accepts it so that - # a newer vllm can still work with an older LMCache. - kwargs: dict[str, Any] = {} - if _adapter_accepts_tp_size(): - kwargs["tp_size"] = tp_size - - return LMCacheMPSchedulerAdapter( - server_url, - zmq_context, - vllm_config.model_config.model, + parallel_strategy = ParallelStrategy( + mla_enabled(vllm_config.model_config), world_size, kv_rank, - vllm_config.cache_config.block_size, + vllm_config.parallel_config.world_size, + vllm_config.parallel_config.rank, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.pipeline_parallel_size, + ) + + return LMCacheMPSchedulerAdapter( + server_url=server_url, + context=zmq_context, + model_name=vllm_config.model_config.model, + vllm_block_size=vllm_config.cache_config.block_size, + parallel_strategy=parallel_strategy, mq_timeout=mq_timeout, heartbeat_interval=heartbeat_interval, - **kwargs, ) @@ -145,13 +151,22 @@ def create_worker_adapter( vllm_config.parallel_config.rank, vllm_config, ) - return LMCacheMPWorkerAdapter( - server_url, - zmq_context, - vllm_config.model_config.model, + parallel_strategy = ParallelStrategy( + mla_enabled(vllm_config.model_config), world_size, kv_rank, - vllm_config.cache_config.block_size, + vllm_config.parallel_config.world_size, + vllm_config.parallel_config.rank, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.pipeline_parallel_size, + ) + + return LMCacheMPWorkerAdapter( + server_url=server_url, + context=zmq_context, + model_name=vllm_config.model_config.model, + vllm_block_size=vllm_config.cache_config.block_size, + parallel_strategy=parallel_strategy, mq_timeout=mq_timeout, heartbeat_interval=heartbeat_interval, ) @@ -200,8 +215,11 @@ class LMCacheMPRequestTracker: # Main state state: LMCacheMPRequestState = LMCacheMPRequestState.PREFETCHING + cache_salt: str = "" + def __init__(self, request: "Request"): self.request_id = request.request_id + self.cache_salt: str = request.cache_salt or "" self.all_token_ids = request.all_token_ids self.block_hashes = ConstantList(request.block_hashes) self.allocated_block_ids = [] @@ -274,6 +292,7 @@ class LMCacheMPRequestMetadata: request_id: str direction: Literal["STORE", "RETRIEVE"] op: LoadStoreOp + cache_salt: str = "" @staticmethod def GetStoreMetadata( @@ -293,10 +312,31 @@ class LMCacheMPRequestMetadata: # NOTE: the invariant here is that `num_stored_blocks` should # always be a multiple of `blocks_in_chunk` # TODO: This should be checked everytime we update the num_stored_blocks + # + # Why computed_blocks uses max(num_vllm_hit_blocks, num_lmcache_hit_blocks): + # + # Both values represent a prefix of blocks whose KV data is already + # available (either from vLLM APC or from LMCache), so they must NOT + # be summed (that would double-count the overlapping prefix). + # + # * num_lmcache_hit_blocks: LMCache-hit blocks are already counted in + # num_stored_blocks (set during lookup), so they must be included + # here to keep the upper bound consistent. They are NOT re-stored. + # * num_vllm_hit_blocks: LMCache stores in units of chunks (N blocks), + # so num_lmcache_hit_blocks is rounded DOWN to the nearest chunk + # boundary. When vLLM APC hits more blocks than that rounded value + # (e.g. APC=44 blocks, LMCache=32 blocks after chunk alignment), + # using only num_lmcache_hit_blocks would set the upper bound too + # low and silently skip the APC-hit blocks that fall between the + # two values, causing under-storing. Taking the max ensures we + # always use the tighter (larger) of the two hit counts. + computed_blocks = tracker.num_scheduled_tokens // vllm_block_size + max( + tracker.num_vllm_hit_blocks, tracker.num_lmcache_hit_blocks + ) min_available_blocks = min( len(tracker.block_hashes), len(tracker.allocated_block_ids), - tracker.num_scheduled_tokens // vllm_block_size, + computed_blocks, ) num_staging_blocks = min_available_blocks - tracker.num_stored_blocks num_chunks = num_staging_blocks // blocks_in_chunk @@ -319,6 +359,7 @@ class LMCacheMPRequestMetadata: request_id=tracker.request_id, direction="STORE", op=op, + cache_salt=tracker.cache_salt, ) # Update the request tracker @@ -385,6 +426,7 @@ class LMCacheMPRequestMetadata: request_id=tracker.request_id, direction="RETRIEVE", op=op, + cache_salt=tracker.cache_salt, ) return ret @@ -533,12 +575,14 @@ class LMCacheMPConnector(KVConnectorBase_V1): request_ids = [] ops = [] + cache_salts = [] for meta in metadata.requests: if meta.direction != "RETRIEVE": continue request_ids.append(meta.request_id) ops.append(meta.op) + cache_salts.append(meta.cache_salt) if len(request_ids) == 0: return @@ -547,7 +591,9 @@ class LMCacheMPConnector(KVConnectorBase_V1): event = torch.cuda.Event(interprocess=True) event.record() - self.worker_adapter.batched_submit_retrieve_requests(request_ids, ops, event) + self.worker_adapter.batched_submit_retrieve_requests( + request_ids, ops, event, cache_salts=cache_salts + ) def wait_for_layer_load(self, layer_name: str) -> None: """ @@ -591,16 +637,26 @@ class LMCacheMPConnector(KVConnectorBase_V1): This prevents overwrites of paged KV buffer before saving done. """ + # In MLA scenario, only the first rank of the pipeline group + # needs to save the KV cache. + if ( + self.worker_adapter.use_mla + and not self.worker_adapter.is_first_rank_of_pp_group + ): + return + metadata = self._get_connector_metadata() assert isinstance(metadata, LMCacheMPConnectorMetadata) request_ids = [] ops = [] + cache_salts = [] for meta in metadata.requests: if meta.direction != "STORE": continue request_ids.append(meta.request_id) ops.append(meta.op) + cache_salts.append(meta.cache_salt) if len(request_ids) == 0: return @@ -609,7 +665,9 @@ class LMCacheMPConnector(KVConnectorBase_V1): event = torch.cuda.Event(interprocess=True) event.record() - self.worker_adapter.batched_submit_store_requests(request_ids, ops, event) + self.worker_adapter.batched_submit_store_requests( + request_ids, ops, event, cache_salts=cache_salts + ) def get_finished( self, finished_req_ids: set[str] @@ -711,6 +769,7 @@ class LMCacheMPConnector(KVConnectorBase_V1): self.scheduler_adapter.maybe_submit_lookup_request( request.request_id, token_ids=list(request.all_token_ids), + cache_salt=tracker.cache_salt, ) ret = self.scheduler_adapter.check_lookup_result(request.request_id) @@ -837,6 +896,9 @@ class LMCacheMPConnector(KVConnectorBase_V1): if len(metadata) > 0: logger.debug("Final connector metadata: %s", metadata) + # Report block allocation deltas to LMCache for observability + self._report_block_allocation_deltas(scheduler_output) + return metadata def update_connector_output(self, connector_output: KVConnectorOutput): @@ -996,8 +1058,9 @@ class LMCacheMPConnector(KVConnectorBase_V1): if request_id not in cached_reqs.resumed_req_ids: request_tracker.append_block_ids(new_block_ids) - # Update new scheduled tokens - num_new_tokens = cached_reqs.num_computed_tokens[idx] + # Use the incremental num_scheduled_tokens to + # stay consistent with _process_new_requests. + num_new_tokens = scheduler_output.num_scheduled_tokens[request_id] request_tracker.increase_num_scheduled_tokens(num_new_tokens) r_meta = LMCacheMPRequestMetadata.GetStoreMetadata( @@ -1007,6 +1070,64 @@ class LMCacheMPConnector(KVConnectorBase_V1): if r_meta is not None: metadata.add_request_metadata(r_meta) + def _report_block_allocation_deltas( + self, + scheduler_output: SchedulerOutput, + ) -> None: + """Gather per-request block allocation deltas and report to LMCache. + + For new requests: all allocated_block_ids and token_ids are new. + For cached requests: only newly appended block_ids and token_ids. + """ + records: list[RequestAllocationRecord] = [] + + # New requests: send all tokens covering all allocated blocks so + # the L0 metrics subscriber can correctly map each block to its + # actual token content (not just the newly-scheduled slice). + for new_request in scheduler_output.scheduled_new_reqs: + tracker = self.request_trackers.get(new_request.req_id) + if tracker is None: + continue + num_blocks = len(tracker.allocated_block_ids) + total_tokens = num_blocks * self.vllm_block_size + records.append( + RequestAllocationRecord( + req_id=new_request.req_id, + new_block_ids=list(tracker.allocated_block_ids), + new_token_ids=list(tracker.all_token_ids[:total_tokens]), + ) + ) + + # Cached requests: only the newly added blocks and their full + # token content. We send all tokens covered by the new blocks + # (not just the tokens scheduled this step) so the L0 subscriber + # can correctly identify block content. + cached_reqs = scheduler_output.scheduled_cached_reqs + for idx, request_id in enumerate(cached_reqs.req_ids): + new_block_ids = reformat_block_ids(cached_reqs.new_block_ids[idx]) + if not new_block_ids: + continue + tracker = self.request_trackers.get(request_id) + if tracker is None: + continue + # The new blocks sit at the end of the request's block list. + # Compute the token range they cover. + total_blocks = len(tracker.allocated_block_ids) + num_new_blocks = len(new_block_ids) + start_token = (total_blocks - num_new_blocks) * self.vllm_block_size + end_token = total_blocks * self.vllm_block_size + new_token_ids = list(tracker.all_token_ids[start_token:end_token]) + records.append( + RequestAllocationRecord( + req_id=request_id, + new_block_ids=new_block_ids, + new_token_ids=new_token_ids, + ) + ) + + if records: + self.scheduler_adapter.report_block_allocations(records) + def _get_request_tracker(self, request_id: str) -> LMCacheMPRequestTracker: assert request_id in self.request_trackers, ( f"Request tracker for request_id {request_id} not found. " diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index b49a016641e..67603e10ff6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import ( EngineId, TpKVTopology, get_current_attn_backend, + get_current_attn_backends, ) from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, @@ -41,11 +42,13 @@ from vllm.distributed.parallel_state import ( ) from vllm.forward_context import ForwardContext from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.request import RequestStatus +from vllm.v1.worker.utils import select_common_block_size logger = init_logger(__name__) @@ -645,6 +648,10 @@ class MooncakeConnectorWorker: logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id) self.vllm_config = vllm_config + # Capture device BEFORE TransferEngine init — MNNVL's NVLink allocator + # may change the current CUDA device during engine.initialize(). + self.device_id = torch.accelerator.current_device_index() + current_platform.set_device(self.device_id) self.engine = TransferEngine() self.hostname = get_ip() @@ -705,9 +712,12 @@ class MooncakeConnectorWorker: # For kv_both, we will act both prefiller and decoder. if not self.is_kv_consumer: # Background threads for sending kvcaches to D. + # Each pool thread must be bound to the correct CUDA device + # because CUDA device selection is thread-local. self._sender_executor = ThreadPoolExecutor( max_workers=self.num_sender_workers, thread_name_prefix="vllm-mooncake-sender", + initializer=self._bind_sender_thread_device, ) logger.debug( "Mooncake Prefiller: use %d workers to send kvcaches", @@ -743,6 +753,7 @@ class MooncakeConnectorWorker: self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config self.use_mla = self.model_config.use_mla + self._sync_block_size_with_kernel() # Get the attention backend from the first layer # NOTE (NickLucche) models with multiple backends are not supported yet @@ -769,6 +780,23 @@ class MooncakeConnectorWorker: self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata) self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse) + def _sync_block_size_with_kernel(self) -> None: + # When speculative decoding (e.g. Eagle) is enabled, the main model + # and draft model may use different attention backends with different + # physical block sizes. Pick the common (smallest) block size so that + # KV-cache registration and transfer work correctly for both models. + backends = get_current_attn_backends(self.vllm_config) + kernel_block_size = select_common_block_size(self.block_size, backends) + if self.block_size != kernel_block_size: + logger.info_once( + "User-specified logical block size (%s) does not match" + " physical kernel block size (%s). Using the latter.", + self.block_size, + kernel_block_size, + ) + assert self.block_size > kernel_block_size + self.block_size = kernel_block_size + def __del__(self): self.shutdown() @@ -1193,6 +1221,12 @@ class MooncakeConnectorWorker: return src_ptrs, dst_ptrs, lengths, err_reqs, err_msg + def _bind_sender_thread_device(self) -> None: + """ThreadPoolExecutor initializer — binds each pool thread to the + correct CUDA device. CUDA device selection is thread-local, so + without this, NVLink transfers fail for TP ranks > 0.""" + current_platform.set_device(self.device_id) + def _send_blocks( self, remote_session: str, @@ -1254,9 +1288,6 @@ class MooncakeConnectorWorker: self.block_len_per_layer.append( curr_tensor_size_bytes // self.num_blocks ) - - kernel_block_size = cache.shape[-2 if self.use_mla else -3] - assert self.block_size == kernel_block_size kv_data_ptrs.append(base_addr) kv_data_lens.append(curr_tensor_size_bytes) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py new file mode 100644 index 00000000000..ed5c892fb9d --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NIXL KV-cache transfer connector (disaggregated prefill / decode).""" + +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import ( + NixlConnector, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlAgentMetadata, + NixlConnectorMetadata, + NixlHandshakePayload, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( + NixlConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( + NixlKVConnectorStats, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, +) + +__all__ = [ + "NixlAgentMetadata", + "NixlConnector", + "NixlConnectorMetadata", + "NixlConnectorScheduler", + "NixlConnectorWorker", + "NixlHandshakePayload", + "NixlKVConnectorStats", +] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py new file mode 100644 index 00000000000..53ad031a4c5 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NixlConnector – thin facade that delegates to scheduler / worker.""" + +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import ( + EngineId, + get_current_attn_backend, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + CopyBlocksOp, + KVConnectorBase_V1, + KVConnectorHandshakeMetadata, + KVConnectorMetadata, + KVConnectorRole, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + NixlConnectorMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import ( + NixlConnectorScheduler, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( + NixlKVConnectorStats, + NixlPromMetrics, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, +) +from vllm.forward_context import ForwardContext +from vllm.logger import init_logger +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.utils import get_kv_cache_layout +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import MambaSpec + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class NixlConnector(KVConnectorBase_V1, SupportsHMA): + @property + def prefer_cross_layer_blocks(self) -> bool: + if any( + [ + isinstance(group.kv_cache_spec, MambaSpec) + for group in self.kv_cache_config.kv_cache_groups + ] + ): + # Hybrid SSM models do not yet support cross-layer layout + return False + + backend = get_current_attn_backend(self._vllm_config) + if backend.get_name() not in ( + "FLASH_ATTN", + "FLASHINFER", + "TRITON_ATTN", + ): + return False + + # For now there is no benefit to run cross layers when backend + # does not support on HND + if get_kv_cache_layout() != "HND": + return False + + extra_config = self.kv_transfer_config.kv_connector_extra_config + return ( + str(extra_config.get("enable_cross_layers_blocks", "False")).lower() + == "true" + ) + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: "KVCacheConfig", + ): + super().__init__(vllm_config, role, kv_cache_config) + assert vllm_config.kv_transfer_config is not None + assert vllm_config.kv_transfer_config.engine_id is not None + self.kv_cache_config = kv_cache_config + self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id + self.kv_transfer_config = vllm_config.kv_transfer_config + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler: NixlConnectorScheduler | None = ( + NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config) + ) + self.connector_worker: NixlConnectorWorker | None = None + elif role == KVConnectorRole.WORKER: + self.connector_scheduler = None + self.connector_worker = NixlConnectorWorker( + vllm_config, self.engine_id, kv_cache_config + ) + + ############################################################ + # Class Methods + ############################################################ + @classmethod + def get_required_kvcache_layout(cls, vllm_config: VllmConfig): + if vllm_config.model_config is None: + logger.warning_once( + "Unable to detect current VLLM config. " + "Fallback to default kv cache layout." + ) + return None + use_mla = vllm_config.model_config.use_mla + if use_mla: + # return None when we have mla + # as the layout should not matter in that case, + # which fallback to the default behavior. + return None + logger.info_once( + "NixlConnector setting KV cache layout to HND for better xfer performance." + ) + return "HND" + + ############################################################ + # Scheduler Side Methods + ############################################################ + + def get_num_new_matched_tokens( + self, request: "Request", num_computed_tokens: int + ) -> tuple[int | None, bool]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_num_new_matched_tokens( + request, num_computed_tokens + ) + + def update_state_after_alloc( + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int + ): + assert self.connector_scheduler is not None + return self.connector_scheduler.update_state_after_alloc( + request, blocks, num_external_tokens + ) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> KVConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta(scheduler_output) + + def request_finished( + self, + request: "Request", + block_ids: list[int], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, (block_ids,)) + + def request_finished_all_groups( + self, + request: "Request", + block_ids: tuple[list[int], ...], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, block_ids) + + def set_xfer_handshake_metadata( + self, metadata: dict[int, KVConnectorHandshakeMetadata] + ) -> None: + """ + Set the KV connector handshake metadata for this connector. + + Args: + metadata (dict): the handshake metadata to set. + """ + assert self.connector_scheduler is not None + self.connector_scheduler.set_xfer_handshake_metadata(metadata) + + ############################################################ + # Worker Side Methods + ############################################################ + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): + assert self.connector_worker is not None + self.connector_worker.register_kv_caches(kv_caches) + + def register_cross_layers_kv_cache( + self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] + ): + assert self.connector_worker is not None + self.connector_worker.register_cross_layers_kv_caches(kv_cache) + + def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): + assert self.connector_worker is not None + self.connector_worker.set_host_xfer_buffer_ops(copy_operation) + + def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: + """Get the finished recving and sending requests.""" + assert self.connector_worker is not None + return self.connector_worker.get_finished() + + def get_block_ids_with_load_errors(self) -> set[int]: + """Get block IDs that failed to load via NIXL.""" + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + + def get_kv_connector_stats(self) -> KVConnectorStats | None: + if self.connector_worker is None: + return None + return self.connector_worker.get_kv_connector_stats() + + @classmethod + def build_kv_connector_stats( + cls, data: dict[str, Any] | None = None + ) -> KVConnectorStats | None: + return ( + NixlKVConnectorStats(data=data) + if data is not None + else NixlKVConnectorStats() + ) + + @classmethod + def build_prom_metrics( + cls, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> KVConnectorPromMetrics: + return NixlPromMetrics( + vllm_config, metric_types, labelnames, per_engine_labelvalues + ) + + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: + assert self.connector_worker is not None + assert isinstance(self._connector_metadata, NixlConnectorMetadata) + self.connector_worker.start_load_kv(self._connector_metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: + """NixlConnector does not do layerwise saving.""" + pass + + def save_kv_layer( + self, + layer_name: str, + kv_layer: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> None: + """NixlConnector does not save explicitly.""" + pass + + def wait_for_save(self): + assert self.connector_worker is not None + assert isinstance(self._connector_metadata, NixlConnectorMetadata) + if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks: + self.connector_worker.save_kv_to_host(self._connector_metadata) + + def shutdown(self): + if self.connector_worker is not None: + self.connector_worker.shutdown() + if self.connector_scheduler is not None: + self.connector_scheduler.shutdown() + + def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: + """ + Get the KVConnector handshake metadata for this connector. + This metadata is used for out-of-band connector handshake + between P/D workers. + + Returns: + KVConnectorHandshakeMetadata: the handshake metadata. + None if no handshake metadata is available. + """ + assert self.connector_worker is not None + return self.connector_worker.xfer_handshake_metadata diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py new file mode 100644 index 00000000000..71ebbf1174f --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Metadata dataclasses and helpers for the NIXL connector.""" + +from dataclasses import dataclass +from typing import Any + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, + KVConnectorMetadata, +) +from vllm.logger import init_logger + +logger = init_logger(__name__) + +TransferHandle = int +ReqId = str + +GET_META_MSG = b"get_meta_msg" +# +# NIXL Connector Version +# +# Increment this version whenever there is an incompatible change to: +# - NixlAgentMetadata schema +# - kv_transfer_params schema or semantics +# - NIXL transfer protocol or wire format +# - KV cache memory layout or block organization +# - Any other change that breaks P/D interoperability +# +# Version History: +# 1: Initial version with compatibility checking +# 2: Add remote_request_id to kv_transfer_params +# +NIXL_CONNECTOR_VERSION: int = 2 + + +@dataclass +class NixlAgentMetadata: + engine_id: str + agent_metadata: bytes + kv_caches_base_addr: list[int] + device_id: int + num_blocks: int + block_lens: list[int] + kv_cache_layout: str + block_size: int + ssm_sizes: tuple[int, int] + attn_backend_name: str + + +@dataclass +class NixlHandshakePayload(KVConnectorHandshakeMetadata): + """ + Wrapper for NIXL handshake sent over the wire. + + Enables two-phase decoding for graceful compatibility checking: + 1. Decode NixlHandshakePayload to get compatibility_hash + 2. Compute local hash and compare + 3. Only if hashes match, decode agent_metadata_bytes + + This prevents decoder errors when NixlAgentMetadata schema is + incompatible, allowing graceful failure with clear error message. + """ + + compatibility_hash: str + agent_metadata_bytes: bytes # NixlAgentMetadata encoded + + +def compute_nixl_compatibility_hash( + vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool +) -> str: + """ + Compute compatibility hash for NIXL KV transfer. + + Hash only the factors that affect whether two NIXL instances can + successfully transfer KV cache data. + + Factors included: + - vLLM version and NIXL connector version + - Model architecture (name, dtype, KV heads, layers) + - KV cache format (dtype, sliding window) + - Attention backend + + Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout + are validated at runtime in _validate_remote_agent_handshake and are not + included in this hash to support heterogeneous deployments. + + Note - the set of factors are likely to evolve significantly over + time to be more or less permissive. + + Returns: + SHA-256 hex digest + """ + from vllm import __version__ as vllm_version + from vllm.config.utils import hash_factors + + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + + factors = { + # Version compatibility + "vllm_version": vllm_version, + "nixl_connector_version": NIXL_CONNECTOR_VERSION, + # Model architecture - affects KV cache shape + "model": model_config.model, + "dtype": str(model_config.dtype), + "num_kv_heads": model_config.get_total_num_kv_heads(), + "head_size": model_config.get_head_size(), + "num_hidden_layers": model_config.get_total_num_hidden_layers(), + # Attention backend and KV cache dtype affect memory layout + "attn_backend_name": attn_backend_name, + "cache_dtype": str(cache_config.cache_dtype), + "cross_layers_blocks": cross_layers_blocks, + "is_hma_enabled": is_hma_enabled, + } + + compat_hash = hash_factors(factors) + logger.debug( + "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, " + "cache_dtype=%s, attn_backend=%s)", + compat_hash, + factors["model"], + factors["dtype"], + factors["num_kv_heads"], + factors["cache_dtype"], + attn_backend_name, + ) + return compat_hash + + +@dataclass +class RemoteMeta: + block_ids: BlockIds + host: str + port: int + engine_id: str + request_id: str + + +@dataclass +class ReqMeta: + local_block_ids: BlockIds + # To be used when logical block size does not match the kernel block size + local_physical_block_ids: BlockIds + tp_size: int + remote: RemoteMeta | None = None + + +class NixlConnectorMetadata(KVConnectorMetadata): + def __init__(self): + self.reqs_to_recv: dict[ReqId, ReqMeta] = {} + self.reqs_to_save: dict[ReqId, ReqMeta] = {} + self.reqs_to_send: dict[ReqId, float] = {} + self.reqs_in_batch: set[ReqId] = set() + self.reqs_not_processed: set[ReqId] = set() + + def _add_new_req( + self, + local_block_ids: BlockIds, + kv_transfer_params: dict[str, Any], + ) -> ReqMeta: + return ReqMeta( + local_block_ids=local_block_ids, + local_physical_block_ids=local_block_ids, + # P workers don't need to receive tp_size from proxy here. + tp_size=kv_transfer_params.get("tp_size", 1), + ) + + def add_new_req_to_save( + self, + request_id: ReqId, + local_block_ids: BlockIds, + kv_transfer_params: dict[str, Any], + ): + self.reqs_to_save[request_id] = self._add_new_req( + local_block_ids, kv_transfer_params + ) + + def add_new_req_to_recv( + self, + request_id: ReqId, + local_block_ids: BlockIds, + kv_transfer_params: dict[str, Any], + ): + req = self._add_new_req(local_block_ids, kv_transfer_params) + req.remote = RemoteMeta( + block_ids=kv_transfer_params["remote_block_ids"], + engine_id=kv_transfer_params["remote_engine_id"], + request_id=kv_transfer_params["remote_request_id"], + host=kv_transfer_params["remote_host"], + port=kv_transfer_params["remote_port"], + ) + self.reqs_to_recv[request_id] = req diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py new file mode 100644 index 00000000000..9f67d0fc525 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Scheduler-side logic for the NIXL connector.""" + +import threading +import time +from typing import TYPE_CHECKING, Any + +import msgspec +import zmq + +from vllm import envs +from vllm.distributed.kv_transfer.kv_connector.utils import ( + BlockIds, + EngineId, + yield_req_data, +) +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, + KVConnectorMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + GET_META_MSG, + NixlConnectorMetadata, + NixlHandshakePayload, + ReqId, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.math_utils import cdiv +from vllm.utils.network_utils import make_zmq_path +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + SlidingWindowSpec, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class NixlConnectorScheduler: + """Implementation of Scheduler side methods""" + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ): + self.vllm_config = vllm_config + self.block_size = vllm_config.cache_config.block_size + self.engine_id: EngineId = engine_id + self.kv_cache_config = kv_cache_config + self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST + self.side_channel_port = ( + envs.VLLM_NIXL_SIDE_CHANNEL_PORT + + vllm_config.parallel_config.data_parallel_index + ) + assert vllm_config.kv_transfer_config is not None + if current_platform.device_type == "cpu": + self.use_host_buffer = False + else: + self.use_host_buffer = ( + vllm_config.kv_transfer_config.kv_buffer_device == "cpu" + ) + self._is_hma_required = ( + not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager + # Also handle unlikely SW-only model case instead of checking num_groups>1. + and any( + not isinstance(g.kv_cache_spec, FullAttentionSpec) + for g in kv_cache_config.kv_cache_groups + ) + ) + self._has_mamba = any( + isinstance(g.kv_cache_spec, MambaSpec) + for g in kv_cache_config.kv_cache_groups + ) + + logger.info("Initializing NIXL Scheduler %s", engine_id) + if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: + logger.info("Hybrid Memory Allocator is enabled with NIXL") + + # Background thread for handling new handshake requests. + self._nixl_handshake_listener_t: threading.Thread | None = None + self._stop_event = threading.Event() + + # Requests that need to start recv/send. + # New requests are added by update_state_after_alloc in + # the scheduler. Used to make metadata passed to Worker. + self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {} + self._reqs_need_save: dict[ReqId, Request] = {} + # Reqs to send and their expiration time + self._reqs_need_send: dict[ReqId, float] = {} + self._reqs_in_batch: set[ReqId] = set() + # Reqs to remove from processed set because they're not to send after + # remote prefill or aborted. + self._reqs_not_processed: set[ReqId] = set() + + # Gather Sliding Window sizes for each kv cache group (if any) in number of + # blocks per KV cache group. This is used to clip the local attention window. + sw_sizes_tokens: list[tuple[int, int]] = [ + (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) + if isinstance(g.kv_cache_spec, SlidingWindowSpec) + else (0, self.block_size) + for g in kv_cache_config.kv_cache_groups + ] + # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively + # account for boundary overlap eg window isn't fully aligned with blocks. + self.blocks_per_sw = [ + cdiv(n_tokens, block_size) + 1 if n_tokens else 0 + for n_tokens, block_size in sw_sizes_tokens + ] + + def shutdown(self): + self._stop_event.set() + if self._nixl_handshake_listener_t is not None: + self._nixl_handshake_listener_t.join() + self._nixl_handshake_listener_t = None + + def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: + """ + Clip the number of blocks to the sliding window size for each kv cache group + that employs SWA. + This is necessary because the KV Cache manager initially allocates blocks for + the entire sequence length, and successively cleans up blocks that are outside + the window prior to the `request_finished_all_groups` hook. + """ + if len(block_ids) == 0 or not self._is_hma_required: + # No blocks to clip eg Full prefix cache hit or not a hybrid model. + return block_ids + # NOTE (NickLucche) This logic is currently handled at the connector level + # because offloading connectors might want to receive the whole sequence even + # for SWA groups. We will abstract this logic once the interface is more stable + assert len(block_ids) == len(self.blocks_per_sw), ( + "Number of KV cache groups must match" + ) + # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged + return tuple( + [ + blocks[-self.blocks_per_sw[i] :] + if self.blocks_per_sw[i] > 0 + else blocks + for i, blocks in enumerate(block_ids) + ] + ) + + def set_xfer_handshake_metadata( + self, metadata: dict[int, KVConnectorHandshakeMetadata] + ) -> None: + """ + Set the KV connector handshake metadata for this connector. + + Args: + metadata (dict): the handshake metadata to set. + """ + encoded_data: dict[int, bytes] = {} + encoder = msgspec.msgpack.Encoder() + for tp_rank, rank_metadata in metadata.items(): + if not isinstance(rank_metadata, NixlHandshakePayload): + raise ValueError( + "NixlConnectorScheduler expects NixlHandshakePayload for " + "handshake metadata." + ) + encoded_data[tp_rank] = encoder.encode(rank_metadata) + logger.debug( + "Tp rank %d: encoded NixlHandshakePayload size: %s bytes", + tp_rank, + str(len(encoded_data[tp_rank])), + ) + + # Only start the listener when we have metadata to serve. + if self._nixl_handshake_listener_t is None: + ready_event = threading.Event() + self._nixl_handshake_listener_t = threading.Thread( + target=self._nixl_handshake_listener, + args=( + encoded_data, + ready_event, + self._stop_event, + self.side_channel_port, + ), + daemon=True, + name="nixl_handshake_listener", + ) + self._nixl_handshake_listener_t.start() + ready_event.wait() # Wait for listener ZMQ socket to be ready. + + @staticmethod + def _nixl_handshake_listener( + encoded_data: dict[int, Any], + ready_event: threading.Event, + stop_event: threading.Event, + port: int, + ): + """Background thread for getting new NIXL handshakes.""" + # NOTE(rob): this is a simple implementation. We will move + # to a better approach via HTTP endpoint soon. + + # Listen for new requests for metadata. + host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST + path = make_zmq_path("tcp", host, port) + logger.debug("Starting listening on path: %s", path) + with zmq_ctx(zmq.ROUTER, path) as sock: + sock.setsockopt(zmq.RCVTIMEO, 1000) + ready_event.set() + while True: + try: + identity, _, msg = sock.recv_multipart() + except zmq.Again: + if stop_event.is_set(): + break + continue + # Decode the message which contains (GET_META_MSG, rank) + msg, target_tp_rank = msgspec.msgpack.decode(msg) + logger.debug( + "Received message for tp rank %s", + target_tp_rank, + ) + if msg != GET_META_MSG: + logger.warning("Connection listener got unexpected message %s", msg) + sock.send_multipart((identity, b"", encoded_data[target_tp_rank])) + + def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int: + """D-side only. Returns N-1 for Mamba models since the decoder + always recomputes the last token and must start from h(N-1).""" + if self._has_mamba and num_prompt_tokens > 1: + return num_prompt_tokens - 1 + return num_prompt_tokens + + def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: + """P-side only: drop the last prompt token so the prefiller computes + h(N-1) instead of h(N). The decoder recomputes the last token to + derive h(N) correctly. + + Guarded by ``_p_side_truncated`` to avoid repeated truncation if the + request is preempted and rescheduled.""" + params = request.kv_transfer_params + if ( + params is not None + # Guard against repeated truncation after preemption/reschedule. + and not params.get("_p_side_truncated") + and request.num_prompt_tokens > 1 + ): + if request.prompt_token_ids is not None: + request.prompt_token_ids.pop() + elif request.prompt_embeds is not None: + request.prompt_embeds = request.prompt_embeds[:-1] + else: + return + + request._all_token_ids.pop() + request.num_prompt_tokens -= 1 + request.max_tokens = 1 + params["_p_side_truncated"] = True + + def get_num_new_matched_tokens( + self, request: "Request", num_computed_tokens: int + ) -> tuple[int, bool]: + """ + For remote prefill, pull all prompt blocks from remote + asynchronously relative to engine execution. + + Args: + request (Request): the request object. + num_computed_tokens (int): the number of locally + computed tokens for this request + Returns: + * the number of tokens that can be loaded from the + external KV cache beyond what is already computed. + * true if the external KV cache tokens will be loaded + asynchronously (between scheduler steps). + """ + + params = request.kv_transfer_params + logger.debug( + "NIXLConnector get_num_new_matched_tokens: " + "num_computed_tokens=%s, kv_transfer_params=%s", + num_computed_tokens, + params, + ) + + if params is not None and params.get("do_remote_prefill"): + # Remote prefill: get all prompt blocks from remote. + token_ids = request.prompt_token_ids or [] + actual = self._mamba_prefill_token_count(len(token_ids)) + count = actual - num_computed_tokens + if count > 0: + return count, True + + if params is not None and params.get("do_remote_decode") and self._has_mamba: + self._truncate_mamba_request_for_prefill(request) + + # No remote prefill for this request. + return 0, False + + def update_state_after_alloc( + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int + ): + params = request.kv_transfer_params + logger.debug( + "NIXLConnector update_state_after_alloc: " + "num_external_tokens=%s, kv_transfer_params=%s", + num_external_tokens, + params, + ) + + if not params: + return + + if params.get("do_remote_decode"): + self._reqs_in_batch.add(request.request_id) + if self.use_host_buffer and params.get("do_remote_decode"): + # NOTE: when accelerator is not directly supported by Nixl, + # prefilled blocks need to be saved to host memory before transfer. + self._reqs_need_save[request.request_id] = request + elif params.get("do_remote_prefill"): + if params.get("remote_block_ids"): + if all( + p in params + for p in ( + "remote_engine_id", + "remote_request_id", + "remote_host", + "remote_port", + ) + ): + # If remote_blocks and num_external_tokens = 0, we have + # a full prefix cache hit on the D worker. We need to call + # send_notif in _read_blocks to free the memory on the P. + + unhashed_local_block_ids: BlockIds = ( + blocks.get_unhashed_block_ids_all_groups() + if num_external_tokens > 0 + else () + ) + local_block_ids = self.get_sw_clipped_blocks( + unhashed_local_block_ids + ) + + # Get unhashed blocks to pull from remote. Mind that a full prefix + # cache hit is indicated with an empty list. + self._reqs_need_recv[request.request_id] = ( + request, + local_block_ids, + ) + + else: + logger.warning( + "Got invalid KVTransferParams: %s. This " + "request will not utilize KVTransfer", + params, + ) + else: + assert num_external_tokens == 0 + # Only trigger 1 KV transfer per request. + params["do_remote_prefill"] = False + + def _build_save_meta( + self, + meta: NixlConnectorMetadata, + scheduler_output: SchedulerOutput, + ) -> None: + # only called when use_host_buffer is True to build the save metadata + + # NOTE: For the prefill side, there might be a chance that an early added + # request is a chunked prefill, so we need to check if new blocks are added + for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output): + req_to_save = self._reqs_need_save.get(req_id) + if req_to_save is None or new_block_id_groups is None: + continue + req = req_to_save + + assert req.kv_transfer_params is not None + clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups) + meta.add_new_req_to_save( + request_id=req_id, + local_block_ids=clipped_block_id_groups, + kv_transfer_params=req.kv_transfer_params, + ) + assert scheduler_output.num_scheduled_tokens is not None + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + is_partial = ( + req.num_computed_tokens + num_scheduled_tokens + ) < req.num_prompt_tokens + if not is_partial: + # For non-partial prefills, once new req_meta is scheduled, it + # can be removed from _reqs_need_save. + # For partial prefill case, we will retain the request in + # _reqs_need_save until all blocks are scheduled with req_meta. + # Therefore, only pop if `not is_partial`. + self._reqs_need_save.pop(req_id) + + def build_connector_meta( + self, + scheduler_output: SchedulerOutput, + ) -> KVConnectorMetadata: + meta = NixlConnectorMetadata() + + # Loop through scheduled reqs and convert to ReqMeta. + for req_id, (req, block_ids) in self._reqs_need_recv.items(): + assert req.kv_transfer_params is not None + meta.add_new_req_to_recv( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=req.kv_transfer_params, + ) + + if self.use_host_buffer: + self._build_save_meta(meta, scheduler_output) + + meta.reqs_to_send = self._reqs_need_send + meta.reqs_in_batch = self._reqs_in_batch + meta.reqs_not_processed = self._reqs_not_processed + + # Clear the list once workers start the transfers + self._reqs_need_recv.clear() + self._reqs_in_batch = set() + self._reqs_not_processed = set() + self._reqs_need_send = {} + + return meta + + def request_finished( + self, + request: "Request", + block_ids: BlockIds, + ) -> tuple[bool, dict[str, Any] | None]: + """ + Once a request is finished, determine whether request blocks + should be freed now or will be sent asynchronously and freed later. + """ + from vllm.v1.request import RequestStatus + + params = request.kv_transfer_params + logger.debug( + "NIXLConnector request_finished(%s), request_status=%s, " + "kv_transfer_params=%s", + request.request_id, + request.status, + params, + ) + if not params: + return False, None + + if params.get("do_remote_prefill"): + # If do_remote_prefill is still True when the request is finished, + # update_state_after_alloc must not have been called (the request + # must have been aborted before it was scheduled). + # To avoid stranding the prefill blocks in the prefill instance, + # we must add empty block_ids to _reqs_need_recv so that our + # worker side will notify and free blocks in the prefill instance. + self._reqs_need_recv[request.request_id] = (request, []) + params["do_remote_prefill"] = False + return False, None + + if not params.get("do_remote_decode"): + return False, None + if request.status != RequestStatus.FINISHED_LENGTH_CAPPED: + # Also include the case of a P/D Prefill request with immediate + # block free (eg abort). Stop tracking this request. + self._reqs_not_processed.add(request.request_id) + # Clear _reqs_need_save if a request is aborted as partial prefill. + self._reqs_need_save.pop(request.request_id, None) + return False, None + + # TODO: check whether block_ids actually ever be 0. If not we could + # remove the conditional below + delay_free_blocks = any(len(group) > 0 for group in block_ids) + + if delay_free_blocks: + # Prefill request on remote. It will be read from D upon completion + logger.debug( + "NIXLConnector request_finished(%s) waiting for %d seconds " + "for remote decode to fetch blocks", + request.request_id, + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT, + ) + self._reqs_need_send[request.request_id] = ( + time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT + ) + # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones), + # trimming down after allocating for the whole sequence length. Empty + # blocks are always at the start of the list. + # Here we "unpad" blocks to send the actual remote blocks to be read. + block_ids = self.get_sw_clipped_blocks(block_ids) + + return delay_free_blocks, dict( + do_remote_prefill=True, + do_remote_decode=False, + remote_block_ids=block_ids, + remote_engine_id=self.engine_id, + remote_request_id=request.request_id, + remote_host=self.side_channel_host, + remote_port=self.side_channel_port, + tp_size=self.vllm_config.parallel_config.tensor_parallel_size, + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py new file mode 100644 index 00000000000..fde99c1e2c2 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stats and Prometheus metrics for the NIXL connector.""" + +import copy +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( + nixlXferTelemetry, +) +from vllm.v1.metrics.utils import create_metric_per_engine + + +@dataclass +class NixlKVConnectorStats(KVConnectorStats): + """Container for transfer performance metrics""" + + def __post_init__(self): + if not self.data: + # Empty container init, no data is passed in. + self.reset() + + def reset(self): + # Must be serializable + self.data: dict[str, list[float | int]] = { + "transfer_duration": [], + "post_duration": [], + "bytes_transferred": [], + "num_descriptors": [], + "num_failed_transfers": [], + "num_failed_notifications": [], + "num_kv_expired_reqs": [], + } + + def record_transfer(self, res: nixlXferTelemetry): + # Keep metrics units consistent with rest of the code: time us->s + self.data["transfer_duration"].append(res.xferDuration / 1e6) + self.data["post_duration"].append(res.postDuration / 1e6) + self.data["bytes_transferred"].append(res.totalBytes) + self.data["num_descriptors"].append(res.descCount) + + def record_failed_transfer(self): + """Record a failed NIXL transfer operation.""" + self.data["num_failed_transfers"].append(1) + + def record_failed_notification(self): + """Record a failed NIXL notification (send_notif).""" + self.data["num_failed_notifications"].append(1) + + def record_kv_expired_req(self): + """Record a request that had its KV blocks expire.""" + self.data["num_kv_expired_reqs"].append(1) + + def clone_and_reset(self) -> "NixlKVConnectorStats": + old = copy.copy(self) + self.reset() + return old + + def is_empty(self) -> bool: + # Do not discard metrics update that are entirely failures related. + return ( + self.num_successful_transfers == 0 + and len(self.data["num_failed_transfers"]) == 0 + and len(self.data["num_failed_notifications"]) == 0 + and len(self.data["num_kv_expired_reqs"]) == 0 + ) + + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: + if not other.is_empty(): + for k, v in other.data.items(): + accumulator = self.data[k] + assert isinstance(accumulator, list) + accumulator.extend(v) + return self + + def reduce(self) -> dict[str, int | float]: + # Compute compact representative stats suitable for CLI logging + if self.num_successful_transfers == 0: + # CLI logging only reports successful transfers stats. If all requests in + # the interval were unsuccessful, Prom will report failures stats instead. + return { + "Num successful transfers": 0, + "Avg xfer time (ms)": 0, + "P90 xfer time (ms)": 0, + "Avg post time (ms)": 0, + "P90 post time (ms)": 0, + "Avg MB per transfer": 0, + "Throughput (MB/s)": 0, + "Avg number of descriptors": 0, + } + + xfer_time = np.asarray(self.data["transfer_duration"]) + post_time = np.asarray(self.data["post_duration"]) + # Convert to MB for CLI logging. + mb = np.asarray(self.data["bytes_transferred"]) / 2**20 + descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32) + n = len(descs) + assert n == self.num_successful_transfers + + total_mb = mb.sum() + avg_mb = total_mb / n + + total_time_seconds = xfer_time.sum() + throughput_mb_s = total_mb / total_time_seconds + + return { + "Num successful transfers": n, + "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3), + "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), + "Avg post time (ms)": round(post_time.mean() * 1e3, 3), + "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3), + "Avg MB per transfer": round(avg_mb, 3), + "Throughput (MB/s)": round(throughput_mb_s, 3), + "Avg number of descriptors": round(descs.mean(), 1), + } + + @property + def num_successful_transfers(self) -> int: + return len(self.data["transfer_duration"]) + + +class NixlPromMetrics(KVConnectorPromMetrics): + def __init__( + self, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ): + super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) + + buckets = [ + 0.001, + 0.005, + 0.01, + 0.025, + 0.05, + 0.075, + 0.1, + 0.2, + 0.3, + 0.5, + 0.75, + 1.0, + 5.0, + ] + nixl_histogram_xfer_time = self._histogram_cls( + name="vllm:nixl_xfer_time_seconds", + documentation="Histogram of transfer duration for NIXL KV Cache transfers.", + buckets=buckets[1:], + labelnames=labelnames, + ) + self.nixl_histogram_xfer_time = create_metric_per_engine( + nixl_histogram_xfer_time, self.per_engine_labelvalues + ) + nixl_histogram_post_time = self._histogram_cls( + name="vllm:nixl_post_time_seconds", + documentation="Histogram of transfer post time for NIXL KV" + " Cache transfers.", + buckets=buckets, + labelnames=labelnames, + ) + self.nixl_histogram_post_time = create_metric_per_engine( + nixl_histogram_post_time, self.per_engine_labelvalues + ) + # uniform 2kb to 16gb range + buckets = [2 ** (10 + i) for i in range(1, 25, 2)] + nixl_histogram_bytes_transferred = self._histogram_cls( + name="vllm:nixl_bytes_transferred", + documentation="Histogram of bytes transferred per NIXL KV Cache transfers.", + buckets=buckets, + labelnames=labelnames, + ) + self.nixl_histogram_bytes_transferred = create_metric_per_engine( + nixl_histogram_bytes_transferred, self.per_engine_labelvalues + ) + buckets = [ + 10, + 20, + 30, + 50, + 75, + 100, + 200, + 400, + 1000, + 2000, + 4000, + 10000, + 20000, + 50000, + ] + nixl_histogram_num_descriptors = self._histogram_cls( + name="vllm:nixl_num_descriptors", + documentation="Histogram of number of descriptors per NIXL" + " KV Cache transfers.", + buckets=buckets, + labelnames=labelnames, + ) + self.nixl_histogram_num_descriptors = create_metric_per_engine( + nixl_histogram_num_descriptors, self.per_engine_labelvalues + ) + counter_nixl_num_failed_transfers = self._counter_cls( + name="vllm:nixl_num_failed_transfers", + documentation="Number of failed NIXL KV Cache transfers.", + labelnames=labelnames, + ) + self.counter_nixl_num_failed_transfers = create_metric_per_engine( + counter_nixl_num_failed_transfers, self.per_engine_labelvalues + ) + counter_nixl_num_failed_notifications = self._counter_cls( + name="vllm:nixl_num_failed_notifications", + documentation="Number of failed NIXL KV Cache notifications.", + labelnames=labelnames, + ) + self.counter_nixl_num_failed_notifications = create_metric_per_engine( + counter_nixl_num_failed_notifications, self.per_engine_labelvalues + ) + + counter_nixl_num_kv_expired_reqs = self._counter_cls( + name="vllm:nixl_num_kv_expired_reqs", + documentation="Number of requests that had their KV expire. " + "NOTE: This metric is tracked on the P instance.", + labelnames=labelnames, + ) + self.counter_nixl_num_kv_expired_reqs = create_metric_per_engine( + counter_nixl_num_kv_expired_reqs, self.per_engine_labelvalues + ) + + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): + for prom_obj, list_item_key in zip( + [ + self.nixl_histogram_xfer_time, + self.nixl_histogram_post_time, + self.nixl_histogram_bytes_transferred, + self.nixl_histogram_num_descriptors, + ], + [ + "transfer_duration", + "post_duration", + "bytes_transferred", + "num_descriptors", + ], + ): + for list_item in transfer_stats_data[list_item_key]: + prom_obj[engine_idx].observe(list_item) + for counter_obj, counter_item_key in zip( + [ + self.counter_nixl_num_failed_transfers, + self.counter_nixl_num_failed_notifications, + self.counter_nixl_num_kv_expired_reqs, + ], + ["num_failed_transfers", "num_failed_notifications", "num_kv_expired_reqs"], + ): + for list_item in transfer_stats_data[counter_item_key]: + counter_obj[engine_idx].inc(list_item) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py new file mode 100644 index 00000000000..514214347ae --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared constants, lazy imports and helpers for the NIXL connector.""" + +import contextlib +import os +import sys +from collections.abc import Iterator +from typing import Any + +import zmq + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.network_utils import make_zmq_socket + +logger = init_logger(__name__) + + +# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used +try: + if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ: + # avoid a memory leak in UCX when using NIXL on some models + # see: https://github.com/vllm-project/vllm/issues/24264 + if "nixl" in sys.modules or "rixl" in sys.modules: + logger.warning( + "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. " + "Please set it to '1024' manually." + ) + else: + logger.info( + "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " + "memory leak in UCX when using NIXL." + ) + os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" + + if not current_platform.is_rocm(): + from nixl._api import nixl_agent as NixlWrapper + from nixl._bindings import nixlXferTelemetry + else: + from rixl._api import nixl_agent as NixlWrapper + from rixl._bindings import nixlXferTelemetry + + logger.info("NIXL is available") +except ImportError: + logger.warning("NIXL is not available") + NixlWrapper = None + nixlXferTelemetry = None + + +try: + if not current_platform.is_rocm(): + from nixl._api import nixl_agent_config + else: + from rixl._api import nixl_agent_config +except ImportError: + nixl_agent_config = None + logger.warning("NIXL agent config is not available") + +# Supported platforms and types of kv transfer buffer. +# {device: tuple of supported kv buffer types} +_NIXL_SUPPORTED_DEVICE = { + "cuda": ( + "cuda", + "cpu", + ), + "tpu": ("cpu",), + "xpu": ( + "cpu", + "xpu", + ), + "cpu": ("cpu",), +} +# support for oot platform by providing mapping in current_platform +_NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices()) + + +# TODO: merge with vllm.utils.network_utils.zmq_socket_ctx +@contextlib.contextmanager +def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: + """Context manager for a ZMQ socket""" + + if socket_type not in (zmq.ROUTER, zmq.REQ): + raise ValueError(f"Unexpected socket type: {socket_type}") + + ctx: zmq.Context | None = None + try: + ctx = zmq.Context() # type: ignore[attr-defined] + yield make_zmq_socket( + ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER + ) + finally: + if ctx is not None: + ctx.destroy(linger=0) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py similarity index 69% rename from vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py rename to vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index c575043fb34..45aa33033e7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -1,18 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import contextlib -import copy +"""Worker-side logic for the NIXL connector.""" + import logging import os import queue -import sys import threading import time import uuid from collections import defaultdict -from collections.abc import Iterator from concurrent.futures import Future, ThreadPoolExecutor -from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast import msgspec @@ -21,32 +18,36 @@ import torch import zmq from vllm import envs -from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.utils import ( BlockIds, EngineId, HeteroTPTransferConfig, TpKVTopology, - get_current_attn_backend, get_current_attn_backends, kv_postprocess_blksize_and_layout_on_receive, kv_postprocess_blksize_on_receive, kv_postprocess_layout_on_receive, - yield_req_data, ) -from vllm.distributed.kv_transfer.kv_connector.v1.base import ( - CopyBlocksOp, - KVConnectorBase_V1, - KVConnectorHandshakeMetadata, - KVConnectorMetadata, - KVConnectorRole, - SupportsHMA, +from vllm.distributed.kv_transfer.kv_connector.v1.base import CopyBlocksOp +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + GET_META_MSG, + NixlAgentMetadata, + NixlConnectorMetadata, + NixlHandshakePayload, + ReqId, + ReqMeta, + TransferHandle, + compute_nixl_compatibility_hash, ) -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( - KVConnectorPromMetrics, - KVConnectorStats, - PromMetric, - PromMetricT, +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( + NixlKVConnectorStats, +) +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( + _NIXL_SUPPORTED_DEVICE, + NixlWrapper, + nixl_agent_config, + zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( MambaConvSplitInfo, @@ -57,960 +58,34 @@ from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from vllm.forward_context import ForwardContext from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.platforms import current_platform -from vllm.utils.math_utils import cdiv -from vllm.utils.network_utils import make_zmq_path, make_zmq_socket -from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.utils.network_utils import make_zmq_path from vllm.v1.attention.backends.utils import get_kv_cache_layout -from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, MambaSpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, ) -from vllm.v1.metrics.utils import create_metric_per_engine from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size if TYPE_CHECKING: - from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.config import VllmConfig from vllm.v1.kv_cache_interface import KVCacheConfig - from vllm.v1.request import Request - -TransferHandle = int -ReqId = str - -# -# NIXL Connector Version -# -# Increment this version whenever there is an incompatible change to: -# - NixlAgentMetadata schema -# - kv_transfer_params schema or semantics -# - NIXL transfer protocol or wire format -# - KV cache memory layout or block organization -# - Any other change that breaks P/D interoperability -# -# Version History: -# 1: Initial version with compatibility checking -# 2: Add remote_request_id to kv_transfer_params -# -NIXL_CONNECTOR_VERSION: int = 2 - -GET_META_MSG = b"get_meta_msg" logger = init_logger(__name__) -# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used -try: - if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ: - # avoid a memory leak in UCX when using NIXL on some models - # see: https://github.com/vllm-project/vllm/issues/24264 - if "nixl" in sys.modules or "rixl" in sys.modules: - logger.warning( - "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. " - "Please set it to '1024' manually." - ) - else: - logger.info( - "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " - "memory leak in UCX when using NIXL." - ) - os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" - - if not current_platform.is_rocm(): - from nixl._api import nixl_agent as NixlWrapper - from nixl._bindings import nixlXferTelemetry - else: - from rixl._api import nixl_agent as NixlWrapper - from rixl._bindings import nixlXferTelemetry - - logger.info("NIXL is available") -except ImportError: - logger.warning("NIXL is not available") - NixlWrapper = None - nixlXferTelemetry = None - - -try: - if not current_platform.is_rocm(): - from nixl._api import nixl_agent_config - else: - from rixl._api import nixl_agent_config -except ImportError: - nixl_agent_config = None - logger.warning("NIXL agent config is not available") - -# Supported platforms and types of kv transfer buffer. -# {device: tuple of supported kv buffer types} -_NIXL_SUPPORTED_DEVICE = { - "cuda": ( - "cuda", - "cpu", - ), - "tpu": ("cpu",), - "xpu": ( - "cpu", - "xpu", - ), - "cpu": ("cpu",), -} -# support for oot platform by providing mapping in current_platform -_NIXL_SUPPORTED_DEVICE.update(current_platform.get_nixl_supported_devices()) - - -@dataclass -class NixlAgentMetadata: - engine_id: str - agent_metadata: bytes - kv_caches_base_addr: list[int] - device_id: int - num_blocks: int - block_lens: list[int] - kv_cache_layout: str - block_size: int - ssm_sizes: tuple[int, int] - - -@dataclass -class NixlHandshakePayload(KVConnectorHandshakeMetadata): - """ - Wrapper for NIXL handshake sent over the wire. - - Enables two-phase decoding for graceful compatibility checking: - 1. Decode NixlHandshakePayload to get compatibility_hash - 2. Compute local hash and compare - 3. Only if hashes match, decode agent_metadata_bytes - - This prevents decoder errors when NixlAgentMetadata schema is - incompatible, allowing graceful failure with clear error message. - """ - - compatibility_hash: str - agent_metadata_bytes: bytes # NixlAgentMetadata encoded - - -def compute_nixl_compatibility_hash( - vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool -) -> str: - """ - Compute compatibility hash for NIXL KV transfer. - - Hash only the factors that affect whether two NIXL instances can - successfully transfer KV cache data. - - Factors included: - - vLLM version and NIXL connector version - - Model architecture (name, dtype, KV heads, layers) - - KV cache format (dtype, sliding window) - - Attention backend - - Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout - are validated at runtime in _validate_remote_agent_handshake and are not - included in this hash to support heterogeneous deployments. - - Note - the set of factors are likely to evolve significantly over - time to be more or less permissive. - - Returns: - SHA-256 hex digest - """ - from vllm import __version__ as vllm_version - from vllm.config.utils import hash_factors - - model_config = vllm_config.model_config - cache_config = vllm_config.cache_config - is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager - - factors = { - # Version compatibility - "vllm_version": vllm_version, - "nixl_connector_version": NIXL_CONNECTOR_VERSION, - # Model architecture - affects KV cache shape - "model": model_config.model, - "dtype": str(model_config.dtype), - "num_kv_heads": model_config.get_total_num_kv_heads(), - "head_size": model_config.get_head_size(), - "num_hidden_layers": model_config.get_total_num_hidden_layers(), - # Attention backend and KV cache dtype affect memory layout - "attn_backend_name": attn_backend_name, - "cache_dtype": str(cache_config.cache_dtype), - "cross_layers_blocks": cross_layers_blocks, - "is_hma_enabled": is_hma_enabled, - } - - compat_hash = hash_factors(factors) - logger.debug( - "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, " - "cache_dtype=%s, attn_backend=%s)", - compat_hash, - factors["model"], - factors["dtype"], - factors["num_kv_heads"], - factors["cache_dtype"], - attn_backend_name, - ) - return compat_hash - - -@dataclass -class RemoteMeta: - block_ids: BlockIds - host: str - port: int - engine_id: str - request_id: str - - -@dataclass -class ReqMeta: - local_block_ids: BlockIds - # To be used when logical block size does not match the kernel block size - local_physical_block_ids: BlockIds - tp_size: int - remote: RemoteMeta | None = None - - -class NixlConnectorMetadata(KVConnectorMetadata): - def __init__(self): - self.reqs_to_recv: dict[ReqId, ReqMeta] = {} - self.reqs_to_save: dict[ReqId, ReqMeta] = {} - self.reqs_to_send: dict[ReqId, float] = {} - self.reqs_in_batch: set[ReqId] = set() - self.reqs_not_processed: set[ReqId] = set() - - def _add_new_req( - self, - local_block_ids: BlockIds, - kv_transfer_params: dict[str, Any], - ) -> ReqMeta: - return ReqMeta( - local_block_ids=local_block_ids, - local_physical_block_ids=local_block_ids, - # P workers don't need to receive tp_size from proxy here. - tp_size=kv_transfer_params.get("tp_size", 1), - ) - - def add_new_req_to_save( - self, - request_id: ReqId, - local_block_ids: BlockIds, - kv_transfer_params: dict[str, Any], - ): - self.reqs_to_save[request_id] = self._add_new_req( - local_block_ids, kv_transfer_params - ) - - def add_new_req_to_recv( - self, - request_id: ReqId, - local_block_ids: BlockIds, - kv_transfer_params: dict[str, Any], - ): - req = self._add_new_req(local_block_ids, kv_transfer_params) - req.remote = RemoteMeta( - block_ids=kv_transfer_params["remote_block_ids"], - engine_id=kv_transfer_params["remote_engine_id"], - request_id=kv_transfer_params["remote_request_id"], - host=kv_transfer_params["remote_host"], - port=kv_transfer_params["remote_port"], - ) - self.reqs_to_recv[request_id] = req - - -class NixlConnector(KVConnectorBase_V1, SupportsHMA): - @property - def prefer_cross_layer_blocks(self) -> bool: - if any( - [ - isinstance(group.kv_cache_spec, MambaSpec) - for group in self.kv_cache_config.kv_cache_groups - ] - ): - # Hybrid SSM models do not yet support cross-layer layout - return False - - backend = get_current_attn_backend(self._vllm_config) - if backend.get_name() not in ( - "FLASH_ATTN", - "FLASHINFER", - "TRITON_ATTN", - ): - return False - - # For now there is no benefit to run cross layers when backend - # does not support on HND - if get_kv_cache_layout() != "HND": - return False - - extra_config = self.kv_transfer_config.kv_connector_extra_config - return ( - str(extra_config.get("enable_cross_layers_blocks", "False")).lower() - == "true" - ) - - def __init__( - self, - vllm_config: VllmConfig, - role: KVConnectorRole, - kv_cache_config: "KVCacheConfig", - ): - super().__init__(vllm_config, role, kv_cache_config) - assert vllm_config.kv_transfer_config is not None - assert vllm_config.kv_transfer_config.engine_id is not None - self.kv_cache_config = kv_cache_config - self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id - self.kv_transfer_config = vllm_config.kv_transfer_config - if role == KVConnectorRole.SCHEDULER: - self.connector_scheduler: NixlConnectorScheduler | None = ( - NixlConnectorScheduler(vllm_config, self.engine_id, kv_cache_config) - ) - self.connector_worker: NixlConnectorWorker | None = None - elif role == KVConnectorRole.WORKER: - self.connector_scheduler = None - self.connector_worker = NixlConnectorWorker( - vllm_config, self.engine_id, kv_cache_config - ) - - ############################################################ - # Class Methods - ############################################################ - @classmethod - def get_required_kvcache_layout(cls, vllm_config: VllmConfig): - if vllm_config.model_config is None: - logger.warning_once( - "Unable to detect current VLLM config. " - "Fallback to default kv cache layout." - ) - return None - use_mla = vllm_config.model_config.use_mla - if use_mla: - # return None when we have mla - # as the layout should not matter in that case, - # which fallback to the default behavior. - return None - logger.info_once( - "NixlConnector setting KV cache layout to HND for better xfer performance." - ) - return "HND" - - ############################################################ - # Scheduler Side Methods - ############################################################ - - def get_num_new_matched_tokens( - self, request: "Request", num_computed_tokens: int - ) -> tuple[int | None, bool]: - assert self.connector_scheduler is not None - return self.connector_scheduler.get_num_new_matched_tokens( - request, num_computed_tokens - ) - - def update_state_after_alloc( - self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int - ): - assert self.connector_scheduler is not None - return self.connector_scheduler.update_state_after_alloc( - request, blocks, num_external_tokens - ) - - def build_connector_meta( - self, - scheduler_output: SchedulerOutput, - ) -> KVConnectorMetadata: - assert self.connector_scheduler is not None - return self.connector_scheduler.build_connector_meta(scheduler_output) - - def request_finished( - self, - request: "Request", - block_ids: list[int], - ) -> tuple[bool, dict[str, Any] | None]: - assert self.connector_scheduler is not None - return self.connector_scheduler.request_finished(request, (block_ids,)) - - def request_finished_all_groups( - self, - request: "Request", - block_ids: tuple[list[int], ...], - ) -> tuple[bool, dict[str, Any] | None]: - assert self.connector_scheduler is not None - return self.connector_scheduler.request_finished(request, block_ids) - - def set_xfer_handshake_metadata( - self, metadata: dict[int, KVConnectorHandshakeMetadata] - ) -> None: - """ - Set the KV connector handshake metadata for this connector. - - Args: - metadata (dict): the handshake metadata to set. - """ - assert self.connector_scheduler is not None - self.connector_scheduler.set_xfer_handshake_metadata(metadata) - - ############################################################ - # Worker Side Methods - ############################################################ - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - assert self.connector_worker is not None - self.connector_worker.register_kv_caches(kv_caches) - - def register_cross_layers_kv_cache( - self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] - ): - assert self.connector_worker is not None - self.connector_worker.register_cross_layers_kv_caches(kv_cache) - - def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): - assert self.connector_worker is not None - self.connector_worker.set_host_xfer_buffer_ops(copy_operation) - - def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: - """Get the finished recving and sending requests.""" - assert self.connector_worker is not None - return self.connector_worker.get_finished() - - def get_block_ids_with_load_errors(self) -> set[int]: - """Get block IDs that failed to load via NIXL.""" - assert self.connector_worker is not None - return self.connector_worker.get_block_ids_with_load_errors() - - def get_kv_connector_stats(self) -> KVConnectorStats | None: - if self.connector_worker is None: - return None - return self.connector_worker.get_kv_connector_stats() - - @classmethod - def build_kv_connector_stats( - cls, data: dict[str, Any] | None = None - ) -> KVConnectorStats | None: - return ( - NixlKVConnectorStats(data=data) - if data is not None - else NixlKVConnectorStats() - ) - - @classmethod - def build_prom_metrics( - cls, - vllm_config: VllmConfig, - metric_types: dict[type[PromMetric], type[PromMetricT]], - labelnames: list[str], - per_engine_labelvalues: dict[int, list[object]], - ) -> KVConnectorPromMetrics: - return NixlPromMetrics( - vllm_config, metric_types, labelnames, per_engine_labelvalues - ) - - def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: - assert self.connector_worker is not None - assert isinstance(self._connector_metadata, NixlConnectorMetadata) - self.connector_worker.start_load_kv(self._connector_metadata) - - def wait_for_layer_load(self, layer_name: str) -> None: - """NixlConnector does not do layerwise saving.""" - pass - - def save_kv_layer( - self, - layer_name: str, - kv_layer: torch.Tensor, - attn_metadata: AttentionMetadata, - **kwargs, - ) -> None: - """NixlConnector does not save explicitly.""" - pass - - def wait_for_save(self): - assert self.connector_worker is not None - assert isinstance(self._connector_metadata, NixlConnectorMetadata) - if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks: - self.connector_worker.save_kv_to_host(self._connector_metadata) - - def shutdown(self): - if self.connector_worker is not None: - self.connector_worker.shutdown() - if self.connector_scheduler is not None: - self.connector_scheduler.shutdown() - - def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: - """ - Get the KVConnector handshake metadata for this connector. - This metadata is used for out-of-band connector handshake - between P/D workers. - - Returns: - KVConnectorHandshakeMetadata: the handshake metadata. - None if no handshake metadata is available. - """ - assert self.connector_worker is not None - return self.connector_worker.xfer_handshake_metadata - - -class NixlConnectorScheduler: - """Implementation of Scheduler side methods""" - - def __init__( - self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: "KVCacheConfig" - ): - self.vllm_config = vllm_config - self.block_size = vllm_config.cache_config.block_size - self.engine_id: EngineId = engine_id - self.kv_cache_config = kv_cache_config - self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST - self.side_channel_port = ( - envs.VLLM_NIXL_SIDE_CHANNEL_PORT - + vllm_config.parallel_config.data_parallel_index - ) - assert vllm_config.kv_transfer_config is not None - if current_platform.device_type == "cpu": - self.use_host_buffer = False - else: - self.use_host_buffer = ( - vllm_config.kv_transfer_config.kv_buffer_device == "cpu" - ) - self._is_hma_required = ( - not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager - # Also handle unlikely SW-only model case instead of checking num_groups>1. - and any( - not isinstance(g.kv_cache_spec, FullAttentionSpec) - for g in kv_cache_config.kv_cache_groups - ) - ) - self._has_mamba = any( - isinstance(g.kv_cache_spec, MambaSpec) - for g in kv_cache_config.kv_cache_groups - ) - - logger.info("Initializing NIXL Scheduler %s", engine_id) - if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager: - logger.info("Hybrid Memory Allocator is enabled with NIXL") - - # Background thread for handling new handshake requests. - self._nixl_handshake_listener_t: threading.Thread | None = None - self._stop_event = threading.Event() - - # Requests that need to start recv/send. - # New requests are added by update_state_after_alloc in - # the scheduler. Used to make metadata passed to Worker. - self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {} - self._reqs_need_save: dict[ReqId, Request] = {} - # Reqs to send and their expiration time - self._reqs_need_send: dict[ReqId, float] = {} - self._reqs_in_batch: set[ReqId] = set() - # Reqs to remove from processed set because they're not to send after - # remote prefill or aborted. - self._reqs_not_processed: set[ReqId] = set() - - # Gather Sliding Window sizes for each kv cache group (if any) in number of - # blocks per KV cache group. This is used to clip the local attention window. - sw_sizes_tokens: list[tuple[int, int]] = [ - (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size) - if isinstance(g.kv_cache_spec, SlidingWindowSpec) - else (0, self.block_size) - for g in kv_cache_config.kv_cache_groups - ] - # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively - # account for boundary overlap eg window isn't fully aligned with blocks. - self.blocks_per_sw = [ - cdiv(n_tokens, block_size) + 1 if n_tokens else 0 - for n_tokens, block_size in sw_sizes_tokens - ] - - def shutdown(self): - self._stop_event.set() - if self._nixl_handshake_listener_t is not None: - self._nixl_handshake_listener_t.join() - self._nixl_handshake_listener_t = None - - def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds: - """ - Clip the number of blocks to the sliding window size for each kv cache group - that employs SWA. - This is necessary because the KV Cache manager initially allocates blocks for - the entire sequence length, and successively cleans up blocks that are outside - the window prior to the `request_finished_all_groups` hook. - """ - if len(block_ids) == 0 or not self._is_hma_required: - # No blocks to clip eg Full prefix cache hit or not a hybrid model. - return block_ids - # NOTE (NickLucche) This logic is currently handled at the connector level - # because offloading connectors might want to receive the whole sequence even - # for SWA groups. We will abstract this logic once the interface is more stable - assert len(block_ids) == len(self.blocks_per_sw), ( - "Number of KV cache groups must match" - ) - # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged - return tuple( - [ - blocks[-self.blocks_per_sw[i] :] - if self.blocks_per_sw[i] > 0 - else blocks - for i, blocks in enumerate(block_ids) - ] - ) - - def set_xfer_handshake_metadata( - self, metadata: dict[int, KVConnectorHandshakeMetadata] - ) -> None: - """ - Set the KV connector handshake metadata for this connector. - - Args: - metadata (dict): the handshake metadata to set. - """ - encoded_data: dict[int, bytes] = {} - encoder = msgspec.msgpack.Encoder() - for tp_rank, rank_metadata in metadata.items(): - if not isinstance(rank_metadata, NixlHandshakePayload): - raise ValueError( - "NixlConnectorScheduler expects NixlHandshakePayload for " - "handshake metadata." - ) - encoded_data[tp_rank] = encoder.encode(rank_metadata) - logger.debug( - "Tp rank %d: encoded NixlHandshakePayload size: %s bytes", - tp_rank, - str(len(encoded_data[tp_rank])), - ) - - # Only start the listener when we have metadata to serve. - if self._nixl_handshake_listener_t is None: - ready_event = threading.Event() - self._nixl_handshake_listener_t = threading.Thread( - target=self._nixl_handshake_listener, - args=( - encoded_data, - ready_event, - self._stop_event, - self.side_channel_port, - ), - daemon=True, - name="nixl_handshake_listener", - ) - self._nixl_handshake_listener_t.start() - ready_event.wait() # Wait for listener ZMQ socket to be ready. - - @staticmethod - def _nixl_handshake_listener( - encoded_data: dict[int, Any], - ready_event: threading.Event, - stop_event: threading.Event, - port: int, - ): - """Background thread for getting new NIXL handshakes.""" - # NOTE(rob): this is a simple implementation. We will move - # to a better approach via HTTP endpoint soon. - - # Listen for new requests for metadata. - host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST - path = make_zmq_path("tcp", host, port) - logger.debug("Starting listening on path: %s", path) - with zmq_ctx(zmq.ROUTER, path) as sock: - sock.setsockopt(zmq.RCVTIMEO, 1000) - ready_event.set() - while True: - try: - identity, _, msg = sock.recv_multipart() - except zmq.Again: - if stop_event.is_set(): - break - continue - # Decode the message which contains (GET_META_MSG, rank) - msg, target_tp_rank = msgspec.msgpack.decode(msg) - logger.debug( - "Received message for tp rank %s", - target_tp_rank, - ) - if msg != GET_META_MSG: - logger.warning("Connection listener got unexpected message %s", msg) - sock.send_multipart((identity, b"", encoded_data[target_tp_rank])) - - def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int: - """D-side only. Returns N-1 for Mamba models since the decoder - always recomputes the last token and must start from h(N-1).""" - if self._has_mamba and num_prompt_tokens > 1: - return num_prompt_tokens - 1 - return num_prompt_tokens - - def _truncate_mamba_request_for_prefill(self, request: "Request") -> None: - """P-side only: drop the last prompt token so the prefiller computes - h(N-1) instead of h(N). The decoder recomputes the last token to - derive h(N) correctly. - - Guarded by ``_p_side_truncated`` to avoid repeated truncation if the - request is preempted and rescheduled.""" - params = request.kv_transfer_params - if ( - params is not None - # Guard against repeated truncation after preemption/reschedule. - and not params.get("_p_side_truncated") - and request.num_prompt_tokens > 1 - ): - if request.prompt_token_ids is not None: - request.prompt_token_ids.pop() - elif request.prompt_embeds is not None: - request.prompt_embeds = request.prompt_embeds[:-1] - else: - return - - request._all_token_ids.pop() - request.num_prompt_tokens -= 1 - request.max_tokens = 1 - params["_p_side_truncated"] = True - - def get_num_new_matched_tokens( - self, request: "Request", num_computed_tokens: int - ) -> tuple[int, bool]: - """ - For remote prefill, pull all prompt blocks from remote - asynchronously relative to engine execution. - - Args: - request (Request): the request object. - num_computed_tokens (int): the number of locally - computed tokens for this request - Returns: - * the number of tokens that can be loaded from the - external KV cache beyond what is already computed. - * true if the external KV cache tokens will be loaded - asynchronously (between scheduler steps). - """ - - params = request.kv_transfer_params - logger.debug( - "NIXLConnector get_num_new_matched_tokens: " - "num_computed_tokens=%s, kv_transfer_params=%s", - num_computed_tokens, - params, - ) - - if params is not None and params.get("do_remote_prefill"): - # Remote prefill: get all prompt blocks from remote. - token_ids = request.prompt_token_ids or [] - actual = self._mamba_prefill_token_count(len(token_ids)) - count = actual - num_computed_tokens - if count > 0: - return count, True - - if params is not None and params.get("do_remote_decode") and self._has_mamba: - self._truncate_mamba_request_for_prefill(request) - - # No remote prefill for this request. - return 0, False - - def update_state_after_alloc( - self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int - ): - params = request.kv_transfer_params - logger.debug( - "NIXLConnector update_state_after_alloc: " - "num_external_tokens=%s, kv_transfer_params=%s", - num_external_tokens, - params, - ) - - if not params: - return - - if params.get("do_remote_decode"): - self._reqs_in_batch.add(request.request_id) - if self.use_host_buffer and params.get("do_remote_decode"): - # NOTE: when accelerator is not directly supported by Nixl, - # prefilled blocks need to be saved to host memory before transfer. - self._reqs_need_save[request.request_id] = request - elif params.get("do_remote_prefill"): - if params.get("remote_block_ids"): - if all( - p in params - for p in ( - "remote_engine_id", - "remote_request_id", - "remote_host", - "remote_port", - ) - ): - # If remote_blocks and num_external_tokens = 0, we have - # a full prefix cache hit on the D worker. We need to call - # send_notif in _read_blocks to free the memory on the P. - - unhashed_local_block_ids: BlockIds = ( - blocks.get_unhashed_block_ids_all_groups() - if num_external_tokens > 0 - else () - ) - local_block_ids = self.get_sw_clipped_blocks( - unhashed_local_block_ids - ) - - # Get unhashed blocks to pull from remote. Mind that a full prefix - # cache hit is indicated with an empty list. - self._reqs_need_recv[request.request_id] = ( - request, - local_block_ids, - ) - - else: - logger.warning( - "Got invalid KVTransferParams: %s. This " - "request will not utilize KVTransfer", - params, - ) - else: - assert num_external_tokens == 0 - # Only trigger 1 KV transfer per request. - params["do_remote_prefill"] = False - - def _build_save_meta( - self, - meta: NixlConnectorMetadata, - scheduler_output: SchedulerOutput, - ) -> None: - # only called when use_host_buffer is True to build the save metadata - - # NOTE: For the prefill side, there might be a chance that an early added - # request is a chunked prefill, so we need to check if new blocks are added - for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output): - req_to_save = self._reqs_need_save.get(req_id) - if req_to_save is None or new_block_id_groups is None: - continue - req = req_to_save - - assert req.kv_transfer_params is not None - clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups) - meta.add_new_req_to_save( - request_id=req_id, - local_block_ids=clipped_block_id_groups, - kv_transfer_params=req.kv_transfer_params, - ) - assert scheduler_output.num_scheduled_tokens is not None - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] - is_partial = ( - req.num_computed_tokens + num_scheduled_tokens - ) < req.num_prompt_tokens - if not is_partial: - # For non-partial prefills, once new req_meta is scheduled, it - # can be removed from _reqs_need_save. - # For partial prefill case, we will retain the request in - # _reqs_need_save until all blocks are scheduled with req_meta. - # Therefore, only pop if `not is_partial`. - self._reqs_need_save.pop(req_id) - - def build_connector_meta( - self, - scheduler_output: SchedulerOutput, - ) -> KVConnectorMetadata: - meta = NixlConnectorMetadata() - - # Loop through scheduled reqs and convert to ReqMeta. - for req_id, (req, block_ids) in self._reqs_need_recv.items(): - assert req.kv_transfer_params is not None - meta.add_new_req_to_recv( - request_id=req_id, - local_block_ids=block_ids, - kv_transfer_params=req.kv_transfer_params, - ) - - if self.use_host_buffer: - self._build_save_meta(meta, scheduler_output) - - meta.reqs_to_send = self._reqs_need_send - meta.reqs_in_batch = self._reqs_in_batch - meta.reqs_not_processed = self._reqs_not_processed - - # Clear the list once workers start the transfers - self._reqs_need_recv.clear() - self._reqs_in_batch = set() - self._reqs_not_processed = set() - self._reqs_need_send = {} - - return meta - - def request_finished( - self, - request: "Request", - block_ids: BlockIds, - ) -> tuple[bool, dict[str, Any] | None]: - """ - Once a request is finished, determine whether request blocks - should be freed now or will be sent asynchronously and freed later. - """ - from vllm.v1.request import RequestStatus - - params = request.kv_transfer_params - logger.debug( - "NIXLConnector request_finished(%s), request_status=%s, " - "kv_transfer_params=%s", - request.request_id, - request.status, - params, - ) - if not params: - return False, None - - if params.get("do_remote_prefill"): - # If do_remote_prefill is still True when the request is finished, - # update_state_after_alloc must not have been called (the request - # must have been aborted before it was scheduled). - # To avoid stranding the prefill blocks in the prefill instance, - # we must add empty block_ids to _reqs_need_recv so that our - # worker side will notify and free blocks in the prefill instance. - self._reqs_need_recv[request.request_id] = (request, []) - params["do_remote_prefill"] = False - return False, None - - if not params.get("do_remote_decode"): - return False, None - if request.status != RequestStatus.FINISHED_LENGTH_CAPPED: - # Also include the case of a P/D Prefill request with immediate - # block free (eg abort). Stop tracking this request. - self._reqs_not_processed.add(request.request_id) - # Clear _reqs_need_save if a request is aborted as partial prefill. - self._reqs_need_save.pop(request.request_id, None) - return False, None - - # TODO: check whether block_ids actually ever be 0. If not we could - # remove the conditional below - delay_free_blocks = any(len(group) > 0 for group in block_ids) - - if delay_free_blocks: - # Prefill request on remote. It will be read from D upon completion - logger.debug( - "NIXLConnector request_finished(%s) waiting for %d seconds " - "for remote decode to fetch blocks", - request.request_id, - envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT, - ) - self._reqs_need_send[request.request_id] = ( - time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT - ) - # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones), - # trimming down after allocating for the whole sequence length. Empty - # blocks are always at the start of the list. - # Here we "unpad" blocks to send the actual remote blocks to be read. - block_ids = self.get_sw_clipped_blocks(block_ids) - - return delay_free_blocks, dict( - do_remote_prefill=True, - do_remote_decode=False, - remote_block_ids=block_ids, - remote_engine_id=self.engine_id, - remote_request_id=request.request_id, - remote_host=self.side_channel_host, - remote_port=self.side_channel_port, - tp_size=self.vllm_config.parallel_config.tensor_parallel_size, - ) - class NixlConnectorWorker: """Implementation of Worker side methods""" def __init__( - self, vllm_config: VllmConfig, engine_id: str, kv_cache_config: "KVCacheConfig" + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", ): if NixlWrapper is None: logger.error("NIXL is not available") @@ -1116,6 +191,7 @@ class NixlConnectorWorker: self.num_blocks = kv_cache_config.num_blocks self.enable_permute_local_kv = False + self.enable_heterogeneous_attn_post_process = False # KV Caches and nixl tracking data. self.device_type = current_platform.device_type @@ -1776,6 +852,7 @@ class NixlConnectorWorker: else self.host_buffer_kv_cache_layout, block_size=self.block_size, ssm_sizes=self._mamba_ssm_size, + attn_backend_name=self.backend_name, ) # Wrap metadata in payload with hash for defensive decoding assert self.compat_hash is not None @@ -2369,6 +1446,21 @@ class NixlConnectorWorker: "Or enable experimental feature to use HND to NHD support by " "setting 'enable_permute_local_kv'=True in --kv-transfer-config." ) + # if remote_agent used attn is not same as local, + # hint heterogenuous attn post process + if ( + nixl_agent_meta.attn_backend_name != self.backend_name + and self.backend_name in ["CPU_ATTN"] + ): + if self._is_hma_required: + raise RuntimeError( + "heterogeneous attn post process is not supported with HMA" + ) + logger.info( + "[Experimental] CPU_ATTN backend is used, " + "hint heterogeneous attn post process" + ) + self.enable_heterogeneous_attn_post_process = True # Heterogeneous TP requires head-splitting, which only works with # HND layout. MLA and replicated-KV cases don't split on heads. @@ -2542,6 +1634,28 @@ class NixlConnectorWorker: cache, indices, block_size_ratio ) + def post_process_device_kv_on_receive_heterogeneous_attn( + self, block_ids: list[int] + ): + """ + Post process device kv cache after receiving from remote + for heterogeneous attention. + """ + assert self.enable_heterogeneous_attn_post_process + + indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long) + + for _, cache_or_caches in self.device_kv_caches.items(): + blocks_to_update = cache_or_caches.index_select(1, indices) + current_platform.pack_kv_cache( + key=blocks_to_update[0], + value=blocks_to_update[1], + key_cache=cache_or_caches[0], + value_cache=cache_or_caches[1], + block_ids=block_ids, + indices=indices, + ) + def get_finished(self) -> tuple[set[str], set[str]]: """ Get requests that are done sending or recving on this specific worker. @@ -2566,6 +1680,7 @@ class NixlConnectorWorker: ) block_ids_for_blocksize_post_process = defaultdict(list) + block_ids_for_heterogeneous_attn_post_process = list[list[int]]() for req_id in done_recving: # clean up metadata for completed requests meta = self._recving_metadata.pop(req_id, None) @@ -2585,12 +1700,20 @@ class NixlConnectorWorker: block_ids_for_blocksize_post_process[block_size_ratio].append( meta.local_physical_block_ids[0] ) + # post processing for heterogeneous attention + if self.enable_heterogeneous_attn_post_process: + block_ids_for_heterogeneous_attn_post_process.append( + meta.local_physical_block_ids[0] + ) for ( block_size_ratio, block_ids_list, ) in block_ids_for_blocksize_post_process.items(): self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list) + for block_ids in block_ids_for_heterogeneous_attn_post_process: + self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids) + # Handle timeout to avoid stranding blocks on remote. now = time.perf_counter() while self._reqs_to_send: @@ -2791,6 +1914,10 @@ class NixlConnectorWorker: meta.remote.block_ids, self._mamba_phys_ratio[meta.remote.engine_id], ) + else: + meta.remote.block_ids = self._logical_to_kernel_block_ids( + meta.remote.block_ids + ) # D may have to perform multiple reads from different remote ranks. for i, remote_rank in enumerate(remote_ranks): if self.use_mla and tp_ratio < 0 and i > 0: @@ -3151,9 +2278,9 @@ class NixlConnectorWorker: the their size differs. Reference diagram: KVCacheTensor (Shared) - / \ - / \ - / \ + / \\ + / \\ + / \\ Attention (FlashInfer) View Mamba View | | | | @@ -3234,266 +2361,3 @@ class NixlConnectorWorker: for desc in self._registered_descs: self.nixl_wrapper.deregister_memory(desc) self._registered_descs.clear() - - -@contextlib.contextmanager -def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: - """Context manager for a ZMQ socket""" - - if socket_type not in (zmq.ROUTER, zmq.REQ): - raise ValueError(f"Unexpected socket type: {socket_type}") - - ctx: zmq.Context | None = None - try: - ctx = zmq.Context() # type: ignore[attr-defined] - yield make_zmq_socket( - ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER - ) - finally: - if ctx is not None: - ctx.destroy(linger=0) - - -@dataclass -class NixlKVConnectorStats(KVConnectorStats): - """Container for transfer performance metrics""" - - def __post_init__(self): - if not self.data: - # Empty container init, no data is passed in. - self.reset() - - def reset(self): - # Must be serializable - self.data: dict[str, list[float | int]] = { - "transfer_duration": [], - "post_duration": [], - "bytes_transferred": [], - "num_descriptors": [], - "num_failed_transfers": [], - "num_failed_notifications": [], - "num_kv_expired_reqs": [], - } - - def record_transfer(self, res: nixlXferTelemetry): - # Keep metrics units consistent with rest of the code: time us->s - self.data["transfer_duration"].append(res.xferDuration / 1e6) - self.data["post_duration"].append(res.postDuration / 1e6) - self.data["bytes_transferred"].append(res.totalBytes) - self.data["num_descriptors"].append(res.descCount) - - def record_failed_transfer(self): - """Record a failed NIXL transfer operation.""" - self.data["num_failed_transfers"].append(1) - - def record_failed_notification(self): - """Record a failed NIXL notification (send_notif).""" - self.data["num_failed_notifications"].append(1) - - def record_kv_expired_req(self): - """Record a request that had its KV blocks expire.""" - self.data["num_kv_expired_reqs"].append(1) - - def clone_and_reset(self) -> "NixlKVConnectorStats": - old = copy.copy(self) - self.reset() - return old - - def is_empty(self) -> bool: - # Do not discard metrics update that are entirely failures related. - return ( - self.num_successful_transfers == 0 - and len(self.data["num_failed_transfers"]) == 0 - and len(self.data["num_failed_notifications"]) == 0 - and len(self.data["num_kv_expired_reqs"]) == 0 - ) - - def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: - if not other.is_empty(): - for k, v in other.data.items(): - accumulator = self.data[k] - assert isinstance(accumulator, list) - accumulator.extend(v) - return self - - def reduce(self) -> dict[str, int | float]: - # Compute compact representative stats suitable for CLI logging - if self.num_successful_transfers == 0: - # CLI logging only reports successful transfers stats. If all requests in - # the interval were unsuccessful, Prom will report failures stats instead. - return { - "Num successful transfers": 0, - "Avg xfer time (ms)": 0, - "P90 xfer time (ms)": 0, - "Avg post time (ms)": 0, - "P90 post time (ms)": 0, - "Avg MB per transfer": 0, - "Throughput (MB/s)": 0, - "Avg number of descriptors": 0, - } - - xfer_time = np.asarray(self.data["transfer_duration"]) - post_time = np.asarray(self.data["post_duration"]) - # Convert to MB for CLI logging. - mb = np.asarray(self.data["bytes_transferred"]) / 2**20 - descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32) - n = len(descs) - assert n == self.num_successful_transfers - - total_mb = mb.sum() - avg_mb = total_mb / n - - total_time_seconds = xfer_time.sum() - throughput_mb_s = total_mb / total_time_seconds - - return { - "Num successful transfers": n, - "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3), - "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), - "Avg post time (ms)": round(post_time.mean() * 1e3, 3), - "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3), - "Avg MB per transfer": round(avg_mb, 3), - "Throughput (MB/s)": round(throughput_mb_s, 3), - "Avg number of descriptors": round(descs.mean(), 1), - } - - @property - def num_successful_transfers(self) -> int: - return len(self.data["transfer_duration"]) - - -class NixlPromMetrics(KVConnectorPromMetrics): - def __init__( - self, - vllm_config: VllmConfig, - metric_types: dict[type[PromMetric], type[PromMetricT]], - labelnames: list[str], - per_engine_labelvalues: dict[int, list[object]], - ): - super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) - - buckets = [ - 0.001, - 0.005, - 0.01, - 0.025, - 0.05, - 0.075, - 0.1, - 0.2, - 0.3, - 0.5, - 0.75, - 1.0, - 5.0, - ] - nixl_histogram_xfer_time = self._histogram_cls( - name="vllm:nixl_xfer_time_seconds", - documentation="Histogram of transfer duration for NIXL KV Cache transfers.", - buckets=buckets[1:], - labelnames=labelnames, - ) - self.nixl_histogram_xfer_time = create_metric_per_engine( - nixl_histogram_xfer_time, self.per_engine_labelvalues - ) - nixl_histogram_post_time = self._histogram_cls( - name="vllm:nixl_post_time_seconds", - documentation="Histogram of transfer post time for NIXL KV" - " Cache transfers.", - buckets=buckets, - labelnames=labelnames, - ) - self.nixl_histogram_post_time = create_metric_per_engine( - nixl_histogram_post_time, self.per_engine_labelvalues - ) - # uniform 2kb to 16gb range - buckets = [2 ** (10 + i) for i in range(1, 25, 2)] - nixl_histogram_bytes_transferred = self._histogram_cls( - name="vllm:nixl_bytes_transferred", - documentation="Histogram of bytes transferred per NIXL KV Cache transfers.", - buckets=buckets, - labelnames=labelnames, - ) - self.nixl_histogram_bytes_transferred = create_metric_per_engine( - nixl_histogram_bytes_transferred, self.per_engine_labelvalues - ) - buckets = [ - 10, - 20, - 30, - 50, - 75, - 100, - 200, - 400, - 1000, - 2000, - 4000, - 10000, - 20000, - 50000, - ] - nixl_histogram_num_descriptors = self._histogram_cls( - name="vllm:nixl_num_descriptors", - documentation="Histogram of number of descriptors per NIXL" - " KV Cache transfers.", - buckets=buckets, - labelnames=labelnames, - ) - self.nixl_histogram_num_descriptors = create_metric_per_engine( - nixl_histogram_num_descriptors, self.per_engine_labelvalues - ) - counter_nixl_num_failed_transfers = self._counter_cls( - name="vllm:nixl_num_failed_transfers", - documentation="Number of failed NIXL KV Cache transfers.", - labelnames=labelnames, - ) - self.counter_nixl_num_failed_transfers = create_metric_per_engine( - counter_nixl_num_failed_transfers, self.per_engine_labelvalues - ) - counter_nixl_num_failed_notifications = self._counter_cls( - name="vllm:nixl_num_failed_notifications", - documentation="Number of failed NIXL KV Cache notifications.", - labelnames=labelnames, - ) - self.counter_nixl_num_failed_notifications = create_metric_per_engine( - counter_nixl_num_failed_notifications, self.per_engine_labelvalues - ) - - counter_nixl_num_kv_expired_reqs = self._counter_cls( - name="vllm:nixl_num_kv_expired_reqs", - documentation="Number of requests that had their KV expire. " - "NOTE: This metric is tracked on the P instance.", - labelnames=labelnames, - ) - self.counter_nixl_num_kv_expired_reqs = create_metric_per_engine( - counter_nixl_num_kv_expired_reqs, self.per_engine_labelvalues - ) - - def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): - for prom_obj, list_item_key in zip( - [ - self.nixl_histogram_xfer_time, - self.nixl_histogram_post_time, - self.nixl_histogram_bytes_transferred, - self.nixl_histogram_num_descriptors, - ], - [ - "transfer_duration", - "post_duration", - "bytes_transferred", - "num_descriptors", - ], - ): - for list_item in transfer_stats_data[list_item_key]: - prom_obj[engine_idx].observe(list_item) - for counter_obj, counter_item_key in zip( - [ - self.counter_nixl_num_failed_transfers, - self.counter_nixl_num_failed_notifications, - self.counter_nixl_num_kv_expired_reqs, - ], - ["num_failed_transfers", "num_failed_notifications", "num_kv_expired_reqs"], - ): - for list_item in transfer_stats_data[counter_item_key]: - counter_obj[engine_idx].inc(list_item) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 1831bd5770c..9fd0bed8d3e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -424,7 +424,10 @@ class OffloadingConnectorScheduler: parent_block_hash=None, token_ids=[], lora_id=None, - block_size=event.block_size, + block_size=0, medium=event.medium, lora_name=None, ) + + def shutdown(self) -> None: + self.manager.shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 23c62b6eca6..cc6d8262c7e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -394,3 +394,13 @@ class OffloadingConnectorWorker: kv_connector_stats = self.kv_connector_stats self.kv_connector_stats = OffloadingConnectorStats() return kv_connector_stats + + def shutdown(self) -> None: + # Drop deferred store jobs: there is no point in submitting + # them during shutdown. + self._unsubmitted_store_jobs.clear() + self._jobs.clear() + self._load_job.clear() + self._store_jobs.clear() + self._finished_reqs_waiting_for_store.clear() + self.worker.shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 547ee2578a1..f11281dcf14 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -64,6 +64,12 @@ class OffloadingConnector(KVConnectorBase_V1): elif role == KVConnectorRole.WORKER: self.connector_worker = OffloadingConnectorWorker(spec) + def shutdown(self) -> None: + if self.connector_worker is not None: + self.connector_worker.shutdown() + if self.connector_scheduler is not None: + self.connector_scheduler.shutdown() + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py index 6d65e006e1b..c8a5e10344b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py @@ -114,7 +114,7 @@ def derive_mamba_conv_split( assert len(conv_shape) == 2, f"Expected 2D conv state shape, got {conv_shape}" # NOTE (ZhanqiuHu): 3-read requires DS layout, which is already asserted - # in nixl_connector __init__. Use it directly instead of heuristic detection. + # in nixl worker __init__. Use it directly instead of heuristic detection. assert is_conv_state_dim_first(), "3-read requires DS conv state layout" local_conv_dim = conv_shape[0] # DS: (conv_dim_local, conv_rows) conv_rows = conv_shape[1] diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 55c87bf356c..7028b12dab3 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -45,6 +45,7 @@ from vllm.config import ( KVTransferConfig, LoadConfig, LoRAConfig, + MambaConfig, ModelConfig, MultiModalConfig, ObservabilityConfig, @@ -72,6 +73,7 @@ from vllm.config.cache import ( from vllm.config.device import Device from vllm.config.kernel import IrOpPriorityConfig, MoEBackend from vllm.config.lora import MaxLoRARanks +from vllm.config.mamba import MambaBackendEnum from vllm.config.model import ( ConvertOption, HfOverrides, @@ -578,6 +580,7 @@ class EngineArgs: pooler_config: PoolerConfig | None = ModelConfig.pooler_config compilation_config: CompilationConfig = get_field(VllmConfig, "compilation_config") attention_config: AttentionConfig = get_field(VllmConfig, "attention_config") + mamba_config: MambaConfig = get_field(VllmConfig, "mamba_config") kernel_config: KernelConfig = get_field(VllmConfig, "kernel_config") enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" @@ -610,10 +613,12 @@ class EngineArgs: mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + + mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON enable_mamba_cache_stochastic_rounding: bool = ( - CacheConfig.enable_mamba_cache_stochastic_rounding + MambaConfig.enable_stochastic_rounding ) - mamba_cache_philox_rounds: int = CacheConfig.mamba_cache_philox_rounds + mamba_cache_philox_rounds: int = MambaConfig.stochastic_rounding_philox_rounds additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config") @@ -655,6 +660,8 @@ class EngineArgs: self.compilation_config = CompilationConfig(**self.compilation_config) if isinstance(self.attention_config, dict): self.attention_config = AttentionConfig(**self.attention_config) + if isinstance(self.mamba_config, dict): + self.mamba_config = MambaConfig(**self.mamba_config) if isinstance(self.kernel_config, dict): self.kernel_config = KernelConfig(**self.kernel_config) if isinstance(self.eplb_config, dict): @@ -825,6 +832,22 @@ class EngineArgs: "--attention-backend", **attention_kwargs["backend"] ) + # Mamba arguments + mamba_kwargs = get_kwargs(MambaConfig) + mamba_group = parser.add_argument_group( + title="MambaConfig", + description=MambaConfig.__doc__, + ) + mamba_group.add_argument("--mamba-backend", **mamba_kwargs["backend"]) + mamba_group.add_argument( + "--enable-mamba-cache-stochastic-rounding", + **mamba_kwargs["enable_stochastic_rounding"], + ) + mamba_group.add_argument( + "--mamba-cache-philox-rounds", + **mamba_kwargs["stochastic_rounding_philox_rounds"], + ) + # Structured outputs arguments structured_outputs_kwargs = get_kwargs(StructuredOutputsConfig) structured_outputs_group = parser.add_argument_group( @@ -1050,13 +1073,6 @@ class EngineArgs: cache_group.add_argument( "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] ) - cache_group.add_argument( - "--enable-mamba-cache-stochastic-rounding", - **cache_kwargs["enable_mamba_cache_stochastic_rounding"], - ) - cache_group.add_argument( - "--mamba-cache-philox-rounds", **cache_kwargs["mamba_cache_philox_rounds"] - ) cache_group.add_argument( "--kv-offloading-size", **cache_kwargs["kv_offloading_size"] ) @@ -1588,10 +1604,7 @@ class EngineArgs: self._check_feature_supported() self._set_default_chunked_prefill_and_prefix_caching_args(model_config) - self._set_default_max_num_seqs_and_batched_tokens_args( - usage_context, model_config - ) - + self._set_default_reasoning_config_args() sliding_window: int | None = None if not is_interleaved(model_config.hf_text_config): # Only set CacheConfig.sliding_window if the model is all sliding @@ -1625,12 +1638,35 @@ class EngineArgs: mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, mamba_cache_mode=self.mamba_cache_mode, - enable_mamba_cache_stochastic_rounding=self.enable_mamba_cache_stochastic_rounding, - mamba_cache_philox_rounds=self.mamba_cache_philox_rounds, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, ) + # TurboQuant: auto-skip first/last 2 layers (boundary protection). + # These layers are most sensitive to quantization error. + # Users can add extra layers via --kv-cache-dtype-skip-layers. + if resolved_cache_dtype.startswith("turboquant_"): + if model_config.is_hybrid: + raise NotImplementedError( + "TurboQuant KV cache is not supported for hybrid " + "(attention + Mamba) models. Boundary layer protection " + "requires uniform attention layers." + ) + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + + num_layers = model_config.hf_text_config.num_hidden_layers + boundary = TurboQuantConfig.get_boundary_skip_layers(num_layers) + existing = set(cache_config.kv_cache_dtype_skip_layers) + merged = sorted(existing | set(boundary), key=lambda x: int(x)) + cache_config.kv_cache_dtype_skip_layers = merged + logger.info( + "TQ: skipping layers %s for boundary protection (num_layers=%d)", + merged, + num_layers, + ) + ray_runtime_env = None if is_ray_initialized(): # Ray Serve LLM calls `create_engine_config` in the context @@ -1846,6 +1882,12 @@ class EngineArgs: target_parallel_config=parallel_config, ) + self._set_default_max_num_seqs_and_batched_tokens_args( + usage_context, + model_config, + parallel_config, + ) + assert self.max_num_batched_tokens is not None, ( "max_num_batched_tokens must be set by this point" ) @@ -1931,6 +1973,35 @@ class EngineArgs: self.attention_backend ) + # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert + # FlashAttentionImpl which fails with TurboQuantAttentionImpl. + if resolved_cache_dtype.startswith("turboquant_") and ( + attention_config.flash_attn_version is None + or attention_config.flash_attn_version >= 3 + ): + logger.warning( + "TurboQuant is not yet compatible with FlashAttention >= 3. " + "Overriding flash_attn_version to 2. To silence this " + "warning, pass --attention-config.flash_attn_version=2" + ) + attention_config.flash_attn_version = 2 + + # Mamba config overrides + mamba_config = copy.deepcopy(self.mamba_config) + # Convert string to enum if needed (CLI parsing returns a string) + if isinstance(self.mamba_backend, str): + mamba_config.backend = MambaBackendEnum[self.mamba_backend.upper()] + else: + mamba_config.backend = self.mamba_backend + if self.enable_mamba_cache_stochastic_rounding: + mamba_config.enable_stochastic_rounding = ( + self.enable_mamba_cache_stochastic_rounding + ) + if self.mamba_cache_philox_rounds: + mamba_config.stochastic_rounding_philox_rounds = ( + self.mamba_cache_philox_rounds + ) + # Kernel config overrides kernel_config = copy.deepcopy(self.kernel_config) if self.enable_flashinfer_autotune is not None: @@ -2029,6 +2100,7 @@ class EngineArgs: load_config=load_config, offload_config=offload_config, attention_config=attention_config, + mamba_config=mamba_config, kernel_config=kernel_config, lora_config=lora_config, speculative_config=speculative_config, @@ -2233,10 +2305,18 @@ class EngineArgs: ) self.enable_prefix_caching = False + def _set_default_reasoning_config_args(self): + if not self.reasoning_parser: + return + if self.reasoning_config is None: + self.reasoning_config = ReasoningConfig() + self.reasoning_config.reasoning_parser = self.reasoning_parser + def _set_default_max_num_seqs_and_batched_tokens_args( self, usage_context: UsageContext | None, model_config: ModelConfig, + parallel_config: ParallelConfig, ): world_size = self.pipeline_parallel_size * self.tensor_parallel_size ( @@ -2248,10 +2328,15 @@ class EngineArgs: orig_max_num_seqs = self.max_num_seqs if self.max_num_batched_tokens is None: - self.max_num_batched_tokens = default_max_num_batched_tokens.get( - usage_context, - SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS, - ) + if parallel_config.use_batched_dp_moe: + self.max_num_batched_tokens = ( + SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP + ) + else: + self.max_num_batched_tokens = default_max_num_batched_tokens.get( + usage_context, + SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS, + ) if self.max_num_seqs is None: self.max_num_seqs = default_max_num_seqs.get( diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 3d466e3fc2a..50013a060a8 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -14,7 +14,6 @@ from vllm.distributed.weight_transfer.base import ( from vllm.inputs import EngineInput, PromptType from vllm.lora.request import LoRARequest from vllm.outputs import PoolingRequestOutput, RequestOutput -from vllm.plugins.io_processors import IOProcessor from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams @@ -44,7 +43,6 @@ class EngineClient(ABC): vllm_config: VllmConfig model_config: ModelConfig renderer: BaseRenderer - io_processor: IOProcessor | None input_processor: InputProcessor @property diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index c1324d2f039..3710473560d 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -290,7 +290,7 @@ class CustomChatCompletionMessageParam(TypedDict, total=False): tool_call_id: str | None """Tool call that this message is responding to.""" - tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None + tool_calls: list[ChatCompletionMessageToolCallParam] | None """The tool calls generated by the model, such as function calls.""" reasoning: str | None @@ -321,7 +321,7 @@ class ConversationMessage(TypedDict, total=False): name: str | None """The name of the function to call""" - tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None + tool_calls: list[ChatCompletionMessageToolCallParam] | None """The tool calls generated by the model, such as function calls.""" reasoning: str | None diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index 6a0505cb94a..8213184ad06 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -280,36 +280,23 @@ def run_multi_api_server(args: argparse.Namespace): vllm_config, executor_class, log_stats, addresses, num_api_servers ) as (local_engine_manager, coordinator, addresses, tensor_queue): # Construct common args for the APIServerProcessManager up-front. - api_server_manager_kwargs = dict( + stats_update_address = None + if coordinator: + stats_update_address = coordinator.get_stats_publish_address() + + # Start API servers. + api_server_manager = APIServerProcessManager( listen_address=listen_address, sock=sock, args=args, num_servers=num_api_servers, input_addresses=addresses.inputs, output_addresses=addresses.outputs, - stats_update_address=coordinator.get_stats_publish_address() - if coordinator - else None, + stats_update_address=stats_update_address, tensor_queue=tensor_queue, ) - # For dp ranks > 0 in external/hybrid DP LB modes, we must delay the - # start of the API servers until the local engine is started - # (after the launcher context manager exits), - # since we get the front-end stats update address from the coordinator - # via the handshake with the local engine. - if dp_rank == 0 or not parallel_config.local_engines_only: - # Start API servers using the manager. - api_server_manager = APIServerProcessManager(**api_server_manager_kwargs) - - # Start API servers now if they weren't already started. - if api_server_manager is None: - api_server_manager_kwargs["stats_update_address"] = ( - addresses.frontend_stats_publish_address - ) - api_server_manager = APIServerProcessManager(**api_server_manager_kwargs) - - # Wait for API servers + # Wait for API servers. try: wait_for_completion_or_failure( api_server_manager=api_server_manager, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 1be2cdd5c74..d296e84d041 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -49,9 +49,7 @@ from vllm.entrypoints.chat_utils import ( load_chat_template, ) from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors -from vllm.entrypoints.pooling.scoring.io_processor import ( - ScoringIOProcessor, -) +from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessor from vllm.entrypoints.pooling.scoring.typing import ScoreInput from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext from vllm.entrypoints.utils import log_non_default_args @@ -398,12 +396,11 @@ class LLM: self.runner_type = self.model_config.runner_type self.renderer = self.llm_engine.renderer self.chat_template = load_chat_template(chat_template) - self.io_processor = self.llm_engine.io_processor self.input_processor = self.llm_engine.input_processor self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template) self.pooling_io_processors = init_pooling_io_processors( supported_tasks=supported_tasks, - model_config=self.model_config, + vllm_config=self.llm_engine.vllm_config, renderer=self.renderer, chat_template_config=self.chat_template_config, ) @@ -1081,118 +1078,55 @@ class LLM: pooled hidden states in the same order as the input prompts. """ - self._verify_pooling_task(pooling_task) - - if isinstance(prompts, dict) and "data" in prompts: - if self.io_processor is None: - raise ValueError( - "No IOProcessor plugin installed. Please refer " - "to the documentation and to the " - "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." - ) - - # Validate the request data is valid for the loaded plugin - prompt_data = prompts.get("data") - if prompt_data is None: - raise ValueError( - "The 'data' field of the prompt is expected to contain " - "the prompt data and it cannot be None. " - "Refer to the documentation of the IOProcessor " - "in use for more details." - ) - validated_prompt = self.io_processor.parse_data(prompt_data) - - # obtain the actual model prompts from the pre-processor - prompts = self.io_processor.pre_process(prompt=validated_prompt) - prompts_seq = prompt_to_seq(prompts) - - params_seq: Sequence[PoolingParams] = [ - self.io_processor.merge_pooling_params(param) - for param in self._params_to_seq( - pooling_params, - len(prompts_seq), - ) - ] - for p in params_seq: - if p.task is None: - p.task = "plugin" - - outputs = self._run_completion( - prompts=prompts_seq, - params=params_seq, - output_type=PoolingRequestOutput, - use_tqdm=use_tqdm, - lora_request=lora_request, - tokenization_kwargs=tokenization_kwargs, + if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin": + raise ValueError( + "The 'data' field is only supported for the 'plugin' pooling task." ) + self._verify_pooling_task(pooling_task) + assert pooling_task is not None and pooling_task in self.pooling_io_processors - # get the post-processed model outputs - assert self.io_processor is not None - processed_outputs = self.io_processor.post_process(outputs) + io_processor = self.pooling_io_processors[pooling_task] - return [ - PoolingRequestOutput[Any]( - request_id="", - outputs=processed_outputs, - num_cached_tokens=getattr( - processed_outputs, "num_cached_tokens", 0 - ), - prompt_token_ids=[], - finished=True, - ) - ] - else: - if pooling_params is None: - # Use default pooling params. - pooling_params = PoolingParams() + if pooling_params is None: + pooling_params = PoolingParams() - prompts_seq = prompt_to_seq(prompts) - params_seq = self._params_to_seq(pooling_params, len(prompts_seq)) + ctx = OfflineInputsContext( + prompts=prompts, + pooling_params=pooling_params, + tokenization_kwargs=tokenization_kwargs, + ) - for param in params_seq: - if param.task is None: - param.task = pooling_task - elif param.task != pooling_task: - msg = ( - f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" - ) - raise ValueError(msg) + engine_inputs = io_processor.pre_process_offline(ctx) + n_inputs = len(engine_inputs) + assert ctx.pooling_params is not None - if pooling_task in self.pooling_io_processors: - io_processor = self.pooling_io_processors[pooling_task] - processor_inputs = io_processor.pre_process_offline( - ctx=OfflineInputsContext( - prompts=prompts_seq, tokenization_kwargs=tokenization_kwargs - ) - ) - seq_lora_requests = self._lora_request_to_seq( - lora_request, len(prompts_seq) - ) - seq_priority = self._priority_to_seq(None, len(prompts)) + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) - self._render_and_add_requests( - prompts=processor_inputs, - params=params_seq, - lora_requests=seq_lora_requests, - priorities=seq_priority, - ) + for param in params_seq: + if param.task is None: + param.task = pooling_task + elif pooling_task == "plugin": + # `plugin` task uses io_processor.parse_request to verify inputs. + # We actually allow plugin to overwrite pooling_task. + pass + elif param.task != pooling_task: + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" + raise ValueError(msg) - outputs = self._run_engine( - use_tqdm=use_tqdm, output_type=PoolingRequestOutput - ) - outputs = io_processor.post_process_offline( - ctx=OfflineOutputsContext(outputs=outputs) - ) - else: - outputs = self._run_completion( - prompts=prompts_seq, - params=params_seq, - output_type=PoolingRequestOutput, - use_tqdm=use_tqdm, - lora_request=lora_request, - tokenization_kwargs=tokenization_kwargs, - ) + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) + seq_priority = self._priority_to_seq(None, n_inputs) + + self._render_and_add_requests( + prompts=engine_inputs, + params=params_seq, + lora_requests=seq_lora_requests, + priorities=seq_priority, + ) + + outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput) + outputs = io_processor.post_process_offline( + ctx=OfflineOutputsContext(outputs=outputs) + ) return outputs def _verify_pooling_task(self, pooling_task: PoolingTask | None): @@ -1254,6 +1188,14 @@ class LLM: pooling_task, ) + if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors: + raise ValueError( + "No IOProcessor plugin installed. Please refer " + "to the documentation and to the " + "'prithvi_geospatial_mae_io_processor' " + "offline inference example for more details." + ) + def embed( self, prompts: PromptType | Sequence[PromptType], @@ -1458,6 +1400,9 @@ class LLM: scoring_data = io_processor.valid_inputs(data_1, data_2) n_queries = len(scoring_data.data_1) + if pooling_params is None: + pooling_params = PoolingParams() + ctx = OfflineInputsContext( prompts=scoring_data, pooling_params=pooling_params, @@ -1466,15 +1411,11 @@ class LLM: n_queries=n_queries, ) - processor_inputs = io_processor.pre_process_offline(ctx) + engine_inputs = io_processor.pre_process_offline(ctx) + n_inputs = len(engine_inputs) - seq_lora_requests = self._lora_request_to_seq( - lora_request, len(processor_inputs) - ) - - if ctx.pooling_params is None: - ctx.pooling_params = PoolingParams() - params_seq = self._params_to_seq(ctx.pooling_params, len(processor_inputs)) + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) for param in params_seq: if param.task is None: @@ -1483,10 +1424,10 @@ class LLM: msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" raise ValueError(msg) - seq_priority = self._priority_to_seq(None, len(processor_inputs)) + seq_priority = self._priority_to_seq(None, n_inputs) self._render_and_add_requests( - prompts=processor_inputs, + prompts=engine_inputs, params=params_seq, lora_requests=seq_lora_requests, priorities=seq_priority, @@ -1579,7 +1520,7 @@ class LLM: if isinstance(params, Sequence): if len(params) != num_requests: raise ValueError( - f"The lengths of prompts ({params}) " + f"The lengths of prompts ({num_requests}) " f"and params ({len(params)}) must be the same." ) diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 2b6cb810ea3..85d2fe43d01 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -370,7 +370,6 @@ async def init_app_state( state.openai_serving_render = OpenAIServingRender( model_config=engine_client.model_config, renderer=engine_client.renderer, - io_processor=engine_client.io_processor, model_registry=state.openai_serving_models.registry, request_logger=request_logger, chat_template=resolved_chat_template, @@ -441,13 +440,12 @@ async def init_render_app_state( Unlike :func:`init_app_state` this function does not require an :class:`~vllm.engine.protocol.EngineClient`; it bootstraps the - preprocessing pipeline (renderer, io_processor, input_processor) + preprocessing pipeline (renderer, input_processor) directly from the :class:`~vllm.config.VllmConfig`. """ from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.openai.models.serving import OpenAIModelRegistry from vllm.entrypoints.serve.render.serving import OpenAIServingRender - from vllm.plugins.io_processors import get_io_processor from vllm.renderers import renderer_from_config served_model_names = args.served_model_name or [args.model] @@ -465,15 +463,11 @@ async def init_render_app_state( request_logger = None renderer = renderer_from_config(vllm_config) - io_processor = get_io_processor( - vllm_config, renderer, vllm_config.model_config.io_processor_plugin - ) resolved_chat_template = load_chat_template(args.chat_template) state.openai_serving_render = OpenAIServingRender( model_config=vllm_config.model_config, renderer=renderer, - io_processor=io_processor, model_registry=model_registry, request_logger=request_logger, chat_template=resolved_chat_template, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 533959df609..2bc1b6e0875 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -357,6 +357,47 @@ class ChatCompletionRequest(OpenAIBaseModel): # --8<-- [end:chat-completion-extra-params] + @model_validator(mode="before") + @classmethod + def _materialize_tool_calls_before(cls, data: Any) -> Any: + """Eagerly convert tool_calls generators/iterators to lists. + + Must run before Pydantic field validation so that one-shot + generators are not consumed during union type matching of + ChatCompletionAssistantMessageParam (which types tool_calls + as Iterable[...]). + """ + if not isinstance(data, dict): + return data + messages = data.get("messages") + if not isinstance(messages, list): + return data + for msg in messages: + if not isinstance(msg, dict): + continue + tool_calls = msg.get("tool_calls") + if tool_calls is not None and not isinstance(tool_calls, list): + msg["tool_calls"] = list(tool_calls) + return data + + @model_validator(mode="after") + def _materialize_tool_calls_after(self) -> "ChatCompletionRequest": + """Convert Pydantic ValidatorIterator wrappers back to lists. + + Even after the "before" validator converts iterables to lists, + Pydantic re-wraps them in a ValidatorIterator when validating + against ChatCompletionAssistantMessageParam's Iterable[...] type. + This "after" pass materialises those wrappers so downstream code + (tokenizers, model_dump_json) always sees plain lists. + """ + for msg in self.messages: + if not isinstance(msg, dict): + continue + tool_calls = msg.get("tool_calls") + if tool_calls is not None and not isinstance(tool_calls, list): + msg["tool_calls"] = list(tool_calls) + return self + def build_chat_params( self, default_template: str | None, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index a426836afd3..0b8dd0aa28e 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -68,11 +68,11 @@ from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.outputs import CompletionOutput, RequestOutput from vllm.parser import ParserManager +from vllm.parser.abstract_parser import Parser from vllm.reasoning import ReasoningParser from vllm.renderers import ChatParams from vllm.sampling_params import BeamSearchParams, SamplingParams from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers import ToolParser from vllm.tool_parsers.mistral_tool_parser import MistralToolCall from vllm.tool_parsers.utils import partial_json_loads from vllm.utils.collection_utils import as_list @@ -134,6 +134,12 @@ class OpenAIServingChat(OpenAIServing): enable_auto_tools=enable_auto_tools, model_name=self.model_config.model, ) + self.parser_cls = ParserManager.get_parser( + tool_parser_name=tool_parser, + reasoning_parser_name=reasoning_parser, + enable_auto_tools=enable_auto_tools, + model_name=self.model_config.model, + ) self.exclude_tools_when_tool_choice_none = exclude_tools_when_tool_choice_none self.enable_prompt_tokens_details = enable_prompt_tokens_details @@ -216,13 +222,12 @@ class OpenAIServingChat(OpenAIServing): # Streaming response tokenizer = self.renderer.tokenizer assert tokenizer is not None + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, + ) reasoning_parser: ReasoningParser | None = None if self.reasoning_parser_cls: - # Pass the same chat template kwargs as used in tokenization - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) reasoning_parser = self.reasoning_parser_cls( tokenizer, chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] @@ -338,6 +343,7 @@ class OpenAIServingChat(OpenAIServing): tokenizer, request_metadata, reasoning_parser, + chat_template_kwargs=chat_template_kwargs, ) return await self.chat_completion_full_generator( @@ -505,6 +511,7 @@ class OpenAIServingChat(OpenAIServing): tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, reasoning_parser: ReasoningParser | None = None, + chat_template_kwargs: dict[str, Any] | None = None, ) -> AsyncGenerator[str, None]: created_time = int(time.time()) chunk_object_type: Final = "chat.completion.chunk" @@ -549,29 +556,29 @@ class OpenAIServingChat(OpenAIServing): if tool_choice_auto or reasoning_parser: # These are only required in "auto" tool choice case all_previous_token_ids = [[] for _ in range(num_choices)] - # For reasoning parser and tool call all enabled - added_content_delta_arr = [False] * num_choices reasoning_end_arr = [False] * num_choices prompt_is_reasoning_end_arr: list[bool | None] = [None] * num_choices else: all_previous_token_ids = None - # Prepare the tool parser if it's needed try: - if tool_choice_auto and self.tool_parser: + if self.parser_cls is not None: if tokenizer is None: raise ValueError( "Tokenizer not available when `skip_tokenizer_init=True`" ) - - tool_parsers: list[ToolParser | None] = [ - self.tool_parser(tokenizer, request.tools) + parsers: list[Parser | None] = [ + self.parser_cls( + tokenizer, + request.tools, + chat_template_kwargs=chat_template_kwargs, + ) for _ in range(num_choices) ] else: - tool_parsers = [None] * num_choices + parsers = [None] * num_choices except Exception as e: - logger.exception("Error in tool parser creation.") + logger.exception("Error in parser creation.") data = self.create_streaming_error_response(e) yield f"data: {data}\n\n" yield "data: [DONE]\n\n" @@ -675,7 +682,8 @@ class OpenAIServingChat(OpenAIServing): for output in res.outputs: i = output.index - tool_parser = tool_parsers[i] + parser = parsers[i] + tool_parser = parser.tool_parser if parser is not None else None if ( reasoning_parser @@ -903,109 +911,16 @@ class OpenAIServingChat(OpenAIServing): history_tool_call_cnt += 1 tools_streamed[i] = True - # handle streaming deltas for tools with "auto" tool choice - # and reasoning parser - elif tool_choice_auto and reasoning_parser: - assert tool_parser is not None - assert added_content_delta_arr is not None - assert reasoning_end_arr is not None - output_token_ids = as_list(output.token_ids) - if not reasoning_end_arr[i]: - # When encountering think end id in prompt_token_ids - # i.e {"enable_thinking": False}, - # set reasoning status to end. - if prompt_is_reasoning_end_arr[i]: - reasoning_end_arr[i] = True - current_token_ids = output_token_ids - # Don't update current_text, keep it as is from delta - else: - delta_message = ( - reasoning_parser.extract_reasoning_streaming( - previous_text, - current_text, - delta_text, - previous_token_ids, - current_token_ids, - output_token_ids, - ) - ) - - # When encountering think end id in delta_token_ids, - # set reasoning status to end. - # Remove the text and token ids related - # to 'reasoning'. - if reasoning_parser.is_reasoning_end(output_token_ids): - reasoning_end_arr[i] = True - current_token_ids = ( - reasoning_parser.extract_content_ids( - output_token_ids - ) - ) - if delta_message and delta_message.content: - current_text = delta_message.content - delta_message.content = None - else: - current_text = "" - - # handle tool calls only after reasoning is done, - if reasoning_end_arr[i]: - delta_token_ids = output_token_ids - # First time to tool call, - # add the remaining text and token ids - # to delta from previous - if not added_content_delta_arr[i]: - added_content_delta_arr[i] = True - previous_text = "" - previous_token_ids = [] - delta_text = current_text - delta_token_ids = current_token_ids - - delta_message = tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=delta_token_ids, - request=request, - ) - if delta_message and delta_message.tool_calls: - tools_streamed[i] = True - # when only tool calls - elif tool_choice_auto: - assert tool_parser is not None - delta_message = tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, + elif parser is not None: + delta_message = parser.parse_delta( delta_text=delta_text, - previous_token_ids=previous_token_ids, - current_token_ids=current_token_ids, - delta_token_ids=output.token_ids, + delta_token_ids=as_list(output.token_ids), request=request, + prompt_token_ids=res.prompt_token_ids, ) if delta_message and delta_message.tool_calls: tools_streamed[i] = True - - # when only reasoning - elif reasoning_parser: - # When encountering think end id in prompt_token_ids - # i.e {"enable_thinking": False}, - # set reasoning status to end. - # Route all generated tokens as content directly. - if prompt_is_reasoning_end_arr[i]: - delta_message = DeltaMessage(content=delta_text) - else: - delta_message = ( - reasoning_parser.extract_reasoning_streaming( - previous_text, - current_text, - delta_text, - previous_token_ids, - current_token_ids, - output.token_ids, - ) - ) - # handle streaming just a content delta + # handle streaming just a content delta (no parsers) else: delta_message = DeltaMessage(content=delta_text) diff --git a/vllm/entrypoints/openai/engine/serving.py b/vllm/entrypoints/openai/engine/serving.py index f5f011a96f2..5bd415b4fbd 100644 --- a/vllm/entrypoints/openai/engine/serving.py +++ b/vllm/entrypoints/openai/engine/serving.py @@ -44,12 +44,6 @@ from vllm.entrypoints.openai.speech_to_text.protocol import ( TranscriptionResponse, TranslationRequest, ) -from vllm.entrypoints.pooling.pooling.protocol import ( - IOProcessorRequest, - PoolingChatRequest, - PoolingCompletionRequest, - PoolingResponse, -) from vllm.entrypoints.serve.disagg.protocol import GenerateRequest, GenerateResponse from vllm.entrypoints.serve.tokenize.protocol import ( DetokenizeRequest, @@ -62,8 +56,7 @@ from vllm.inputs import EngineInput, PromptType from vllm.logger import init_logger from vllm.logprobs import Logprob, PromptLogprobs from vllm.lora.request import LoRARequest -from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput -from vllm.pooling_params import PoolingParams +from vllm.outputs import CompletionOutput, RequestOutput from vllm.renderers import ChatParams, TokenizeParams from vllm.renderers.inputs.preprocess import ( extract_prompt_components, @@ -78,10 +71,7 @@ from vllm.tracing import ( log_tracing_disabled_warning, ) from vllm.utils import random_uuid -from vllm.utils.async_utils import ( - collect_from_async_generator, - merge_async_iterators, -) +from vllm.utils.async_utils import collect_from_async_generator logger = init_logger(__name__) @@ -101,17 +91,11 @@ class RendererChatRequest(RendererRequest, Protocol): CompletionLikeRequest: TypeAlias = ( - CompletionRequest - | TokenizeCompletionRequest - | DetokenizeRequest - | PoolingCompletionRequest + CompletionRequest | TokenizeCompletionRequest | DetokenizeRequest ) ChatLikeRequest: TypeAlias = ( - ChatCompletionRequest - | BatchChatCompletionRequest - | TokenizeChatRequest - | PoolingChatRequest + ChatCompletionRequest | BatchChatCompletionRequest | TokenizeChatRequest ) SpeechToTextRequest: TypeAlias = TranscriptionRequest | TranslationRequest @@ -121,7 +105,6 @@ AnyRequest: TypeAlias = ( | ChatLikeRequest | SpeechToTextRequest | ResponsesRequest - | IOProcessorRequest | GenerateRequest ) @@ -130,7 +113,6 @@ AnyResponse: TypeAlias = ( | ChatCompletionResponse | TranscriptionResponse | TokenizeResponse - | PoolingResponse | GenerateResponse ) @@ -146,12 +128,6 @@ class ServeContext(Generic[RequestT]): created_time: int = field(default_factory=lambda: int(time.time())) lora_request: LoRARequest | None = None engine_inputs: list[EngineInput] | None = None - - result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = ( - None - ) - final_res_batch: list[PoolingRequestOutput] = field(default_factory=list) - model_config = ConfigDict(arbitrary_types_allowed=True) @@ -171,7 +147,6 @@ class OpenAIServing: super().__init__() self.engine_client = engine_client - self.models = models self.request_logger = request_logger @@ -179,7 +154,6 @@ class OpenAIServing: self.model_config = engine_client.model_config self.renderer = engine_client.renderer - self.io_processor = engine_client.io_processor self.input_processor = engine_client.input_processor async def beam_search( @@ -381,155 +355,6 @@ class OpenAIServing: prompt_logprobs=None, ) - async def _preprocess( - self, - ctx: ServeContext, - ) -> ErrorResponse | None: - """ - Default preprocessing hook. Subclasses may override to prepare `ctx`. - """ - return None - - def _build_response( - self, - ctx: ServeContext, - ) -> AnyResponse | ErrorResponse: - """ - Default response builder. Subclass may override this method - to return the appropriate response object. - """ - return self.create_error_response("unimplemented endpoint") - - async def handle( - self, - ctx: ServeContext, - ) -> AnyResponse | ErrorResponse: - async for response in self._pipeline(ctx): - return response - - return self.create_error_response("No response yielded from pipeline") - - async def _pipeline( - self, - ctx: ServeContext, - ) -> AsyncGenerator[AnyResponse | ErrorResponse, None]: - """Execute the request processing pipeline yielding responses.""" - if error := await self._check_model(ctx.request): - yield error - if error := self._validate_request(ctx): - yield error - - preprocess_ret = await self._preprocess(ctx) - if isinstance(preprocess_ret, ErrorResponse): - yield preprocess_ret - - generators_ret = await self._prepare_generators(ctx) - if isinstance(generators_ret, ErrorResponse): - yield generators_ret - - collect_ret = await self._collect_batch(ctx) - if isinstance(collect_ret, ErrorResponse): - yield collect_ret - - yield self._build_response(ctx) - - def _validate_request(self, ctx: ServeContext) -> ErrorResponse | None: - truncate_prompt_tokens = getattr(ctx.request, "truncate_prompt_tokens", None) - - if ( - truncate_prompt_tokens is not None - and truncate_prompt_tokens > self.model_config.max_model_len - ): - return self.create_error_response( - "truncate_prompt_tokens value is " - "greater than max_model_len." - " Please request a smaller truncation size." - ) - return None - - def _create_pooling_params( - self, - ctx: ServeContext, - ) -> PoolingParams | ErrorResponse: - if not hasattr(ctx.request, "to_pooling_params"): - return self.create_error_response( - "Request type does not support pooling parameters" - ) - - return ctx.request.to_pooling_params() - - async def _prepare_generators( - self, - ctx: ServeContext, - ) -> ErrorResponse | None: - """Schedule the request and get the result generator.""" - generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - - trace_headers = ( - None - if ctx.raw_request is None - else await self._get_trace_headers(ctx.raw_request.headers) - ) - - pooling_params = self._create_pooling_params(ctx) - if isinstance(pooling_params, ErrorResponse): - return pooling_params - - if ctx.engine_inputs is None: - return self.create_error_response("Engine prompts not available") - - for i, engine_input in enumerate(ctx.engine_inputs): - request_id_item = f"{ctx.request_id}-{i}" - - self._log_inputs( - request_id_item, - engine_input, - params=pooling_params, - lora_request=ctx.lora_request, - ) - - generator = self.engine_client.encode( - engine_input, - pooling_params, - request_id_item, - lora_request=ctx.lora_request, - trace_headers=trace_headers, - priority=getattr(ctx.request, "priority", 0), - ) - - generators.append(generator) - - ctx.result_generator = merge_async_iterators(*generators) - - return None - - async def _collect_batch( - self, - ctx: ServeContext, - ) -> ErrorResponse | None: - """Collect batch results from the result generator.""" - if ctx.engine_inputs is None: - return self.create_error_response("Engine prompts not available") - - num_prompts = len(ctx.engine_inputs) - final_res_batch: list[PoolingRequestOutput | None] - final_res_batch = [None] * num_prompts - - if ctx.result_generator is None: - return self.create_error_response("Result generator not available") - - async for i, res in ctx.result_generator: - final_res_batch[i] = res - - if None in final_res_batch: - return self.create_error_response( - "Failed to generate results for all prompts" - ) - - ctx.final_res_batch = [res for res in final_res_batch if res is not None] - - return None - @staticmethod def create_error_response( message: str | Exception, @@ -719,7 +544,7 @@ class OpenAIServing: self, request_id: str, inputs: PromptType | EngineInput, - params: SamplingParams | PoolingParams | BeamSearchParams | None, + params: SamplingParams | BeamSearchParams | None, lora_request: LoRARequest | None, ) -> None: if self.request_logger is None: diff --git a/vllm/entrypoints/openai/models/serving.py b/vllm/entrypoints/openai/models/serving.py index dd7a8687f2b..ba1902d5b77 100644 --- a/vllm/entrypoints/openai/models/serving.py +++ b/vllm/entrypoints/openai/models/serving.py @@ -112,7 +112,6 @@ class OpenAIServingModels: self.model_config = self.engine_client.model_config self.renderer = self.engine_client.renderer - self.io_processor = self.engine_client.io_processor self.input_processor = self.engine_client.input_processor async def init_static_loras(self): diff --git a/vllm/entrypoints/openai/responses/api_router.py b/vllm/entrypoints/openai/responses/api_router.py index 88d82126094..61077f1a7c5 100644 --- a/vllm/entrypoints/openai/responses/api_router.py +++ b/vllm/entrypoints/openai/responses/api_router.py @@ -39,7 +39,8 @@ async def _convert_stream_to_sse_events( event_type = getattr(event, "type", "unknown") # https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format event_data = ( - f"event: {event_type}\ndata: {event.model_dump_json(indent=None)}\n\n" + f"event: {event_type}\ndata: " + f"{event.model_dump_json(indent=None, by_alias=True)}\n\n" ) yield event_data @@ -65,10 +66,11 @@ async def create_responses(request: ResponsesRequest, raw_request: Request): if isinstance(generator, ErrorResponse): return JSONResponse( - content=generator.model_dump(), status_code=generator.error.code + content=generator.model_dump(mode="json", by_alias=True), + status_code=generator.error.code, ) elif isinstance(generator, ResponsesResponse): - return JSONResponse(content=generator.model_dump()) + return JSONResponse(content=generator.model_dump(mode="json", by_alias=True)) return StreamingResponse( content=_convert_stream_to_sse_events(generator), media_type="text/event-stream" @@ -95,10 +97,11 @@ async def retrieve_responses( if isinstance(response, ErrorResponse): return JSONResponse( - content=response.model_dump(), status_code=response.error.code + content=response.model_dump(mode="json", by_alias=True), + status_code=response.error.code, ) elif isinstance(response, ResponsesResponse): - return JSONResponse(content=response.model_dump()) + return JSONResponse(content=response.model_dump(mode="json", by_alias=True)) return StreamingResponse( content=_convert_stream_to_sse_events(response), media_type="text/event-stream" ) @@ -115,9 +118,10 @@ async def cancel_responses(response_id: str, raw_request: Request): if isinstance(response, ErrorResponse): return JSONResponse( - content=response.model_dump(), status_code=response.error.code + content=response.model_dump(mode="json", by_alias=True), + status_code=response.error.code, ) - return JSONResponse(content=response.model_dump()) + return JSONResponse(content=response.model_dump(mode="json", by_alias=True)) def attach_router(app: FastAPI): diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index d34ba2d75bb..79f5894fb91 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -106,7 +106,7 @@ def serialize_message(msg): return msg.to_dict() else: # fallback to pydantic dump - return msg.model_dump_json(by_alias=True) + return msg.model_dump(mode="json", by_alias=True) def serialize_messages(msgs): diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 7edeb8ab0cb..6a0c4c1e9b6 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -1341,6 +1341,7 @@ class OpenAIServingResponses(OpenAIServing): current_content_index = 0 current_output_index = 0 current_item_id = "" + current_tool_call_index: int | None = None parser = self.parser(tokenizer, request.tools) if self.parser else None first_delta_sent = False previous_delta_messages: list[DeltaMessage] = [] @@ -1368,6 +1369,7 @@ class OpenAIServingResponses(OpenAIServing): ) if not delta_message: continue + tool_call_item_started = False if not first_delta_sent: current_item_id = random_uuid() if delta_message.tool_calls: @@ -1384,6 +1386,7 @@ class OpenAIServingResponses(OpenAIServing): current_tool_call_name = delta_message.tool_calls[ 0 ].function.name + current_tool_call_index = delta_message.tool_calls[0].index yield _increment_sequence_number_and_return( ResponseOutputItemAddedEvent( type="response.output_item.added", @@ -1394,13 +1397,12 @@ class OpenAIServingResponses(OpenAIServing): id=current_item_id, call_id=current_tool_call_id, name=current_tool_call_name, - arguments=delta_message.tool_calls[ - 0 - ].function.arguments, + arguments="", status="in_progress", ), ) ) + tool_call_item_started = True elif delta_message.reasoning: yield _increment_sequence_number_and_return( ResponseOutputItemAddedEvent( @@ -1572,6 +1574,79 @@ class OpenAIServingResponses(OpenAIServing): # reset previous delta messages previous_delta_messages = [] if delta_message.tool_calls and delta_message.tool_calls[0].function: + tool_call = delta_message.tool_calls[0] + tool_call_function = tool_call.function + if ( + current_tool_call_index is not None + and tool_call.index is not None + and tool_call.index != current_tool_call_index + and tool_call_function is not None + and tool_call_function.name is not None + ): + # From one tool call to another, finalize the previous + # function-call item before opening the next one. + parts = [] + for pm in previous_delta_messages: + if pm.tool_calls: + previous_tool_call = pm.tool_calls[0] + if previous_tool_call.function is not None: + parts.append( + previous_tool_call.function.arguments or "" + ) + + tool_call_arguments = "".join(parts) + yield _increment_sequence_number_and_return( + ResponseFunctionCallArgumentsDoneEvent( + type="response.function_call_arguments.done", + sequence_number=-1, + output_index=current_output_index, + item_id=current_item_id, + arguments=tool_call_arguments, + name=current_tool_call_name, + ) + ) + function_call_item = ResponseFunctionToolCall( + type="function_call", + name=current_tool_call_name, + arguments=tool_call_arguments, + status="completed", + id=current_item_id, + call_id=current_tool_call_id, + ) + yield _increment_sequence_number_and_return( + ResponseOutputItemDoneEvent( + type="response.output_item.done", + sequence_number=-1, + output_index=current_output_index, + item=function_call_item, + ) + ) + # Reset previous delta messages so the next tool call + # does not reuse arguments from the completed item. + previous_delta_messages = [] + current_output_index += 1 + current_item_id = random_uuid() + current_tool_call_name = tool_call_function.name + current_tool_call_id = f"call_{random_uuid()}" + current_tool_call_index = tool_call.index + yield _increment_sequence_number_and_return( + ResponseOutputItemAddedEvent( + type="response.output_item.added", + sequence_number=-1, + output_index=current_output_index, + item=ResponseFunctionToolCallItem( + type="function_call", + id=current_item_id, + call_id=current_tool_call_id, + name=current_tool_call_name, + arguments="", + status="in_progress", + ), + ) + ) + current_content_index = 0 + tool_call_item_started = True + if delta_message.tool_calls[0].function.arguments: yield _increment_sequence_number_and_return( ResponseFunctionCallArgumentsDeltaEvent( @@ -1583,7 +1658,10 @@ class OpenAIServingResponses(OpenAIServing): ) ) # tool call initiated with no arguments - elif delta_message.tool_calls[0].function.name: + elif ( + delta_message.tool_calls[0].function.name + and not tool_call_item_started + ): # send done with current content part # and add new function call item yield _increment_sequence_number_and_return( @@ -1628,11 +1706,11 @@ class OpenAIServingResponses(OpenAIServing): ) current_output_index += 1 current_item_id = random_uuid() - assert delta_message.tool_calls[0].function is not None current_tool_call_name = delta_message.tool_calls[ 0 ].function.name current_tool_call_id = f"call_{random_uuid()}" + current_tool_call_index = delta_message.tool_calls[0].index yield _increment_sequence_number_and_return( ResponseOutputItemAddedEvent( type="response.output_item.added", @@ -1909,7 +1987,7 @@ class OpenAIServingResponses(OpenAIServing): output=[], status="in_progress", usage=None, - ).model_dump() + ).model_dump(mode="json", by_alias=True) yield _increment_sequence_number_and_return( ResponseCreatedEvent( type="response.created", diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 789a0e0b6be..66239289f73 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -94,8 +94,10 @@ def construct_input_messages( # Prepend the conversation history. if prev_msg is not None: - # Add the previous messages. - messages.extend(prev_msg) + # Filter out system messages from previous conversation -- per the + # OpenAI spec, instructions should NOT carry over across responses. + # The current request's instructions (if any) were already added above. + messages.extend(m for m in prev_msg if m.get("role") != "system") if prev_response_output is not None: # Add the previous output. for output_item in prev_response_output: diff --git a/vllm/entrypoints/openai/server_utils.py b/vllm/entrypoints/openai/server_utils.py index 02b8c335262..f8fe3366b06 100644 --- a/vllm/entrypoints/openai/server_utils.py +++ b/vllm/entrypoints/openai/server_utils.py @@ -69,7 +69,10 @@ class AuthenticationMiddleware: return token_match def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]: - if scope["type"] not in ("http", "websocket") or scope["method"] == "OPTIONS": + if ( + scope["type"] not in ("http", "websocket") + or scope.get("method") == "OPTIONS" + ): # scope["type"] can be "lifespan" or "startup" for example, # in which case we don't need to do anything return self.app(scope, receive, send) diff --git a/vllm/entrypoints/openai/speech_to_text/serving.py b/vllm/entrypoints/openai/speech_to_text/serving.py index 28e798a986f..bacd6d794de 100644 --- a/vllm/entrypoints/openai/speech_to_text/serving.py +++ b/vllm/entrypoints/openai/speech_to_text/serving.py @@ -86,6 +86,7 @@ class OpenAIServingTranscription(OpenAISpeechToText): request_id: str, request_metadata: RequestResponseMetadata, audio_duration_s: float, + separator: str, ) -> AsyncGenerator[str, None]: generator = self._speech_to_text_stream_generator( request=request, @@ -96,6 +97,7 @@ class OpenAIServingTranscription(OpenAISpeechToText): chunk_object_type="transcription.chunk", response_stream_choice_class=TranscriptionResponseStreamChoice, stream_response_class=TranscriptionStreamResponse, + separator=separator, ) async for chunk in generator: yield chunk @@ -157,6 +159,7 @@ class OpenAIServingTranslation(OpenAISpeechToText): request_id: str, request_metadata: RequestResponseMetadata, audio_duration_s: float, + separator: str, ) -> AsyncGenerator[str, None]: generator = self._speech_to_text_stream_generator( request=request, @@ -167,6 +170,7 @@ class OpenAIServingTranslation(OpenAISpeechToText): chunk_object_type="translation.chunk", response_stream_choice_class=TranslationResponseStreamChoice, stream_response_class=TranslationStreamResponse, + separator=separator, ) async for chunk in generator: yield chunk diff --git a/vllm/entrypoints/openai/speech_to_text/speech_to_text.py b/vllm/entrypoints/openai/speech_to_text/speech_to_text.py index e0a3cf0dc0d..4ebc612a415 100644 --- a/vllm/entrypoints/openai/speech_to_text/speech_to_text.py +++ b/vllm/entrypoints/openai/speech_to_text/speech_to_text.py @@ -5,7 +5,7 @@ import io import math import time import zlib -from collections.abc import AsyncGenerator, Callable +from collections.abc import AsyncGenerator, Callable, Set from functools import cached_property from typing import Final, Literal, TypeAlias, TypeVar, cast @@ -69,6 +69,17 @@ ResponseType: TypeAlias = ( logger = init_logger(__name__) +def asr_inter_chunk_separator( + language: str | None, no_space_languages: Set[str] +) -> str: + """Space to insert between ASR text chunks for streaming and non-streaming join. + + Languages in ``no_space_languages`` (e.g. Chinese, Japanese) use an empty + separator; others use a single ASCII space. + """ + return "" if language and language.lower() in no_space_languages else " " + + class OpenAISpeechToText(OpenAIServing): """Base class for speech-to-text operations like transcription and translation.""" @@ -378,6 +389,9 @@ class OpenAISpeechToText(OpenAIServing): if error_check_ret is not None: return error_check_ret + if not request.model: + request.model = self.models.model_name() + # If the engine is dead, raise the engine's DEAD_ERROR. # This is required for the streaming case, where we return a # success status before we actually start generating text :). @@ -486,9 +500,18 @@ class OpenAISpeechToText(OpenAIServing): list_result_generator.append(generator) + separator = asr_inter_chunk_separator( + request.language, self.model_cls.no_space_languages + ) + if request.stream: return stream_generator_method( - request, list_result_generator, request_id, request_metadata, duration_s + request, + list_result_generator, + request_id, + request_metadata, + duration_s, + separator, ) # Non-streaming response. total_segments = [] @@ -500,7 +523,6 @@ class OpenAISpeechToText(OpenAIServing): "translate": TranslationSegment, } segment_class: type[SpeechToTextSegment] = segments_types[self.task_type] - text = "" chunk_size_in_s = self.asr_config.max_audio_clip_s if chunk_size_in_s is None: assert len(list_result_generator) == 1, ( @@ -528,7 +550,7 @@ class OpenAISpeechToText(OpenAIServing): else: raw_text = op.outputs[0].text text_parts.append(self.model_cls.post_process_output(raw_text)) - text = "".join(text_parts) + text = separator.join(text_parts) if self.task_type == "transcribe": final_response: ResponseType # add usage in TranscriptionResponse. @@ -581,6 +603,7 @@ class OpenAISpeechToText(OpenAIServing): | type[TranslationResponseStreamChoice], stream_response_class: type[TranscriptionStreamResponse] | type[TranslationStreamResponse], + separator: str, ) -> AsyncGenerator[str, None]: created_time = int(time.time()) model_name = request.model @@ -597,6 +620,7 @@ class OpenAISpeechToText(OpenAIServing): try: for result_generator in list_result_generator: + beginning_of_chunk = True async for res in result_generator: # On first result. if res.prompt_token_ids is not None: @@ -614,6 +638,14 @@ class OpenAISpeechToText(OpenAIServing): assert len(res.outputs) == 1 output = res.outputs[0] + # dont add separator to the first chunk + if ( + result_generator is not list_result_generator[0] + and beginning_of_chunk + ): + output.text = separator + output.text + beginning_of_chunk = False + # TODO: For models that output structured formats (e.g., # Qwen3-ASR with "language X" prefix), streaming # would need buffering to strip the prefix properly since diff --git a/vllm/entrypoints/pooling/__init__.py b/vllm/entrypoints/pooling/__init__.py index fb0c10e6f4f..1980750ec20 100644 --- a/vllm/entrypoints/pooling/__init__.py +++ b/vllm/entrypoints/pooling/__init__.py @@ -67,20 +67,18 @@ def init_pooling_state( from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.pooling.classify.serving import ServingClassification from vllm.entrypoints.pooling.embed.serving import ServingEmbedding - from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling + from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores from vllm.tasks import POOLING_TASKS model_config = engine_client.model_config - resolved_chat_template = load_chat_template(args.chat_template) state.serving_pooling = ( ( - OpenAIServingPooling( + ServingPooling( engine_client, state.openai_serving_models, - state.openai_serving_render, supported_tasks=supported_tasks, request_logger=request_logger, chat_template=resolved_chat_template, diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index fd4c076cdda..83e82664ef1 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -4,8 +4,8 @@ from collections.abc import Sequence from typing import Any, Final -from vllm import PoolingRequestOutput, PromptType -from vllm.config import ModelConfig +from vllm import PoolingParams, PoolingRequestOutput, PromptType +from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateConfig, @@ -33,11 +33,12 @@ class PoolingIOProcessor: def __init__( self, - model_config: ModelConfig, + vllm_config: VllmConfig, renderer: BaseRenderer, chat_template_config: ChatTemplateConfig, ): - self.model_config = model_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config self.renderer = renderer self.chat_template = chat_template_config.chat_template @@ -48,12 +49,12 @@ class PoolingIOProcessor: chat_template_config.trust_request_chat_template ) - def create_pooling_params(self, request): - return request.to_pooling_params() - ####################################### # online APIs + def create_pooling_params(self, request): + return request.to_pooling_params() + def pre_process_online(self, ctx: PoolingServeContext): request = ctx.request @@ -71,7 +72,7 @@ class PoolingIOProcessor: default_template_kwargs=None, ) elif isinstance(request, PoolingCompletionLikeRequest): - engine_inputs = self._preprocess_completion_online( + engine_inputs = self._preprocess_cmpl_online( request, prompt_input=request.input, prompt_embeds=None, @@ -81,35 +82,25 @@ class PoolingIOProcessor: ctx.engine_inputs = engine_inputs - async def pre_process_online_async(self, ctx: PoolingServeContext): - self.pre_process_online(ctx) - def post_process_online( self, ctx: PoolingServeContext, ): pass - async def post_process_online_async( - self, - ctx: PoolingServeContext, - ): - self.post_process_online(ctx) - ####################################### # offline APIs def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: - assert not isinstance(ctx.prompts, ScoringData) + assert not isinstance(ctx.prompts, ScoringData) and not ( + isinstance(ctx.prompts, dict) and "data" in ctx.prompts + ) + + prompts_seq = prompt_to_seq(ctx.prompts) tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) - return self._preprocess_completion_offline( - prompts=ctx.prompts, tok_params=tok_params - ) - - async def pre_process_offline_async(self, ctx: OfflineInputsContext): - return self.pre_process_offline(ctx) + return self._preprocess_cmpl_offline(prompts=prompts_seq, tok_params=tok_params) def post_process_offline( self, @@ -117,16 +108,10 @@ class PoolingIOProcessor: ) -> list[PoolingRequestOutput]: return ctx.outputs - async def post_process_offline_async( - self, - ctx: OfflineOutputsContext, - ) -> list[PoolingRequestOutput]: - return self.post_process_offline(ctx) - ####################################### # helpers - def _preprocess_completion_online( + def _preprocess_cmpl_online( self, request: RendererRequest, prompt_input: str | list[str] | list[int] | list[list[int]] | None, @@ -204,7 +189,7 @@ class PoolingIOProcessor: return conversation, [engine_input] - def _preprocess_completion_offline( + def _preprocess_cmpl_offline( self, prompts: PromptType | Sequence[PromptType], tok_params: TokenizeParams, @@ -243,3 +228,19 @@ class PoolingIOProcessor: "Refused request with untrusted chat template." ) return None + + def _params_to_seq( + self, + params: PoolingParams | Sequence[PoolingParams], + num_requests: int, + ) -> Sequence[PoolingParams]: + if isinstance(params, Sequence): + if len(params) != num_requests: + raise ValueError( + f"The lengths of prompts ({num_requests}) " + f"and params ({len(params)}) must be the same." + ) + + return params + + return [params] * num_requests diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index 90554aa634b..c65bedc70f1 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -1,15 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod from collections.abc import AsyncGenerator, Mapping +from concurrent.futures import Executor from http import HTTPStatus from typing import ClassVar +import torch from fastapi import Request from fastapi.responses import Response from starlette.datastructures import Headers from vllm import PoolingParams, PoolingRequestOutput, envs -from vllm.config import ModelConfig +from vllm.config import VllmConfig from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatTemplateConfig, @@ -30,12 +34,12 @@ from vllm.tracing import ( log_tracing_disabled_warning, ) from vllm.utils import random_uuid -from vllm.utils.async_utils import merge_async_iterators +from vllm.utils.async_utils import make_async, merge_async_iterators from .io_processor import PoolingIOProcessor -class PoolingServing: +class PoolingServingBase(ABC): request_id_prefix: ClassVar[str] def __init__( @@ -50,10 +54,11 @@ class PoolingServing: return_tokens_as_token_ids: bool = False, log_error_stack: bool = False, ): - super().__init__() self.engine_client = engine_client self.models = models self.model_config = models.model_config + self.renderer = models.renderer + self.vllm_config = engine_client.vllm_config self.max_model_len = self.model_config.max_model_len self.request_logger = request_logger self.return_tokens_as_token_ids = return_tokens_as_token_ids @@ -63,34 +68,48 @@ class PoolingServing: chat_template_content_format=chat_template_content_format, trust_request_chat_template=trust_request_chat_template, ) - self.io_processor = self.init_io_processor( - model_config=models.model_config, - renderer=models.renderer, - chat_template_config=self.chat_template_config, - ) - def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> PoolingIOProcessor: - raise NotImplementedError + # Shared thread pool executor for preprocessing and postprocessing. + self._executor: Executor = models.renderer._executor + self._preprocessing_async = make_async( + self._preprocessing, executor=self._executor + ) + self._postprocessing_async = make_async( + self._postprocessing, executor=self._executor + ) async def __call__( self, request: AnyPoolingRequest, raw_request: Request | None = None, ) -> Response: - ctx = await self._init_ctx(request, raw_request) - await self.io_processor.pre_process_online_async(ctx) + io_processor = self.get_io_processor(request) + ctx = await self._init_ctx(io_processor, request, raw_request) + await self._preprocessing_async(io_processor, ctx) await self._prepare_generators(ctx) await self._collect_batch(ctx) - await self.io_processor.post_process_online_async(ctx) - return await self._build_response(ctx) + return await self._postprocessing_async(io_processor, ctx) + + @abstractmethod + def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: + raise NotImplementedError + + @torch.inference_mode() + def _preprocessing( + self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext + ): + return io_processor.pre_process_online(ctx) + + @torch.inference_mode() + def _postprocessing( + self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext + ): + io_processor.post_process_online(ctx) + return self._build_response(ctx) async def _init_ctx( self, + io_processor: PoolingIOProcessor, request: AnyPoolingRequest, raw_request: Request | None = None, ): @@ -98,10 +117,12 @@ class PoolingServing: request_id = f"{self.request_id_prefix}-{self._base_request_id(raw_request)}" await self._check_model(request) + pooling_params = io_processor.create_pooling_params(request) ctx = PoolingServeContext( request=request, raw_request=raw_request, model_name=model_name, + pooling_params=pooling_params, request_id=request_id, ) @@ -124,10 +145,8 @@ class PoolingServing: else await self._get_trace_headers(ctx.raw_request.headers) ) - if ctx.pooling_params is None: - pooling_params = self.io_processor.create_pooling_params(ctx.request) - else: - pooling_params = ctx.pooling_params + assert ctx.pooling_params is not None + pooling_params = ctx.pooling_params if isinstance(pooling_params, list): for params in pooling_params: @@ -190,7 +209,8 @@ class PoolingServing: ctx.final_res_batch = [res for res in final_res_batch if res is not None] - async def _build_response( + @abstractmethod + def _build_response( self, ctx: PoolingServeContext, ) -> Response: @@ -247,6 +267,7 @@ class PoolingServing: "greater than max_model_len." " Please request a smaller truncation size." ) + return None async def _get_trace_headers( @@ -355,3 +376,26 @@ class PoolingServing: params=params, lora_request=lora_request, ) + + +class PoolingServing(PoolingServingBase, ABC): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.io_processor = self.init_io_processor( + vllm_config=self.vllm_config, + renderer=self.renderer, + chat_template_config=self.chat_template_config, + ) + + @abstractmethod + def init_io_processor( + self, + vllm_config: VllmConfig, + renderer: BaseRenderer, + chat_template_config: ChatTemplateConfig, + ) -> PoolingIOProcessor: + raise NotImplementedError + + def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: + return self.io_processor diff --git a/vllm/entrypoints/pooling/classify/io_processor.py b/vllm/entrypoints/pooling/classify/io_processor.py index ee73207dff5..9bb3774ab0b 100644 --- a/vllm/entrypoints/pooling/classify/io_processor.py +++ b/vllm/entrypoints/pooling/classify/io_processor.py @@ -5,4 +5,8 @@ from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor class ClassifyIOProcessor(PoolingIOProcessor): - name = "classification" + name = "classify" + + +class TokenClassifyIOProcessor(PoolingIOProcessor): + name = "token_classify" diff --git a/vllm/entrypoints/pooling/classify/serving.py b/vllm/entrypoints/pooling/classify/serving.py index 24d4f9aacff..a48ec819b7f 100644 --- a/vllm/entrypoints/pooling/classify/serving.py +++ b/vllm/entrypoints/pooling/classify/serving.py @@ -6,14 +6,11 @@ from typing import TypeAlias import numpy as np from fastapi.responses import JSONResponse -from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.entrypoints.pooling.typing import PoolingServeContext from vllm.logger import init_logger from vllm.outputs import ClassificationOutput -from vllm.renderers import BaseRenderer from .io_processor import ClassifyIOProcessor from .protocol import ( @@ -31,19 +28,10 @@ ClassificationServeContext: TypeAlias = PoolingServeContext[ClassificationReques class ServingClassification(PoolingServing): request_id_prefix = "classify" - def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> ClassifyIOProcessor: - return ClassifyIOProcessor( - model_config=model_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) + def init_io_processor(self, *args, **kwargs) -> ClassifyIOProcessor: + return ClassifyIOProcessor(*args, **kwargs) - async def _build_response( + def _build_response( self, ctx: ClassificationServeContext, ) -> JSONResponse: diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index 614f8e0d9d0..09016253f09 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -24,7 +24,13 @@ from vllm.entrypoints.pooling.embed.protocol import ( EmbeddingChatRequest, EmbeddingCompletionRequest, ) -from vllm.entrypoints.pooling.typing import PoolingServeContext +from vllm.entrypoints.pooling.scoring.io_processor import JinaRankingIOProcessorMixin +from vllm.entrypoints.pooling.typing import ( + OfflineInputsContext, + PoolingChatLikeRequest, + PoolingCompletionLikeRequest, + PoolingServeContext, +) from vllm.inputs import EngineInput, tokens_input from vllm.logger import init_logger from vllm.outputs import PoolingOutput, PoolingRequestOutput @@ -37,7 +43,7 @@ logger = init_logger(__name__) class EmbedIOProcessor(PoolingIOProcessor): - name = "embedding" + name = "embed" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -464,7 +470,7 @@ class EmbedIOProcessor(PoolingIOProcessor): truncate_prompt_tokens=truncate_prompt_tokens, truncation_side=truncation_side, ) - return self._preprocess_completion_online( + return self._preprocess_cmpl_online( proxy, prompt_input=proxy.input, prompt_embeds=None ) @@ -549,3 +555,52 @@ class EmbedIOProcessor(PoolingIOProcessor): request = ctx.request if request.truncate == "NONE" and request.max_tokens is not None: self._check_cohere_max_tokens(ctx.final_res_batch, request.max_tokens) + + +class TokenEmbedIOProcessor(PoolingIOProcessor): + name = "token_embed" + + +class JinaRankingTokenEmbedIOProcessor( + TokenEmbedIOProcessor, JinaRankingIOProcessorMixin +): + def pre_process_online(self, ctx: PoolingServeContext): + request = ctx.request + if isinstance(request, PoolingCompletionLikeRequest): + prompts = request.input + if not isinstance(prompts, Sequence) or len(prompts) < 2: + raise ValueError("The JinaForRanking model requires at least 2 inputs.") + + text_prompts = self.ensure_str(prompts) + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + prompt_input = self.format_docs_prompts_func( + query=text_prompts[-1], docs=text_prompts[:-1] + ) + + engine_inputs = self._preprocess_cmpl_online( + request, + prompt_input=prompt_input, + prompt_embeds=None, + ) + elif isinstance(request, PoolingChatLikeRequest): + raise ValueError("The JinaForRanking does not support chat Request.") + else: + raise ValueError(f"Invalid {self.name} request type") + + ctx.engine_inputs = engine_inputs + + def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2: + raise ValueError("The JinaForRanking model requires at least 2 inputs.") + + text_prompts = self.ensure_str(ctx.prompts) + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + ctx.prompts = self.format_docs_prompts_func( + query=text_prompts[-1], docs=text_prompts[:-1] + ) + + return super().pre_process_offline(ctx) diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index f0c33164591..9389309efc7 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -8,8 +8,6 @@ from typing import Literal, TypeAlias, cast from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never -from vllm.config import ModelConfig -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor @@ -33,12 +31,10 @@ from vllm.entrypoints.pooling.utils import ( ) from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput -from vllm.renderers import BaseRenderer from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) -JSONResponseCLS = get_json_response_cls() EmbeddingServeContext: TypeAlias = PoolingServeContext[EmbeddingRequest] @@ -49,27 +45,23 @@ class ServingEmbedding(PoolingServing): request_id_prefix = "embd" io_processor: EmbedIOProcessor - def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> EmbedIOProcessor: - return EmbedIOProcessor( - model_config=model_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) - async def _build_response( + self.json_response_cls = get_json_response_cls() + + def init_io_processor(self, *args, **kwargs) -> EmbedIOProcessor: + return EmbedIOProcessor(*args, **kwargs) + + def _build_response( self, ctx: PoolingServeContext, ) -> Response: if isinstance(ctx.request, CohereEmbedRequest): return self._build_cohere_response_from_ctx(ctx) - return await self._build_openai_response(ctx) + return self._build_openai_response(ctx) - async def _build_openai_response( + def _build_openai_response( self, ctx: EmbeddingServeContext, ) -> JSONResponse | StreamingResponse: @@ -149,7 +141,7 @@ class ServingEmbedding(PoolingServing): data=items, usage=usage, ) - return JSONResponseCLS(content=response.model_dump()) + return self.json_response_cls(content=response.model_dump()) def _openai_bytes_response( self, @@ -190,8 +182,8 @@ class ServingEmbedding(PoolingServing): media_type=response.media_type, ) - @staticmethod def _build_cohere_response_from_ctx( + self, ctx: PoolingServeContext, ) -> JSONResponse: request = ctx.request @@ -218,4 +210,4 @@ class ServingEmbedding(PoolingServing): ), ), ) - return JSONResponse(content=response.model_dump(exclude_none=True)) + return self.json_response_cls(content=response.model_dump(exclude_none=True)) diff --git a/vllm/entrypoints/pooling/io_processor_factories.py b/vllm/entrypoints/pooling/io_processor_factories.py index 71033bd2398..5e67d069da7 100644 --- a/vllm/entrypoints/pooling/io_processor_factories.py +++ b/vllm/entrypoints/pooling/io_processor_factories.py @@ -1,42 +1,76 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from vllm.config import ModelConfig +from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ChatTemplateConfig -from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor -from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessors -from vllm.entrypoints.pooling.utils import enable_scoring_api +from vllm.plugins.io_processors import has_io_processor from vllm.renderers import BaseRenderer from vllm.tasks import SupportedTask +from .base.io_processor import PoolingIOProcessor +from .utils import enable_scoring_api + def init_pooling_io_processors( supported_tasks: tuple[SupportedTask, ...], - model_config: ModelConfig, + vllm_config: VllmConfig, renderer: BaseRenderer, chat_template_config: ChatTemplateConfig, ) -> dict[str, PoolingIOProcessor]: - processors: list[tuple[str, type[PoolingIOProcessor]]] = [] + model_config = vllm_config.model_config + processors: dict[str, type[PoolingIOProcessor]] = {} + if "classify" in supported_tasks: - from vllm.entrypoints.pooling.classify.io_processor import ClassifyIOProcessor + from .classify.io_processor import ClassifyIOProcessor + + processors["classify"] = ClassifyIOProcessor + + if "token_classify" in supported_tasks: + from .classify.io_processor import TokenClassifyIOProcessor + + processors["token_classify"] = TokenClassifyIOProcessor - processors.append(("classify", ClassifyIOProcessor)) if "embed" in supported_tasks: - from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor + from .embed.io_processor import EmbedIOProcessor - processors.append(("embed", EmbedIOProcessor)) + processors["embed"] = EmbedIOProcessor + + if "token_embed" in supported_tasks: + from .embed.io_processor import TokenEmbedIOProcessor + + processors["token_embed"] = TokenEmbedIOProcessor + + if has_io_processor( + vllm_config, + model_config.io_processor_plugin, + ): + from .pooling.io_processor import PluginWithIOProcessorPlugins + + processors["plugin"] = PluginWithIOProcessorPlugins + elif "plugin" in supported_tasks: + from .pooling.io_processor import PluginWithoutIOProcessorPlugins + + processors["plugin"] = PluginWithoutIOProcessorPlugins if enable_scoring_api(supported_tasks, model_config): score_type = model_config.score_type + from .scoring.io_processor import ScoringIOProcessors + if score_type is not None and score_type in ScoringIOProcessors: - processors.append((score_type, ScoringIOProcessors[score_type])) + processors[score_type] = ScoringIOProcessors[score_type] + + if model_config.architecture == "JinaForRanking": + from .embed.io_processor import JinaRankingTokenEmbedIOProcessor + from .scoring.io_processor import ScoringIOProcessors + + processors["token_embed"] = JinaRankingTokenEmbedIOProcessor + processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"] return { task: processor_cls( - model_config=model_config, + vllm_config=vllm_config, renderer=renderer, chat_template_config=chat_template_config, ) - for task, processor_cls in processors + for task, processor_cls in processors.items() } diff --git a/vllm/entrypoints/pooling/pooling/api_router.py b/vllm/entrypoints/pooling/pooling/api_router.py index f63a8edf6ca..a08570038c3 100644 --- a/vllm/entrypoints/pooling/pooling/api_router.py +++ b/vllm/entrypoints/pooling/pooling/api_router.py @@ -3,24 +3,17 @@ from http import HTTPStatus from fastapi import APIRouter, Depends, Request -from fastapi.responses import JSONResponse, StreamingResponse -from typing_extensions import assert_never from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.pooling.protocol import ( - IOProcessorResponse, - PoolingBytesResponse, - PoolingRequest, - PoolingResponse, -) -from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling +from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest +from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.utils import load_aware_call, with_cancellation router = APIRouter() -def pooling(request: Request) -> OpenAIServingPooling | None: +def pooling(request: Request) -> ServingPooling | None: return request.app.state.serving_pooling @@ -39,19 +32,4 @@ async def create_pooling(request: PoolingRequest, raw_request: Request): if handler is None: raise NotImplementedError("The model does not support Pooling API") - generator = await handler.create_pooling(request, raw_request) - - if isinstance(generator, ErrorResponse): - return JSONResponse( - content=generator.model_dump(), status_code=generator.error.code - ) - elif isinstance(generator, (PoolingResponse, IOProcessorResponse)): - return JSONResponse(content=generator.model_dump()) - elif isinstance(generator, PoolingBytesResponse): - return StreamingResponse( - content=generator.content, - headers=generator.headers, - media_type=generator.media_type, - ) - - assert_never(generator) + return await handler(request, raw_request) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py new file mode 100644 index 00000000000..31f860144ea --- /dev/null +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence +from typing import Any + +from vllm import PoolingParams, PoolingRequestOutput +from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.inputs import EngineInput +from vllm.logger import init_logger +from vllm.plugins.io_processors import get_io_processor +from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq + +from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext +from .protocol import IOProcessorRequest, IOProcessorResponse + +logger = init_logger(__name__) + + +class PluginWithoutIOProcessorPlugins(PoolingIOProcessor): + name = "plugin" + + +class PluginWithIOProcessorPlugins(PoolingIOProcessor): + """IO Processor plugins are a feature that allows pre- and post-processing + of the model input and output for pooling models.""" + + name = "plugin" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + io_processor = get_io_processor( + self.vllm_config, + self.renderer, + self.model_config.io_processor_plugin, + ) + + assert io_processor is not None + self.io_processor = io_processor + + ####################################### + # online APIs + + def pre_process_online(self, ctx: PoolingServeContext): + assert isinstance(ctx.request, IOProcessorRequest) + + validated_prompt = self.io_processor.parse_data(ctx.request.data) + + raw_prompts = self.io_processor.pre_process( + prompt=validated_prompt, request_id=ctx.request_id + ) + + parsed_prompts = [ + ( + prompt + if isinstance(prompt, bytes) + else parse_model_prompt(self.model_config, prompt) + ) + for prompt in prompt_to_seq(raw_prompts) + ] + + tok_params = ctx.request.build_tok_params(self.model_config) + + ctx.engine_inputs = self.renderer.render_cmpl( + parsed_prompts, + tok_params, + prompt_extras={ + k: v + for k in ("mm_processor_kwargs", "cache_salt") + if (v := getattr(ctx.request, k, None)) is not None + }, + ) + + pooling_params = self.io_processor.merge_pooling_params() + if pooling_params.task is None: + pooling_params.task = "plugin" + ctx.pooling_params = pooling_params + + def post_process_online( + self, + ctx: PoolingServeContext, + ): + output = self.io_processor.post_process( + ctx.final_res_batch, + request_id=ctx.request_id, + ) + + if callable( + output_to_response := getattr(self.io_processor, "output_to_response", None) + ): + logger.warning_once( + "`IOProcessor.output_to_response` is deprecated. To ensure " + "consistency between offline and online APIs, " + "`IOProcessorResponse` will become a transparent wrapper " + "around output data from v0.19 onwards.", + ) + + if hasattr(output, "request_id") and output.request_id is None: + output.request_id = ctx.request_id # type: ignore + + ctx.response = output_to_response(output) # type: ignore + else: + ctx.response = IOProcessorResponse(request_id=ctx.request_id, data=output) + + ####################################### + # offline APIs + + def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts + assert ctx.pooling_params is not None + + # Validate the request data is valid for the loaded plugin + prompt_data = ctx.prompts.get("data") + if prompt_data is None: + raise ValueError( + "The 'data' field of the prompt is expected to contain " + "the prompt data and it cannot be None. " + "Refer to the documentation of the IOProcessor " + "in use for more details." + ) + validated_prompt = self.io_processor.parse_data(prompt_data) + + # obtain the actual model prompts from the pre-processor + prompts = self.io_processor.pre_process(prompt=validated_prompt) + prompts_seq = prompt_to_seq(prompts) + + params_seq: list[PoolingParams] = [ + self.io_processor.merge_pooling_params(param) + for param in self._params_to_seq( + ctx.pooling_params, + len(prompts_seq), + ) + ] + for p in params_seq: + if p.task is None: + p.task = "plugin" + + ctx.pooling_params = params_seq + ctx.prompts = prompts_seq + return super().pre_process_offline(ctx) + + def post_process_offline( + self, + ctx: OfflineOutputsContext, + ) -> list[PoolingRequestOutput]: + processed_outputs = self.io_processor.post_process(ctx.outputs) + + return [ + PoolingRequestOutput[Any]( + request_id="", + outputs=processed_outputs, + num_cached_tokens=getattr(processed_outputs, "num_cached_tokens", 0), + prompt_token_ids=[], + finished=True, + ) + ] diff --git a/vllm/entrypoints/pooling/pooling/serving.py b/vllm/entrypoints/pooling/pooling/serving.py index 4706684f363..ea8eff6eb30 100644 --- a/vllm/entrypoints/pooling/pooling/serving.py +++ b/vllm/entrypoints/pooling/pooling/serving.py @@ -1,252 +1,140 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import asyncio -import json -import time -from collections.abc import AsyncGenerator, Callable, Sequence -from functools import partial -from typing import Final, Literal, cast -from fastapi import Request +import json +from collections.abc import Callable +from functools import partial +from typing import Literal, cast + +from fastapi.responses import JSONResponse, Response, StreamingResponse from typing_extensions import assert_never -from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption -from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo -from vllm.entrypoints.openai.engine.serving import OpenAIServing -from vllm.entrypoints.openai.models.serving import OpenAIServingModels +from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor +from vllm.entrypoints.pooling.base.serving import PoolingServingBase +from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors from vllm.entrypoints.pooling.pooling.protocol import ( IOProcessorRequest, - IOProcessorResponse, PoolingBytesResponse, - PoolingChatRequest, - PoolingCompletionRequest, PoolingRequest, PoolingResponse, PoolingResponseData, ) +from vllm.entrypoints.pooling.typing import AnyPoolingRequest, PoolingServeContext from vllm.entrypoints.pooling.utils import ( encode_pooling_bytes, encode_pooling_output_base64, encode_pooling_output_float, + get_json_response_cls, ) -from vllm.entrypoints.serve.render.serving import OpenAIServingRender -from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput -from vllm.renderers.inputs.preprocess import prompt_to_seq from vllm.tasks import SupportedTask -from vllm.utils.async_utils import merge_async_iterators -from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness +from vllm.utils.serial_utils import EmbedDType, Endianness logger = init_logger(__name__) -class OpenAIServingPooling(OpenAIServing): +class ServingPooling(PoolingServingBase): + request_id_prefix = "pooling" + def __init__( self, - engine_client: EngineClient, - models: OpenAIServingModels, - openai_serving_render: OpenAIServingRender, + *args, supported_tasks: tuple[SupportedTask, ...], - *, - request_logger: RequestLogger | None, - chat_template: str | None, - chat_template_content_format: ChatTemplateContentFormatOption, - trust_request_chat_template: bool = False, - ) -> None: - super().__init__( - engine_client=engine_client, - models=models, - request_logger=request_logger, - ) + **kwargs, + ): + super().__init__(*args, **kwargs) + self.supported_tasks = supported_tasks self.pooling_task = self.model_config.get_pooling_task(supported_tasks) - self.openai_serving_render = openai_serving_render - self.chat_template = chat_template - self.chat_template_content_format: Final = chat_template_content_format - self.trust_request_chat_template = trust_request_chat_template + self.io_processors = init_pooling_io_processors( + supported_tasks=supported_tasks, + vllm_config=self.vllm_config, + renderer=self.renderer, + chat_template_config=self.chat_template_config, + ) + self.json_response_cls = get_json_response_cls() - async def create_pooling( - self, - request: PoolingRequest, - raw_request: Request | None = None, - ) -> PoolingResponse | IOProcessorResponse | PoolingBytesResponse | ErrorResponse: - """ - See https://platform.openai.com/docs/api-reference/embeddings/create - for the API specification. This API mimics the OpenAI Embedding API. - """ - error_check_ret = await self._check_model(request) - if error_check_ret is not None: - return error_check_ret + def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: + assert isinstance(request, PoolingRequest) + pooling_task = self._verify_pooling_task(request) + return self.io_processors[pooling_task] - model_name = self.models.model_name() - - request_id = f"pool-{self._base_request_id(raw_request)}" - created_time = int(time.time()) - - lora_request = self._maybe_get_adapters(request) + def _verify_pooling_task(self, request: PoolingRequest) -> str: + if getattr(request, "dimensions", None) is not None: + raise ValueError("dimensions is currently not supported") if request.task is None: request.task = self.pooling_task - if getattr(request, "dimensions", None) is not None: - return self.create_error_response("dimensions is currently not supported") + if isinstance(request, IOProcessorRequest): + request.task = "plugin" + + assert request.task is not None + pooling_task = request.task # plugin task uses io_processor.parse_request to verify inputs - if request.task != "plugin" and request.task != self.pooling_task: - if request.task not in self.supported_tasks: + if pooling_task != "plugin" and pooling_task != self.pooling_task: + if pooling_task not in self.io_processors: raise ValueError( - f"Unsupported task: {request.task!r} " + f"Unsupported task: {pooling_task!r} " f"Supported tasks: {self.supported_tasks}" ) else: logger.warning_once( "Pooling multitask support is deprecated and will be removed " "in v0.20. When the default pooling task is not what you want, you " - 'need to manually specify it via --pooler-config.task "%s". ', - request.task, + "need to manually specify it via --pooler-config.task %s. ", + pooling_task, ) - engine_inputs: Sequence[EngineInput] - if use_io_processor := isinstance(request, IOProcessorRequest): - if self.io_processor is None: - raise ValueError( - "No IOProcessor plugin installed. Please refer " - "to the documentation and to the " - "'prithvi_geospatial_mae_io_processor' " - "offline inference example for more details." - ) - - validated_prompt = self.io_processor.parse_data(request.data) - - raw_prompts = await self.io_processor.pre_process_async( - prompt=validated_prompt, request_id=request_id - ) - engine_inputs = await self.openai_serving_render.preprocess_cmpl( - request, - prompt_to_seq(raw_prompts), - ) - elif isinstance(request, PoolingChatRequest): - error_check_ret = self.openai_serving_render.validate_chat_template( - request_chat_template=request.chat_template, - chat_template_kwargs=request.chat_template_kwargs, - trust_request_chat_template=self.trust_request_chat_template, - ) - if error_check_ret is not None: - return error_check_ret - - _, engine_inputs = await self.openai_serving_render.preprocess_chat( - request, - request.messages, - default_template=self.chat_template, - default_template_content_format=self.chat_template_content_format, - default_template_kwargs=None, - ) - elif isinstance(request, PoolingCompletionRequest): - engine_inputs = await self.openai_serving_render.preprocess_completion( - request, - prompt_input=request.input, - prompt_embeds=None, - ) - else: - raise ValueError(f"Unsupported request of type {type(request)}") - - # Schedule the request and get the result generator. - generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] - if use_io_processor: - assert self.io_processor is not None - - pooling_params = self.io_processor.merge_pooling_params() - if pooling_params.task is None: - pooling_params.task = "plugin" - else: - pooling_params = request.to_pooling_params() # type: ignore - - for i, engine_input in enumerate(engine_inputs): - request_id_item = f"{request_id}-{i}" - - self._log_inputs( - request_id_item, - engine_input, - params=pooling_params, - lora_request=lora_request, + if pooling_task == "plugin" and "plugin" not in self.io_processors: + raise ValueError( + "No IOProcessor plugin installed. Please refer " + "to the documentation and to the " + "'prithvi_geospatial_mae_io_processor' " + "offline inference example for more details." ) - trace_headers = ( - None - if raw_request is None - else await self._get_trace_headers(raw_request.headers) + return pooling_task + + def _build_response( + self, + ctx: PoolingServeContext, + ) -> Response: + if ctx.response is not None: + # for IOProcessorResponse + return self.json_response_cls(content=ctx.response.model_dump()) + + encoding_format = ctx.request.encoding_format + embed_dtype = ctx.request.embed_dtype + endianness = ctx.request.endianness + + if encoding_format == "float" or encoding_format == "base64": + return self.request_output_to_pooling_json_response( + ctx.final_res_batch, + ctx.request_id, + ctx.created_time, + ctx.model_name, + encoding_format, + embed_dtype, + endianness, ) - generator = self.engine_client.encode( - engine_input, - pooling_params, - request_id_item, - lora_request=lora_request, - trace_headers=trace_headers, - priority=request.priority, + if encoding_format == "bytes" or encoding_format == "bytes_only": + return self.request_output_to_pooling_bytes_response( + ctx.final_res_batch, + ctx.request_id, + ctx.created_time, + ctx.model_name, + encoding_format, + embed_dtype, + endianness, ) - generators.append(generator) - - result_generator = merge_async_iterators(*generators) - - if use_io_processor: - assert self.io_processor is not None - output = await self.io_processor.post_process_async( - result_generator, - request_id=request_id, - ) - - if callable( - output_to_response := getattr( - self.io_processor, "output_to_response", None - ) - ): - logger.warning_once( - "`IOProcessor.output_to_response` is deprecated. To ensure " - "consistency between offline and online APIs, " - "`IOProcessorResponse` will become a transparent wrapper " - "around output data from v0.19 onwards.", - ) - - if hasattr(output, "request_id") and output.request_id is None: - output.request_id = request_id # type: ignore - - return output_to_response(output) # type: ignore - - return IOProcessorResponse(request_id=request_id, data=output) - - assert isinstance(request, (PoolingCompletionRequest, PoolingChatRequest)) - num_prompts = len(engine_inputs) - - # Non-streaming response - final_res_batch: list[PoolingRequestOutput | None] - final_res_batch = [None] * num_prompts - try: - async for i, res in result_generator: - final_res_batch[i] = res - - assert all(final_res is not None for final_res in final_res_batch) - - final_res_batch_checked = cast(list[PoolingRequestOutput], final_res_batch) - - response = self.request_output_to_pooling_response( - final_res_batch_checked, - request_id, - created_time, - model_name, - request.encoding_format, - request.embed_dtype, - request.endianness, - ) - except asyncio.CancelledError: - return self.create_error_response("Client disconnected") - - return response + assert_never(encoding_format) def request_output_to_pooling_json_response( self, @@ -257,7 +145,7 @@ class OpenAIServingPooling(OpenAIServing): encoding_format: Literal["float", "base64"], embed_dtype: EmbedDType, endianness: Endianness, - ) -> PoolingResponse: + ) -> JSONResponse: encode_fn = cast( Callable[[PoolingRequestOutput], list[float] | str], ( @@ -289,13 +177,14 @@ class OpenAIServingPooling(OpenAIServing): total_tokens=num_prompt_tokens, ) - return PoolingResponse( + response = PoolingResponse( id=request_id, created=created_time, model=model_name, data=items, usage=usage, ) + return self.json_response_cls(content=response.model_dump()) def request_output_to_pooling_bytes_response( self, @@ -306,7 +195,7 @@ class OpenAIServingPooling(OpenAIServing): encoding_format: Literal["bytes", "bytes_only"], embed_dtype: EmbedDType, endianness: Endianness, - ) -> PoolingBytesResponse: + ) -> StreamingResponse: content, items, usage = encode_pooling_bytes( pooling_outputs=final_res_batch, embed_dtype=embed_dtype, @@ -329,38 +218,10 @@ class OpenAIServingPooling(OpenAIServing): } ) - return PoolingBytesResponse(content=content, headers=headers) + response = PoolingBytesResponse(content=content, headers=headers) - def request_output_to_pooling_response( - self, - final_res_batch: list[PoolingRequestOutput], - request_id: str, - created_time: int, - model_name: str, - encoding_format: EncodingFormat, - embed_dtype: EmbedDType, - endianness: Endianness, - ) -> PoolingResponse | PoolingBytesResponse: - if encoding_format == "float" or encoding_format == "base64": - return self.request_output_to_pooling_json_response( - final_res_batch, - request_id, - created_time, - model_name, - encoding_format, - embed_dtype, - endianness, - ) - - if encoding_format == "bytes" or encoding_format == "bytes_only": - return self.request_output_to_pooling_bytes_response( - final_res_batch, - request_id, - created_time, - model_name, - encoding_format, - embed_dtype, - endianness, - ) - - assert_never(encoding_format) + return StreamingResponse( + content=response.content, + headers=response.headers, + media_type=response.media_type, + ) diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index c520eb5ceb3..549bae2775d 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -25,8 +25,10 @@ from .typing import ScoreData, ScoreInput, ScoringData from .utils import ( compress_token_type_ids, compute_maxsim_score, + get_num_special_tokens_for_pair, parse_score_data, score_data_to_prompts, + truncate_text_to_tokens, validate_score_input, ) @@ -48,6 +50,64 @@ class ScoringIOProcessor(PoolingIOProcessor): def create_pooling_params(self, request): return request.to_pooling_params(self.pooling_task) + def _validate_token_limit(self, value: int, name: str) -> None: + if value < 0: + raise ValueError(f"{name} must be a non-negative integer") + if value >= self.model_config.max_model_len: + raise ValueError( + f"{name} ({value}) must be less " + f"than max_model_len ({self.model_config.max_model_len})." + ) + + def _get_token_limits( + self, + request: ScoringRequest | None = None, + pooling_params: PoolingParams | None = None, + ) -> tuple[int, int]: + """Extract and validate token limits from request or pooling_params.""" + if request is not None: + max_tokens_per_query = getattr(request, "max_tokens_per_query", 0) + max_tokens_per_doc = getattr(request, "max_tokens_per_doc", 0) + else: + extra = ( + (pooling_params.extra_kwargs or {}) + if pooling_params is not None + else {} + ) + max_tokens_per_query = extra.get("max_tokens_per_query", 0) + max_tokens_per_doc = extra.get("max_tokens_per_doc", 0) + + if max_tokens_per_query != 0: + self._validate_token_limit(max_tokens_per_query, "max_tokens_per_query") + if max_tokens_per_doc != 0: + self._validate_token_limit(max_tokens_per_doc, "max_tokens_per_doc") + return max_tokens_per_query, max_tokens_per_doc + + def _truncate_scoring_data( + self, + scoring_data: ScoringData, + max_tokens_per_query: int = 0, + max_tokens_per_doc: int = 0, + ) -> ScoringData: + """Truncate query/document texts to token limits.""" + data_1 = scoring_data.data_1 + data_2 = scoring_data.data_2 + if max_tokens_per_query > 0: + data_1 = [ + truncate_text_to_tokens(d, self.tokenizer, max_tokens_per_query) + if isinstance(d, str) + else d + for d in data_1 + ] + if max_tokens_per_doc > 0: + data_2 = [ + truncate_text_to_tokens(d, self.tokenizer, max_tokens_per_doc) + if isinstance(d, str) + else d + for d in data_2 + ] + return ScoringData(data_1=data_1, data_2=data_2) + def valid_inputs( self, data_1: ScoreInput | list[ScoreInput], @@ -82,6 +142,15 @@ class BiEncoderIOProcessor(ScoringIOProcessor): raise ValueError(f"Invalid {self.name} request type") scoring_data = self.valid_inputs(data_1, data_2) + + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + request=request + ) + if max_tokens_per_query > 0 or max_tokens_per_doc > 0: + scoring_data = self._truncate_scoring_data( + scoring_data, max_tokens_per_query, max_tokens_per_doc + ) + tok_params = request.build_tok_params(self.model_config) engine_inputs = self._pre_process( scoring_data, @@ -112,10 +181,23 @@ class BiEncoderIOProcessor(ScoringIOProcessor): def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: assert isinstance(ctx.prompts, ScoringData) + assert not isinstance(ctx.pooling_params, Sequence) + tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) - return self._pre_process(ctx.prompts, tok_params) + + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + pooling_params=ctx.pooling_params + ) + + scoring_data = ctx.prompts + if max_tokens_per_query > 0 or max_tokens_per_doc > 0: + scoring_data = self._truncate_scoring_data( + scoring_data, max_tokens_per_query, max_tokens_per_doc + ) + + return self._pre_process(scoring_data, tok_params) def post_process_offline( self, @@ -138,7 +220,7 @@ class BiEncoderIOProcessor(ScoringIOProcessor): scoring_data.data_2, "document", self.model_config ) - return self._preprocess_completion_offline( + return self._preprocess_cmpl_offline( prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras ) @@ -217,8 +299,38 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor): class FlashLateInteractionIOProcessor(LateInteractionIOProcessor): name = "flash-late-interaction" - def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): - return outputs + def post_process_online( + self, + ctx: ScoringServeContext, + ): + assert ctx.query_final_res_batch is not None + assert ctx.final_res_batch is not None + assert isinstance(ctx.n_queries, int) + + # Expand queries if 1:N scoring + if len(ctx.query_final_res_batch) == 1: + ctx.query_final_res_batch = ctx.query_final_res_batch * len( + ctx.final_res_batch + ) + + final_res_batch: list[PoolingRequestOutput] = [] + for d1, d2 in zip(ctx.query_final_res_batch, ctx.final_res_batch): + padding: list[int] = [] + if (pad_token_id := self.pad_token_id) is not None: + padding = [pad_token_id] + + tokens = d1.prompt_token_ids + padding + d2.prompt_token_ids + + final_res_batch.append( + PoolingRequestOutput( + request_id=f"{d1.request_id}_{d2.request_id}", + outputs=d2.outputs, + prompt_token_ids=tokens, + num_cached_tokens=d1.num_cached_tokens + d2.num_cached_tokens, + finished=True, + ) + ) + ctx.final_res_batch = final_res_batch class CrossEncoderIOProcessor(ScoringIOProcessor): @@ -255,6 +367,11 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): raise ValueError(f"Invalid {self.name} request type") scoring_data = self.valid_inputs(data_1, data_2) + + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + request=request + ) + tok_params = request.build_tok_params(self.model_config) pooling_params = self.create_pooling_params(request) @@ -263,6 +380,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): tok_params, pooling_params, chat_template=self.chat_template, + max_tokens_per_query=max_tokens_per_query, + max_tokens_per_doc=max_tokens_per_doc, prompt_extras={ k: v for k in ("mm_processor_kwargs", "cache_salt") @@ -278,13 +397,23 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: assert isinstance(ctx.prompts, ScoringData) - assert not isinstance(ctx.pooling_params, list) + assert not isinstance(ctx.pooling_params, Sequence) tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) + + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( + pooling_params=ctx.pooling_params + ) + engine_inputs, pooling_params_list = self._pre_process( - ctx.prompts, tok_params, ctx.pooling_params, ctx.chat_template + ctx.prompts, + tok_params, + ctx.pooling_params, + ctx.chat_template, + max_tokens_per_query=max_tokens_per_query, + max_tokens_per_doc=max_tokens_per_doc, ) ctx.pooling_params = pooling_params_list return engine_inputs @@ -298,6 +427,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): tok_params: TokenizeParams, pooling_params: PoolingParams | None, chat_template: str | None = None, + max_tokens_per_query: int = 0, + max_tokens_per_doc: int = 0, prompt_extras: dict[str, Any] | None = None, ) -> tuple[Sequence[EngineInput], list[PoolingParams]]: # todo: support prompt_extras @@ -320,6 +451,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): data_2=d, encode_kwargs=tok_params.get_encode_kwargs(), chat_template=chat_template, + max_tokens_per_query=max_tokens_per_query, + max_tokens_per_doc=max_tokens_per_doc, ) if token_type_ids := engine_prompt.pop("token_type_ids", None): @@ -342,6 +475,8 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): data_2: ScoreData, encode_kwargs: dict[str, Any], chat_template: str | None = None, + max_tokens_per_query: int = 0, + max_tokens_per_doc: int = 0, ): model_config = self.model_config tokenizer = self.tokenizer @@ -352,25 +487,61 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): model_config, ) + # Apply truncation before defining closures + if max_tokens_per_query > 0 and isinstance(prompt_1, str): + prompt_1 = truncate_text_to_tokens( + prompt_1, tokenizer, max_tokens_per_query + ) + if max_tokens_per_doc > 0 and isinstance(prompt_2, str): + prompt_2 = truncate_text_to_tokens(prompt_2, tokenizer, max_tokens_per_doc) + def default_tokenizer_encode(): + local_kwargs = encode_kwargs.copy() + if self.supports_score_template: assert self.model is not None full_prompt = self.model.get_score_template(prompt_1, prompt_2) if full_prompt is None: raise ValueError("Get empty score template from model") - prompt_inputs = tokenizer(full_prompt, **encode_kwargs) + prompt_inputs = tokenizer(full_prompt, **local_kwargs) else: if self.use_sep_token: # cross_encoder models defaults to using separating token. + if max_tokens_per_doc > 0 and isinstance(prompt_2, str): + query_tokens = tokenizer.encode( + prompt_1, add_special_tokens=False + ) + num_special = get_num_special_tokens_for_pair(tokenizer) + doc_limit_max_length = ( + len(query_tokens) + max_tokens_per_doc + num_special + ) + existing_max_length = local_kwargs.get("max_length") + if existing_max_length is not None: + effective_max_length = min( + doc_limit_max_length, existing_max_length + ) + else: + effective_max_length = doc_limit_max_length + local_kwargs["truncation"] = "only_second" + local_kwargs["max_length"] = effective_max_length + prompt_inputs = tokenizer( - text=prompt_1, text_pair=prompt_2, **encode_kwargs + text=prompt_1, text_pair=prompt_2, **local_kwargs ) full_prompt = tokenizer.decode(prompt_inputs["input_ids"]) else: # `llm as reranker` defaults to not using separating token. - full_prompt = prompt_1 + prompt_2 - prompt_inputs = tokenizer(text=full_prompt, **encode_kwargs) + if max_tokens_per_doc > 0 and isinstance(prompt_2, str): + query_ids = tokenizer.encode(prompt_1, add_special_tokens=False) + doc_ids = tokenizer.encode(prompt_2, add_special_tokens=False) + doc_ids = doc_ids[:max_tokens_per_doc] + input_ids = query_ids + doc_ids + full_prompt = tokenizer.decode(input_ids) + prompt_inputs = {"input_ids": input_ids} + else: + full_prompt = prompt_1 + prompt_2 + prompt_inputs = tokenizer(text=full_prompt, **local_kwargs) return full_prompt, prompt_inputs # FIXME: For now, we only apply a template when one is explicitly provided. @@ -416,11 +587,137 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): return full_prompt, engine_prompt +class JinaRankingIOProcessorMixin: + @staticmethod + def sanitize_input(text: str, special_tokens: dict[str, str]) -> str: + for token in special_tokens.values(): + text = text.replace(token, "") + return text + + @staticmethod + def format_docs_prompts_func( + query: str, + docs: list[str], + special_tokens: dict[str, str] | None = None, + instruction: str | None = None, + no_thinking: bool = True, + ) -> str: + # TODO: Try converting the code below into a chat template. + + default_special_tokens = { + "query_embed_token": "<|rerank_token|>", + "doc_embed_token": "<|embed_token|>", + } + if special_tokens is None: + special_tokens = default_special_tokens + + query = JinaRankingIOProcessorMixin.sanitize_input(query, special_tokens) + docs = [ + JinaRankingIOProcessorMixin.sanitize_input(doc, special_tokens) + for doc in docs + ] + + prefix = ( + "<|im_start|>system\n" + "You are a search relevance expert who can determine a ranking of the passages based on how relevant they are to the query. " # noqa: E501 + "If the query is a question, how relevant a passage is depends on how well it answers the question. " # noqa: E501 + "If not, try to analyze the intent of the query and assess how well each passage satisfies the intent. " # noqa: E501 + "If an instruction is provided, you should follow the instruction when determining the ranking." # noqa: E501 + "<|im_end|>\n<|im_start|>user\n" + ) + suffix = "<|im_end|>\n<|im_start|>assistant\n" + if no_thinking: + suffix += "\n\n\n\n" + + doc_emb_token = special_tokens["doc_embed_token"] + query_emb_token = special_tokens["query_embed_token"] + + prompt = ( + f"I will provide you with {len(docs)} passages, each indicated by a numerical identifier. " # noqa: E501 + f"Rank the passages based on their relevance to query: {query}\n" + ) + + if instruction: + prompt += f"\n{instruction}\n\n" + + doc_prompts = [ + f'\n{doc}{doc_emb_token}\n' + for i, doc in enumerate(docs) + ] + prompt += "\n".join(doc_prompts) + "\n" + prompt += f"\n{query}{query_emb_token}\n" + + return prefix + prompt + suffix + + @staticmethod + def ensure_str(data: Sequence[Any]) -> list[str]: + text: list[str] = [] + for prompt in data: + if not isinstance(prompt, str): + raise ValueError( + "The JinaForRanking model only supports text as input." + ) + text.append(prompt) + return text + + +class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorMixin): + name = "jina-reranking-scoring" + pooling_task: PoolingTask = "token_embed" + + def _pre_process( + self, + scoring_data: ScoringData, + tok_params: TokenizeParams, + prompt_extras: dict[str, Any] | None = None, + ) -> Sequence[EngineInput]: + queries = self.ensure_str(scoring_data.data_1) + docs = self.ensure_str(scoring_data.data_2) + + if len(queries) == 1: + prompts = [self.format_docs_prompts_func(query=queries[0], docs=docs)] + else: + prompts = [ + self.format_docs_prompts_func(query=q, docs=[d]) + for q, d in zip(queries, docs) + ] + + return self._preprocess_cmpl_offline( + prompts=prompts, tok_params=tok_params, prompt_extras=prompt_extras + ) + + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): + final_res_batch: list[PoolingRequestOutput] = [] + + for i in range(len(outputs)): + embeds = outputs[i].outputs.data.float() + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + query_embeds = embeds[-1] + doc_embeds = embeds[:-1] + + scores = F.cosine_similarity(query_embeds, doc_embeds) + + for score in scores: + final_res_batch.append( + PoolingRequestOutput( + request_id=outputs[i].request_id, + outputs=score, + prompt_token_ids=outputs[i].prompt_token_ids, + num_cached_tokens=outputs[i].num_cached_tokens, + finished=True, + ) + ) + return final_res_batch + + ScoringIOProcessors: dict[str, type[ScoringIOProcessor]] = { p.name: p for p in [ BiEncoderIOProcessor, LateInteractionIOProcessor, + JinaRankingIOProcessor, FlashLateInteractionIOProcessor, CrossEncoderIOProcessor, ] diff --git a/vllm/entrypoints/pooling/scoring/protocol.py b/vllm/entrypoints/pooling/scoring/protocol.py index 9fbfbed3732..83fdafb1458 100644 --- a/vllm/entrypoints/pooling/scoring/protocol.py +++ b/vllm/entrypoints/pooling/scoring/protocol.py @@ -20,6 +20,24 @@ from .typing import ScoreContentPartParam, ScoreInput class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin): + max_tokens_per_query: int = Field( + default=0, + description=( + "Maximum number of tokens per query. Queries longer than " + "this will be truncated to this length. 0 means no " + "query-level truncation is applied." + ), + ) + max_tokens_per_doc: int = Field( + default=0, + description=( + "Maximum number of tokens per document. Documents longer than " + "this will be truncated to this length. 0 means no " + "document-level truncation is applied (only truncate_prompt_tokens " + "applies to the combined query+document)." + ), + ) + def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: encoder_config = model_config.encoder_config or {} @@ -91,29 +109,11 @@ ScoreRequest: TypeAlias = ( ) -class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin): +class RerankRequest(ScoreRequestMixin): query: ScoreInput documents: ScoreInput | list[ScoreInput] top_n: int = Field(default_factory=lambda: 0) - def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams: - encoder_config = model_config.encoder_config or {} - - return TokenizeParams( - max_total_tokens=model_config.max_model_len, - max_output_tokens=0, - truncate_prompt_tokens=self.truncate_prompt_tokens, - truncation_side=self.truncation_side, - do_lower_case=encoder_config.get("do_lower_case", False), - max_total_tokens_param="max_model_len", - ) - - def to_pooling_params(self, task: PoolingTask = "classify"): - return PoolingParams( - task=task, - use_activation=self.use_activation, - ) - ScoringRequest: TypeAlias = ScoreRequest | RerankRequest diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index de5b5797ce4..df866efd56e 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -4,15 +4,12 @@ from fastapi.responses import JSONResponse, Response from vllm import PoolingParams -from vllm.config import ModelConfig from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import ChatTemplateConfig from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor from vllm.entrypoints.pooling.base.serving import PoolingServing from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput -from vllm.renderers import BaseRenderer from vllm.v1.pool.late_interaction import ( build_late_interaction_doc_params, build_late_interaction_query_params, @@ -44,30 +41,23 @@ class ServingScores(PoolingServing): enable_flash_late_interaction: bool = True, **kwargs, ): - self.score_type = engine_client.model_config.score_type + self.io_processor_name: str = engine_client.model_config.score_type self.enable_flash_late_interaction = ( - self.score_type == "late-interaction" and enable_flash_late_interaction + self.io_processor_name == "late-interaction" + and enable_flash_late_interaction ) + if self.enable_flash_late_interaction: + self.io_processor_name = "flash-late-interaction" + + if engine_client.model_config.architecture == "JinaForRanking": + self.io_processor_name = "jina-reranking-scoring" + self.enable_flash_late_interaction = False + super().__init__(engine_client, *args, **kwargs) - def init_io_processor( - self, - model_config: ModelConfig, - renderer: BaseRenderer, - chat_template_config: ChatTemplateConfig, - ) -> PoolingIOProcessor: - score_type: str = model_config.score_type - if self.enable_flash_late_interaction: - score_type = "flash-late-interaction" - - assert score_type in ScoringIOProcessors - processor_cls = ScoringIOProcessors[score_type] - return processor_cls( - model_config=model_config, - renderer=renderer, - chat_template_config=chat_template_config, - ) + def init_io_processor(self, *args, **kwargs) -> PoolingIOProcessor: + return ScoringIOProcessors[self.io_processor_name](*args, **kwargs) async def __call__(self, *args, **kwargs) -> Response: if not self.enable_flash_late_interaction: @@ -75,7 +65,7 @@ class ServingScores(PoolingServing): return await self.flash_late_interaction(*args, **kwargs) - async def _build_response( + def _build_response( self, ctx: ScoringServeContext, ) -> JSONResponse: @@ -193,17 +183,15 @@ class ServingScores(PoolingServing): ### Can significantly improve late-interaction scoring performance. async def flash_late_interaction(self, *args, **kwargs) -> Response: - ctx = await self._init_ctx(*args, **kwargs) - ctx.pooling_params = self.io_processor.create_pooling_params(ctx.request) - await self.io_processor.pre_process_online_async(ctx) + ctx = await self._init_ctx(self.io_processor, *args, **kwargs) + await self._preprocessing_async(self.io_processor, ctx) # stage 1: encode queries and cache token embeddings on workers. await self._flash_late_interaction_encode_queries(ctx) # stage 2: encode docs and return scalar scores from workers. await self._flash_late_interaction_encode_docs(ctx) - await self.io_processor.post_process_online_async(ctx) - return await self._build_response(ctx) + return await self._postprocessing_async(self.io_processor, ctx) async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext): assert ctx.n_queries is not None @@ -247,6 +235,7 @@ class ServingScores(PoolingServing): await self._prepare_generators(query_ctx) await self._collect_batch(query_ctx) + ctx.query_final_res_batch = query_ctx.final_res_batch async def _flash_late_interaction_encode_docs(self, ctx: ScoringServeContext): assert ctx.n_queries is not None diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index 01b8514eb7b..13db389b823 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -25,6 +25,37 @@ from .typing import ( ) +def get_num_special_tokens_for_pair(tokenizer) -> int: + """Get number of special tokens added for a text pair encoding.""" + method = getattr(tokenizer, "num_special_tokens_to_add", None) + if method is not None: + try: + return method(pair=True) + except TypeError: + pass + # Fallback: compute by tokenizing empty strings + empty_encoding = tokenizer("", text_pair="", add_special_tokens=True) + return len(empty_encoding["input_ids"]) + + +def truncate_text_to_tokens( + text: str, + tokenizer, + max_tokens: int, +) -> str: + """Truncate text to a maximum number of content tokens. + + Uses offset_mapping to slice the original text at the exact character + boundary, avoiding lossy encode→decode round-trips that can shift + the token count by 1-3 tokens due to BPE merge boundary changes. + """ + encoding = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) + if len(encoding["input_ids"]) <= max_tokens: + return text + char_end = encoding["offset_mapping"][max_tokens - 1][1] + return text[:char_end] + + def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor: """ Compute ColBERT MaxSim score. diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 66dd9dd4d2b..0ce9d5c8384 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -30,7 +30,7 @@ from vllm.entrypoints.pooling.pooling.protocol import ( ) from vllm.entrypoints.pooling.scoring.protocol import ScoringRequest, ScoringResponse from vllm.entrypoints.pooling.scoring.typing import ScoringData -from vllm.inputs import EngineInput +from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest PoolingCompletionLikeRequest: TypeAlias = ( @@ -69,9 +69,9 @@ class PoolingServeContext(Generic[PoolingRequestT]): raw_request: Request | None = None model_name: str request_id: str + pooling_params: PoolingParams | list[PoolingParams] created_time: int = field(default_factory=lambda: int(time.time())) lora_request: LoRARequest | None = None - pooling_params: PoolingParams | list[PoolingParams] | None = None engine_inputs: Sequence[EngineInput] | None = None prompt_request_ids: list[str] | None = None intermediates: Any | None = None @@ -86,11 +86,17 @@ class PoolingServeContext(Generic[PoolingRequestT]): ## for bi-encoder & late-interaction n_queries: int | None = None + ## for IOProcessorResponse + response: Any | None = None + + ## for flash-late-interaction + query_final_res_batch: list[PoolingRequestOutput] | None = None + @dataclass class OfflineInputsContext: - prompts: PromptType | Sequence[PromptType] | ScoringData - pooling_params: PoolingParams | list[PoolingParams] | None = None + prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData + pooling_params: PoolingParams | Sequence[PoolingParams] tokenization_kwargs: dict[str, Any] | None = None chat_template: str | None = None diff --git a/vllm/entrypoints/sagemaker/api_router.py b/vllm/entrypoints/sagemaker/api_router.py index 45f5613bf5e..1d63793df39 100644 --- a/vllm/entrypoints/sagemaker/api_router.py +++ b/vllm/entrypoints/sagemaker/api_router.py @@ -14,7 +14,7 @@ from vllm.config import ModelConfig from vllm.entrypoints.openai.engine.protocol import ErrorResponse from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.entrypoints.openai.utils import validate_json_request -from vllm.entrypoints.pooling.base.serving import PoolingServing +from vllm.entrypoints.pooling.base.serving import PoolingServingBase from vllm.entrypoints.pooling.utils import enable_scoring_api from vllm.entrypoints.serve.instrumentator.basic import base from vllm.entrypoints.serve.instrumentator.health import health @@ -23,7 +23,7 @@ from vllm.tasks import POOLING_TASKS, SupportedTask # TODO: RequestType = TypeForm[BaseModel] when recognized by type checkers # (requires typing_extensions >= 4.13) RequestType = Any -GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServing | None] +GetHandlerFn = Callable[[Request], OpenAIServing | PoolingServingBase | None] EndpointFn = Callable[[RequestType, Request], Awaitable[Any]] diff --git a/vllm/entrypoints/serve/render/serving.py b/vllm/entrypoints/serve/render/serving.py index 2aaa83e7564..5aa2449797b 100644 --- a/vllm/entrypoints/serve/render/serving.py +++ b/vllm/entrypoints/serve/render/serving.py @@ -65,7 +65,6 @@ class OpenAIServingRender: self, model_config: ModelConfig, renderer: BaseRenderer, - io_processor: Any, model_registry: OpenAIModelRegistry, *, request_logger: RequestLogger | None, @@ -81,7 +80,6 @@ class OpenAIServingRender: ) -> None: self.model_config = model_config self.renderer = renderer - self.io_processor = io_processor self.model_registry = model_registry self.request_logger = request_logger self.chat_template = chat_template @@ -550,7 +548,9 @@ class OpenAIServingRender: if reasoning_parser is not None: tokenizer = renderer.get_tokenizer() - request = reasoning_parser(tokenizer).adjust_request(request=request) + request = reasoning_parser( + tokenizer, model_config=self.model_config + ).adjust_request(request=request) # tool parsing is done only if a tool_parser has been set and if # tool_choice is not "none" (if tool_choice is "none" but a tool_parser diff --git a/vllm/env_override.py b/vllm/env_override.py index 55dd5099dac..2e3f02866ab 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -586,3 +586,52 @@ if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0"): return runtime_env GraphCaptureOutput.get_runtime_env = _patched_get_runtime_env + +# =================================================== +# torch 2.10 FxGraphCachePickler.dumps ValueError fix +# =================================================== +# PyTorch 2.10's FxGraphCachePickler.dumps() doesn't catch ValueError, +# causing torch.compile cache failures when tensors with non-standard +# layouts (e.g. blocked-layout prepacked weights) are serialized. +# PyTorch mainline fixed this in pytorch/pytorch#176557 (merged 2026-03-04). +# This is a thin backport for 2.10 users; remove once 2.10 is dropped. + + +def _apply_fxgraphcache_pickle_patch(pickler_cls, bypass_cls): + """Wrap pickler_cls.dumps to convert ValueError into bypass_cls. + + Idempotent: sets `_vllm_fxgraph_dumps_patched` on the class after the + first apply to prevent re-application. The wrapper function is also + marked with `_vllm_patched` as an additional safeguard. + """ + if getattr(pickler_cls, "_vllm_fxgraph_dumps_patched", False): + return + + original_dumps = pickler_cls.dumps + if hasattr(original_dumps, "_vllm_patched"): + return + + def patched_dumps(self, obj): + try: + return original_dumps(self, obj) + except ValueError as e: + raise bypass_cls("Failed to pickle cache key") from e + + patched_dumps._vllm_patched = True # type: ignore[attr-defined] + pickler_cls.dumps = patched_dumps + pickler_cls._vllm_fxgraph_dumps_patched = True # type: ignore[attr-defined] + + +def _patch_fxgraphcache_pickle_if_needed(): + """Apply FxGraphCachePickler.dumps ValueError backport when on torch 2.10.x.""" + from vllm.utils.torch_utils import is_torch_equal_or_newer + + if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"): + return + + from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler + + _apply_fxgraphcache_pickle_patch(FxGraphCachePickler, BypassFxGraphCache) + + +_patch_fxgraphcache_pickle_if_needed() diff --git a/vllm/envs.py b/vllm/envs.py index d2af9e64d66..8ed1d33434c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -129,6 +129,7 @@ if TYPE_CHECKING: VLLM_ENABLE_V1_MULTIPROCESSING: bool = True VLLM_LOG_BATCHSIZE_INTERVAL: float = -1 VLLM_DISABLE_COMPILE_CACHE: bool = False + VLLM_USE_LAYERNAME: bool = True Q_SCALE_CONSTANT: int = 200 K_SCALE_CONSTANT: int = 200 V_SCALE_CONSTANT: int = 100 @@ -145,8 +146,6 @@ if TYPE_CHECKING: VLLM_ENABLE_PREGRAD_PASSES: bool = False VLLM_DP_MASTER_IP: str = "" VLLM_DP_MASTER_PORT: int = 0 - VLLM_MOE_DP_CHUNK_SIZE: int = 256 - VLLM_ENABLE_MOE_DP_CHUNK: bool = True VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False VLLM_RAY_DP_PACK_STRATEGY: Literal["strict", "fill", "span"] = "strict" VLLM_RAY_EXTRA_ENV_VAR_PREFIXES_TO_COPY: str = "" @@ -257,6 +256,7 @@ if TYPE_CHECKING: VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = False VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False + VLLM_LORA_ENABLE_DUAL_STREAM: bool = False def get_default_cache_root(): @@ -298,10 +298,7 @@ def use_aot_compile() -> bool: else "0" ) - return ( - not bool(int(os.getenv("VLLM_BATCH_INVARIANT", "0"))) - and os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1" - ) + return os.environ.get("VLLM_USE_AOT_COMPILE", default_value) == "1" def use_mega_aot_artifact(): @@ -495,8 +492,9 @@ environment_variables: dict[str, Callable[[], Any]] = { # rocm, cpu] "VLLM_TARGET_DEVICE": lambda: os.getenv("VLLM_TARGET_DEVICE", "cuda").lower(), # Main CUDA version of vLLM. This follows PyTorch but can be overridden. - "VLLM_MAIN_CUDA_VERSION": lambda: os.getenv("VLLM_MAIN_CUDA_VERSION", "").lower() - or "12.9", + "VLLM_MAIN_CUDA_VERSION": lambda: ( + os.getenv("VLLM_MAIN_CUDA_VERSION", "").lower() or "12.9" + ), # Controls PyTorch float32 matmul precision mode within vLLM workers. # Valid options mirror torch.set_float32_matmul_precision "VLLM_FLOAT32_MATMUL_PRECISION": env_with_choices( @@ -516,21 +514,19 @@ environment_variables: dict[str, Callable[[], Any]] = { # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU. "NVCC_THREADS": lambda: os.getenv("NVCC_THREADS", None), # If set, vllm will use precompiled binaries (*.so) - "VLLM_USE_PRECOMPILED": lambda: os.environ.get("VLLM_USE_PRECOMPILED", "") - .strip() - .lower() - in ("1", "true") - or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")), + "VLLM_USE_PRECOMPILED": lambda: ( + os.environ.get("VLLM_USE_PRECOMPILED", "").strip().lower() in ("1", "true") + or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")) + ), # If set, skip adding +precompiled suffix to version string "VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX": lambda: bool( int(os.environ.get("VLLM_SKIP_PRECOMPILED_VERSION_SUFFIX", "0")) ), # Used to mark that setup.py is running in a Docker build context, # in order to force the use of precompiled binaries. - "VLLM_DOCKER_BUILD_CONTEXT": lambda: os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "") - .strip() - .lower() - in ("1", "true"), + "VLLM_DOCKER_BUILD_CONTEXT": lambda: ( + os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "").strip().lower() in ("1", "true") + ), # CMake build type # If not set, defaults to "Debug" or "RelWithDebInfo" # Available options: "Debug", "Release", "RelWithDebInfo" @@ -576,10 +572,9 @@ environment_variables: dict[str, Callable[[], Any]] = { ), # If true, will load models from ModelScope instead of Hugging Face Hub. # note that the value is true or false, not numbers - "VLLM_USE_MODELSCOPE": lambda: os.environ.get( - "VLLM_USE_MODELSCOPE", "False" - ).lower() - == "true", + "VLLM_USE_MODELSCOPE": lambda: ( + os.environ.get("VLLM_USE_MODELSCOPE", "False").lower() == "true" + ), # Interval in seconds to log a warning message when the ring buffer is full "VLLM_RINGBUFFER_WARNING_INTERVAL": lambda: int( os.environ.get("VLLM_RINGBUFFER_WARNING_INTERVAL", "60") @@ -600,19 +595,17 @@ environment_variables: dict[str, Callable[[], Any]] = { # Feature flag to enable/disable Inductor standalone compile. # In torch <= 2.7 we ignore this flag; in torch >= 2.9 this is # enabled by default. - "VLLM_USE_STANDALONE_COMPILE": lambda: os.environ.get( - "VLLM_USE_STANDALONE_COMPILE", "1" - ) - == "1", + "VLLM_USE_STANDALONE_COMPILE": lambda: ( + os.environ.get("VLLM_USE_STANDALONE_COMPILE", "1") == "1" + ), # Inductor's pre-grad passes don't do anything for vLLM. # The pre-grad passes get run even on cache-hit and negatively impact # vllm cold compile times by O(1s) # Can remove this after the following issue gets fixed # https://github.com/pytorch/pytorch/issues/174502 - "VLLM_ENABLE_PREGRAD_PASSES": lambda: os.environ.get( - "VLLM_ENABLE_PREGRAD_PASSES", "0" - ) - == "1", + "VLLM_ENABLE_PREGRAD_PASSES": lambda: ( + os.environ.get("VLLM_ENABLE_PREGRAD_PASSES", "0") == "1" + ), # Debug pattern matching inside custom passes. # Should be set to the fx.Node name (e.g. 'getitem_34' or 'scaled_mm_3'). "VLLM_PATTERN_MATCH_DEBUG": lambda: os.environ.get( @@ -655,10 +648,9 @@ environment_variables: dict[str, Callable[[], Any]] = { # API key for vLLM API server "VLLM_API_KEY": lambda: os.environ.get("VLLM_API_KEY", None), # Whether to log responses from API Server for debugging - "VLLM_DEBUG_LOG_API_SERVER_RESPONSE": lambda: os.environ.get( - "VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False" - ).lower() - == "true", + "VLLM_DEBUG_LOG_API_SERVER_RESPONSE": lambda: ( + os.environ.get("VLLM_DEBUG_LOG_API_SERVER_RESPONSE", "False").lower() == "true" + ), # S3 access information, used for tensorizer to load model from S3 "S3_ACCESS_KEY_ID": lambda: os.environ.get("S3_ACCESS_KEY_ID", None), "S3_SECRET_ACCESS_KEY": lambda: os.environ.get("S3_SECRET_ACCESS_KEY", None), @@ -669,11 +661,13 @@ environment_variables: dict[str, Callable[[], Any]] = { ), "VLLM_NO_USAGE_STATS": lambda: os.environ.get("VLLM_NO_USAGE_STATS", "0") == "1", "VLLM_DO_NOT_TRACK": lambda: ( - os.environ.get("VLLM_DO_NOT_TRACK", None) - or os.environ.get("DO_NOT_TRACK", None) - or "0" - ) - == "1", + ( + os.environ.get("VLLM_DO_NOT_TRACK", None) + or os.environ.get("DO_NOT_TRACK", None) + or "0" + ) + == "1" + ), "VLLM_USAGE_SOURCE": lambda: os.environ.get("VLLM_USAGE_SOURCE", "production"), # Logging configuration # If set to 0, vllm will not configure logging @@ -696,36 +690,40 @@ environment_variables: dict[str, Callable[[], Any]] = { "NO_COLOR": lambda: os.getenv("NO_COLOR", "0") != "0", # If set, vllm will log stats at this interval in seconds # If not set, vllm will log stats every 10 seconds. - "VLLM_LOG_STATS_INTERVAL": lambda: val - if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0 - else 10.0, + "VLLM_LOG_STATS_INTERVAL": lambda: ( + val + if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0 + else 10.0 + ), # Trace function calls # If set to 1, vllm will trace function calls # Useful for debugging "VLLM_TRACE_FUNCTION": lambda: int(os.getenv("VLLM_TRACE_FUNCTION", "0")), # If set, vllm will use flashinfer sampler - "VLLM_USE_FLASHINFER_SAMPLER": lambda: bool( - int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"]) - ) - if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ - else None, + "VLLM_USE_FLASHINFER_SAMPLER": lambda: ( + bool(int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"])) + if "VLLM_USE_FLASHINFER_SAMPLER" in os.environ + else None + ), # Pipeline stage partition strategy "VLLM_PP_LAYER_PARTITION": lambda: os.getenv("VLLM_PP_LAYER_PARTITION", None), # (CPU backend only) CPU key-value cache space. # default is None and will be set as 4 GB - "VLLM_CPU_KVCACHE_SPACE": lambda: int(os.getenv("VLLM_CPU_KVCACHE_SPACE", "0")) - if "VLLM_CPU_KVCACHE_SPACE" in os.environ - else None, + "VLLM_CPU_KVCACHE_SPACE": lambda: ( + int(os.getenv("VLLM_CPU_KVCACHE_SPACE", "0")) + if "VLLM_CPU_KVCACHE_SPACE" in os.environ + else None + ), # (CPU backend only) CPU core ids bound by OpenMP threads, e.g., "0-31", # "0,1,2", "0-31,33". CPU cores of different ranks are separated by '|'. "VLLM_CPU_OMP_THREADS_BIND": lambda: os.getenv("VLLM_CPU_OMP_THREADS_BIND", "auto"), # (CPU backend only) CPU cores not used by OMP threads . # Those CPU cores will not be used by OMP threads of a rank. - "VLLM_CPU_NUM_OF_RESERVED_CPU": lambda: int( - os.getenv("VLLM_CPU_NUM_OF_RESERVED_CPU", "0") - ) - if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ - else None, + "VLLM_CPU_NUM_OF_RESERVED_CPU": lambda: ( + int(os.getenv("VLLM_CPU_NUM_OF_RESERVED_CPU", "0")) + if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ + else None + ), # (CPU backend only) whether to use SGL kernels, optimized for small batch. "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))), # (CPU backend only) whether to enable attention spilt KV. @@ -919,9 +917,11 @@ environment_variables: dict[str, Callable[[], Any]] = { # a list of plugin names to load, separated by commas. # if this is not set, it means all plugins will be loaded # if this is set to an empty string, no plugins will be loaded - "VLLM_PLUGINS": lambda: None - if "VLLM_PLUGINS" not in os.environ - else os.environ["VLLM_PLUGINS"].split(","), + "VLLM_PLUGINS": lambda: ( + None + if "VLLM_PLUGINS" not in os.environ + else os.environ["VLLM_PLUGINS"].split(",") + ), # a local directory to look in for unrecognized LoRA adapters. # only works if plugins are enabled and # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled. @@ -953,9 +953,11 @@ environment_variables: dict[str, Callable[[], Any]] = { # and performance comparisons. Currently only affects MPLinearKernel # selection # (kernels: MacheteLinearKernel, MarlinLinearKernel, ExllamaLinearKernel) - "VLLM_DISABLED_KERNELS": lambda: [] - if "VLLM_DISABLED_KERNELS" not in os.environ - else os.environ["VLLM_DISABLED_KERNELS"].split(","), + "VLLM_DISABLED_KERNELS": lambda: ( + [] + if "VLLM_DISABLED_KERNELS" not in os.environ + else os.environ["VLLM_DISABLED_KERNELS"].split(",") + ), "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool( int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1")) ), @@ -1090,6 +1092,9 @@ environment_variables: dict[str, Callable[[], Any]] = { os.getenv("VLLM_LOG_BATCHSIZE_INTERVAL", "-1") ), "VLLM_DISABLE_COMPILE_CACHE": disable_compile_cache, + # If set to "0", disable LayerName opaque type for layer_name + # parameters in custom ops. Defaults to enabled on torch >= 2.11. + "VLLM_USE_LAYERNAME": lambda: bool(int(os.getenv("VLLM_USE_LAYERNAME", "1"))), # If set, vllm will run in development mode, which will enable # some additional endpoints for developing and debugging, # e.g. `/reset_prefix_cache` @@ -1133,20 +1138,10 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_DP_MASTER_IP": lambda: os.getenv("VLLM_DP_MASTER_IP", "127.0.0.1"), # Port of the master node in the data parallel setting "VLLM_DP_MASTER_PORT": lambda: int(os.getenv("VLLM_DP_MASTER_PORT", "0")), - # In the context of executing MoE models with Data-Parallel, Expert-Parallel - # and Batched All-to-All dispatch/combine kernels, VLLM_MOE_DP_CHUNK_SIZE - # dictates the quantum of tokens that can be dispatched from a DP - # rank. All DP ranks process the activations in VLLM_MOE_DP_CHUNK_SIZE - # units. - "VLLM_MOE_DP_CHUNK_SIZE": lambda: int(os.getenv("VLLM_MOE_DP_CHUNK_SIZE", "256")), - "VLLM_ENABLE_MOE_DP_CHUNK": lambda: bool( - int(os.getenv("VLLM_ENABLE_MOE_DP_CHUNK", "1")) - ), # Randomize inputs during dummy runs when using Data Parallel - "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: os.environ.get( - "VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0" - ) - == "1", + "VLLM_RANDOMIZE_DP_DUMMY_INPUTS": lambda: ( + os.environ.get("VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1" + ), # Strategy to pack the data parallel ranks for Ray. # Available options: # - "fill": @@ -1186,10 +1181,9 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MODEL_REDIRECT_PATH", None ), # Whether to use atomicAdd reduce in gptq/awq marlin kernel. - "VLLM_MARLIN_USE_ATOMIC_ADD": lambda: os.environ.get( - "VLLM_MARLIN_USE_ATOMIC_ADD", "0" - ) - == "1", + "VLLM_MARLIN_USE_ATOMIC_ADD": lambda: ( + os.environ.get("VLLM_MARLIN_USE_ATOMIC_ADD", "0") == "1" + ), # Whether to use marlin kernel in mxfp4 quantization method "VLLM_MXFP4_USE_MARLIN": lambda: maybe_convert_bool( os.environ.get("VLLM_MXFP4_USE_MARLIN", None) @@ -1207,17 +1201,16 @@ environment_variables: dict[str, Callable[[], Any]] = { # Whether to turn on the outlines cache for V1 # This cache is unbounded and on disk, so it's not safe to use in # an environment with potentially malicious users. - "VLLM_V1_USE_OUTLINES_CACHE": lambda: os.environ.get( - "VLLM_V1_USE_OUTLINES_CACHE", "0" - ) - == "1", + "VLLM_V1_USE_OUTLINES_CACHE": lambda: ( + os.environ.get("VLLM_V1_USE_OUTLINES_CACHE", "0") == "1" + ), # Gap between padding buckets for the forward pass. So we have # 8, we will run forward pass with [16, 24, 32, ...]. - "VLLM_TPU_BUCKET_PADDING_GAP": lambda: int( - os.environ["VLLM_TPU_BUCKET_PADDING_GAP"] - ) - if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ - else 0, + "VLLM_TPU_BUCKET_PADDING_GAP": lambda: ( + int(os.environ["VLLM_TPU_BUCKET_PADDING_GAP"]) + if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ + else 0 + ), "VLLM_TPU_MOST_MODEL_LEN": lambda: maybe_convert_int( os.environ.get("VLLM_TPU_MOST_MODEL_LEN", None) ), @@ -1710,6 +1703,10 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool( int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0")) ), + # Whether to enable dual cuda streams for LoRA computation + "VLLM_LORA_ENABLE_DUAL_STREAM": lambda: bool( + int(os.getenv("VLLM_LORA_ENABLE_DUAL_STREAM", "0")) + ), } diff --git a/vllm/forward_context.py b/vllm/forward_context.py index fa568c33f36..537a28a4252 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -70,27 +70,8 @@ def _compute_sp_num_tokens( return sp_tokens.tolist() -def _compute_chunked_local_num_tokens( - num_tokens_across_dp_cpu: torch.Tensor, - sequence_parallel_size: int, - max_num_tokens: int, - chunk_idx: int, -) -> list[int]: - sp_tokens = _compute_sp_num_tokens(num_tokens_across_dp_cpu, sequence_parallel_size) - sp_size = len(sp_tokens) - - local_size = [-1] * sp_size - for i in range(sp_size): - # Take into account sharding if MoE activation is sequence parallel. - local_size[i] = min(max_num_tokens, sp_tokens[i] - (max_num_tokens * chunk_idx)) - if local_size[i] <= 0: - local_size[i] = 1 # ensure lockstep even if done - return local_size - - @dataclass class DPMetadata: - max_tokens_across_dp_cpu: torch.Tensor num_tokens_across_dp_cpu: torch.Tensor # NOTE: local_sizes should only be set by the chunked_sizes context manager @@ -113,47 +94,7 @@ class DPMetadata: assert num_tokens_across_dp_cpu[dp_rank] == batchsize, ( f"{num_tokens_across_dp_cpu[dp_rank]} {batchsize}" ) - max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp_cpu) - return DPMetadata(max_tokens_across_dp_cpu, num_tokens_across_dp_cpu) - - @contextmanager - def chunked_sizes( - self, sequence_parallel_size: int, max_chunk_size_per_rank: int, chunk_idx: int - ): - """ - Context manager to compute and temporarily set the per-rank local token - sizes for a specific chunk during chunked forward execution. - - This is necessary to ensure each DP (data parallel) rank processes its - designated portion of tokens in lockstep with others, even when the - token counts are uneven or some ranks have completed their input early. - - For chunked execution, we break up the total tokens on each rank into - multiple chunks (of at most `max_chunk_size_per_rank`), and for a given - `chunk_idx`, this context manager sets `self.local_sizes` to the number - of tokens to process in that chunk on each rank. - - `self.local_sizes` is only valid inside the context. - - Args: - sequence_parallel_size: When Attn is TP and MoE layers are EP, - we use SP between the layers to avoid - redundant ops. We need this value to - compute the chunked sizes. - max_chunk_size_per_rank: The max number of tokens each rank is - allowed to process in this chunk. - chunk_idx: The index of the chunk to compute sizes for. - """ - self.local_sizes = _compute_chunked_local_num_tokens( - self.num_tokens_across_dp_cpu, - sequence_parallel_size, - max_chunk_size_per_rank, - chunk_idx, - ) - try: - yield self.local_sizes - finally: - self.local_sizes = None + return DPMetadata(num_tokens_across_dp_cpu) @contextmanager def sp_local_sizes(self, sequence_parallel_size: int): diff --git a/vllm/lora/layers/base_linear.py b/vllm/lora/layers/base_linear.py index 1b666dcb790..4ea6b1ec8f0 100644 --- a/vllm/lora/layers/base_linear.py +++ b/vllm/lora/layers/base_linear.py @@ -5,8 +5,15 @@ import torch from transformers import PretrainedConfig +from vllm import envs +from vllm.config import get_current_vllm_config from vllm.config.lora import LoRAConfig from vllm.distributed.utils import divide +from vllm.forward_context import ( + ForwardContext, + get_forward_context, + is_forward_context_available, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, LinearBase, @@ -14,24 +21,88 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.platforms import current_platform +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import direct_register_custom_op from .base import BaseLayerWithLoRA from .utils import _get_lora_device +if envs.VLLM_LORA_ENABLE_DUAL_STREAM: + _lora_aux_cuda_stream: torch.cuda.Stream | None = None + + def _get_lora_aux_cuda_stream() -> torch.cuda.Stream | None: + global _lora_aux_cuda_stream + if _lora_aux_cuda_stream is None and current_platform.is_cuda_alike(): + _lora_aux_cuda_stream = torch.cuda.Stream() + return _lora_aux_cuda_stream + + def lora_linear_async( + layer_name: str, + output_size: int, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + forward_context: ForwardContext = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + return self._apply_async_impl(x, bias) + + def lora_linear_async_fake( + layer_name: str, + output_size: int, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + # The real function reshapes output back to the original 3D shape + # when the input has an extra batch dimension (transformers backend). + if x.ndim == 3: + return torch.empty( + (x.size(0), x.size(1), output_size), + device=x.device, + dtype=x.dtype, + ) + return torch.empty( + (x.size(0), output_size), + device=x.device, + dtype=x.dtype, + ) + + direct_register_custom_op( + op_name="lora_linear_async", + op_func=lora_linear_async, + fake_impl=lora_linear_async_fake, + ) + class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): def __init__(self, base_layer: LinearBase): super().__init__() + + self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM self.base_layer = base_layer self.input_size = self.base_layer.input_size # Ensure tp_size and tp_rank consistency with the base_layer. self.tp_size = self.base_layer.tp_size self.tp_rank = self.base_layer.tp_rank self.device = _get_lora_device(self.base_layer) + self._init_lora_stream_context() self.output_slices: tuple[int, ...] self.output_size: int self.n_slices: int + def _init_lora_stream_context(self) -> None: + if not self._enable_aux_cuda_stream: + return + vllm_config = get_current_vllm_config() + self._lora_stream = _get_lora_aux_cuda_stream() + assert current_platform.is_cuda_alike() + self._events = [torch.cuda.Event(), torch.cuda.Event()] + # lora_linear avoids prefix conflicts with the base layer + self.layer_name = self.base_layer.prefix + ".lora_linear_async" + compilation_config = vllm_config.compilation_config + if self.layer_name in compilation_config.static_forward_context: + raise ValueError("Duplicate layer name: {}".format(self.layer_name)) + compilation_config.static_forward_context[self.layer_name] = self + def create_lora_weights( self, max_loras: int, @@ -39,7 +110,6 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): model_config: PretrainedConfig | None = None, ) -> None: self.lora_config = lora_config - # if isinstance(self.base_layer, ReplicatedLinear): lora_a_out_size = lora_config.max_lora_rank lora_b_out_size = self.output_size @@ -120,6 +190,18 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): ) def apply(self, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: + # is_forward_context_available for tower modules + if self._enable_aux_cuda_stream and is_forward_context_available(): + output_size = sum(self.output_slices) + return torch.ops.vllm.lora_linear_async( + self.layer_name, output_size, x, bias + ) + else: + return self._apply_sync(x, bias) + + def _apply_sync( + self, x: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: output = self.base_layer.quant_method.apply(self.base_layer, x, bias) original_shape = output.shape if output.ndim == 3 else None @@ -144,6 +226,72 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA): return output + def _apply_async_impl( + self, x: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: + """ + Forward pass with base linear and LoRA on separate CUDA streams + for overlap, using maybe_execute_in_parallel. + Base layer runs on default stream; LoRA runs on aux stream. + """ + assert envs.VLLM_LORA_ENABLE_DUAL_STREAM + assert x.ndim in (2, 3) + num_tokens = x.size(0) if x.ndim == 2 else x.size(1) + output_size = sum(self.output_slices) + + def base_fn() -> torch.Tensor: + return self.base_layer.quant_method.apply(self.base_layer, x, bias) + + def lora_fn() -> torch.Tensor: + # Must be zeros, not empty: _lora_expand_kernel exits early (without + # writing) when lora_id == -1 (no active LoRA). If uninitialized, + # output.add_(lora_result) below would corrupt the base output. + lora_output = torch.zeros( + (num_tokens, output_size), + device=self.device, + dtype=x.dtype, + ) + + # Flatten the batch dimension for the transformers backend + # (which uses shape (1, seq_len, hidden)), matching _apply_sync. + x_2d = x.flatten(0, 1) if x.ndim == 3 else x + self.punica_wrapper.add_lora_linear( + lora_output, + x_2d, + self.lora_a_stacked, + self.lora_b_stacked, + 1.0, + self.output_slices, + add_inputs=False, + ) + return lora_output + + output, lora_result = maybe_execute_in_parallel( + base_fn, + lora_fn, + self._events[0], + self._events[1], + self._lora_stream, + ) + + original_shape = output.shape if output.ndim == 3 else None + + # In transformers backend, x and output have extra batch dimension like + # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim), + # therefore we need to flatten the batch dimensions. + if x.ndim == 3 and output.ndim == 3: + output = output.flatten(0, 1) + x = x.flatten(0, 1) + + output.add_(lora_result) + + # Reshape the flattened output back to its original shape, + # as some MM encoders cannot handle flattened inputs. + if original_shape is not None: + output = output.reshape(original_shape) + + return output + @property def weight(self) -> torch.Tensor: # unquantizedLinear diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 01efe3e4731..835bffe58ca 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -19,6 +19,9 @@ from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.fused_moe.config import ( _get_config_dtype_str, ) +from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( + UnfusedOAITritonExperts, +) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( MarlinExperts, ) @@ -28,9 +31,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe import ( from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import ( FusedMoEModularMethod, ) -from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( - UnfusedOAITritonExperts, -) from vllm.model_executor.layers.fused_moe.modular_kernel import ( FusedMoEKernel, ) diff --git a/vllm/lora/ops/triton_ops/lora_expand_op.py b/vllm/lora/ops/triton_ops/lora_expand_op.py index 343e0c81080..7f8ed577ecb 100644 --- a/vllm/lora/ops/triton_ops/lora_expand_op.py +++ b/vllm/lora/ops/triton_ops/lora_expand_op.py @@ -9,8 +9,13 @@ https://arxiv.org/abs/2310.18547 import torch +from vllm import envs from vllm.lora.ops.triton_ops.kernel_utils import do_expand_kernel -from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr, get_lora_op_configs +from vllm.lora.ops.triton_ops.utils import ( + _get_lora_b_ptr, + get_lora_op_configs, + supports_pdl, +) from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op @@ -237,9 +242,9 @@ def _lora_expand( NUM_SLICES, num_active_loras.item(), ) - # We disable PDL temporarily because LoRA kernels are not launching back-to-back, - # making PDL invalid and affecting the kernel performance. - use_gdc = False # supports_pdl(inputs.device) + + # PDL only works when dual-stream is being used. + use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM _lora_expand_kernel[grid]( inputs, lora_ptr_tensor, diff --git a/vllm/lora/ops/triton_ops/lora_shrink_op.py b/vllm/lora/ops/triton_ops/lora_shrink_op.py index ea850baa253..88c24c740db 100644 --- a/vllm/lora/ops/triton_ops/lora_shrink_op.py +++ b/vllm/lora/ops/triton_ops/lora_shrink_op.py @@ -9,8 +9,13 @@ https://arxiv.org/abs/2310.18547 import torch +from vllm import envs from vllm.lora.ops.triton_ops.kernel_utils import do_shrink_kernel -from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr, get_lora_op_configs +from vllm.lora.ops.triton_ops.utils import ( + _get_lora_a_ptr, + get_lora_op_configs, + supports_pdl, +) from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op @@ -220,9 +225,9 @@ def _lora_shrink( NUM_SLICES, num_active_loras.item(), ) - # We disable PDL temporarily because LoRA kernels are not launching back-to-back, - # making PDL invalid and affecting the kernel performance. - use_gdc = False # supports_pdl(inputs.device) + + # PDL only works when dual-stream is being used. + use_gdc = supports_pdl(inputs.device) and envs.VLLM_LORA_ENABLE_DUAL_STREAM _lora_shrink_kernel[grid]( inputs, lora_ptr_tensor, diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 5f2604892ce..321cbfcab7c 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -144,7 +144,9 @@ class PunicaWrapperGPU(PunicaWrapperBase): x (torch.Tensor): Input tensors lora_b_stacked (tuple[torch.Tensor, ...]): lora_b's weight output_slices (tuple[int, ...]): Every slice's size - add_inputs (bool): Defaults to True. + add_inputs (bool): If True, add LoRA output to y; if False, write + LoRA-only output to y (used for dual-stream when base and LoRA + run on different CUDA streams). Defaults to True. """ y_org = y y = y.view(-1, y.shape[-1]) @@ -161,7 +163,7 @@ class PunicaWrapperGPU(PunicaWrapperBase): num_tokens, self.lora_config.specialize_active_lora ), offset_start=offset_start, - add_inputs=True, + add_inputs=add_inputs, ) y = y.view_as(y_org) @@ -244,7 +246,7 @@ class PunicaWrapperGPU(PunicaWrapperBase): buffer = torch.empty( (len(output_slices), x.size(0), r), dtype=torch.float32, device=x.device ) - + add_inputs = kwargs.pop("add_inputs", True) self.add_shrink( buffer, # type: ignore x, @@ -257,7 +259,7 @@ class PunicaWrapperGPU(PunicaWrapperBase): buffer, # type: ignore lora_b_stacked, output_slices, - add_inputs=True, + add_inputs=add_inputs, **kwargs, ) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 774e92c228b..e0a9272c890 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -51,10 +51,50 @@ from vllm.model_executor.kernels.linear.mixed_precision.machete import ( from vllm.model_executor.kernels.linear.mixed_precision.marlin import ( MarlinLinearKernel, ) +from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( + TritonW4A16LinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.xpu import ( XPUW4A8IntLinearKernel, XPUwNa16LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp8 import ( + Mxfp8LinearKernel, + Mxfp8LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mxfp8.emulation import ( + EmulationMxfp8LinearKernel, +) +from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutlassMxfp8LinearKernel, +) +from vllm.model_executor.kernels.linear.mxfp8.marlin import ( + MarlinMxfp8LinearKernel, +) +from vllm.model_executor.kernels.linear.mxfp8.xpu import ( + XPUMxFp8LinearKernel, +) +from vllm.model_executor.kernels.linear.nvfp4 import ( + NvFp4LinearKernel, + NvFp4LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, +) +from vllm.model_executor.kernels.linear.nvfp4.emulation import ( + EmulationNvFp4LinearKernel, +) +from vllm.model_executor.kernels.linear.nvfp4.fbgemm import ( + FbgemmNvFp4LinearKernel, +) +from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCudnnNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + FlashInferTrtllmNvFp4LinearKernel, +) +from vllm.model_executor.kernels.linear.nvfp4.marlin import ( + MarlinNvFp4LinearKernel, +) from vllm.model_executor.kernels.linear.scaled_mm import ( Fp8BlockScaledMMLinearKernel, FP8ScaledMMLinearKernel, @@ -156,6 +196,21 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ ], } +_POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { + PlatformEnum.CUDA: [ + MarlinFP8ScaledMMLinearKernel, + ], + PlatformEnum.ROCM: [ + # To be added + ], + PlatformEnum.CPU: [ + # To be added + ], + PlatformEnum.XPU: [ + XPUFP8ScaledMMLinearKernel, + ], +} + # in priority/performance order (when available) _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { PlatformEnum.CUDA: [ @@ -167,6 +222,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { ExllamaLinearKernel, ], PlatformEnum.ROCM: [ + TritonW4A16LinearKernel, ConchLinearKernel, ExllamaLinearKernel, ], @@ -180,6 +236,37 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = { ], } +# in priority/performance order (when available) +_POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = { + PlatformEnum.CUDA: [ + FlashInferCutlassMxfp8LinearKernel, + MarlinMxfp8LinearKernel, + EmulationMxfp8LinearKernel, + ], + PlatformEnum.ROCM: [ + EmulationMxfp8LinearKernel, + ], + PlatformEnum.XPU: [ + XPUMxFp8LinearKernel, + EmulationMxfp8LinearKernel, + ], +} + +_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { + PlatformEnum.CUDA: [ + FlashInferCutlassNvFp4LinearKernel, + CutlassNvFp4LinearKernel, + MarlinNvFp4LinearKernel, + FlashInferTrtllmNvFp4LinearKernel, + FlashInferCudnnNvFp4LinearKernel, + FbgemmNvFp4LinearKernel, + EmulationNvFp4LinearKernel, + ], + PlatformEnum.ROCM: [ + EmulationNvFp4LinearKernel, + ], +} + # TODO make all kernels inherit from MMLinearKernel # then bound _KernelT only to MMLinearKernel _KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel | MMLinearKernel) @@ -426,6 +513,164 @@ def choose_mp_linear_kernel( ) +def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel: + """Select and instantiate the best MXFP8 linear kernel for the + current platform.""" + config = Mxfp8LinearLayerConfig() + + platform = current_platform._enum + possible = _POSSIBLE_MXFP8_KERNELS.get(platform, []) + + failure_reasons = [] + for kernel_cls in possible: + if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: + failure_reasons.append( + f" {kernel_cls.__name__} disabled by environment variable" + ) + continue + + is_supported, reason = kernel_cls.is_supported() + if not is_supported: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + can_implement, reason = kernel_cls.can_implement(config) + if not can_implement: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + logger.info_once("Using %s for MXFP8 GEMM", kernel_cls.__name__) + return kernel_cls(config) + + raise ValueError( + "Failed to find a kernel that can implement the " + "MXFP8 linear layer. Reasons: \n" + "\n".join(failure_reasons) + ) + + +def init_wfp8_a16_linear_kernel( + weight_quant_key: QuantKey, + activation_quant_key: QuantKey, + weight_shape: tuple[int, int], + input_dtype: torch.dtype, + out_dtype: torch.dtype, + force_kernel: type[FP8ScaledMMLinearKernel] | None = None, + module_name: str | None = None, +) -> FP8ScaledMMLinearKernel: + config = FP8ScaledMMLinearLayerConfig( + weight_quant_key=weight_quant_key, + activation_quant_key=activation_quant_key, + weight_shape=weight_shape, + input_dtype=input_dtype, + out_dtype=out_dtype, + ) + + kernel_type = choose_scaled_mm_linear_kernel( + config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel + ) + + if module_name: + logger.info_once( + "Selected %s for %s", + kernel_type.__name__, + module_name, + scope="global", + ) + + return kernel_type( + config, + layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"], + ) + + +# Maps VLLM_NVFP4_GEMM_BACKEND env var values to kernel classes. +_NVFP4_BACKEND_TO_KERNEL: dict[str, type[NvFp4LinearKernel]] = { + "flashinfer-cutlass": FlashInferCutlassNvFp4LinearKernel, + "cutlass": CutlassNvFp4LinearKernel, + "marlin": MarlinNvFp4LinearKernel, + "flashinfer-trtllm": FlashInferTrtllmNvFp4LinearKernel, + "flashinfer-cudnn": FlashInferCudnnNvFp4LinearKernel, + "emulation": EmulationNvFp4LinearKernel, +} + + +def init_nvfp4_linear_kernel() -> NvFp4LinearKernel: + """Select and instantiate the best NVFP4 linear kernel for the + current platform.""" + config = NvFp4LinearLayerConfig() + + # Env-var overrides. + force_kernel: type[NvFp4LinearKernel] | None = None + if envs.VLLM_BATCH_INVARIANT: + logger.info_once( + "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the " + "emulation backend for deterministic execution." + ) + force_kernel = EmulationNvFp4LinearKernel + elif envs.VLLM_USE_FBGEMM: + force_kernel = FbgemmNvFp4LinearKernel + elif envs.VLLM_USE_NVFP4_CT_EMULATIONS: + force_kernel = EmulationNvFp4LinearKernel + elif envs.VLLM_NVFP4_GEMM_BACKEND is not None: + backend_name = envs.VLLM_NVFP4_GEMM_BACKEND + force_kernel = _NVFP4_BACKEND_TO_KERNEL.get(backend_name) + if force_kernel is None: + raise ValueError( + f"Unknown VLLM_NVFP4_GEMM_BACKEND={backend_name!r}. " + f"Valid choices: {list(_NVFP4_BACKEND_TO_KERNEL.keys())}" + ) + + if force_kernel is not None: + is_supported, reason = force_kernel.is_supported() + if not is_supported: + raise ValueError( + f"Forced NVFP4 kernel {force_kernel.__name__} is not " + f"supported: {reason}" + ) + logger.info_once("Using %s for NVFP4 GEMM", force_kernel.__name__) + return force_kernel(config) + + # Auto-select from registry. + platform = current_platform._enum + possible = _POSSIBLE_NVFP4_KERNELS.get(platform, []) + + failure_reasons = [] + for kernel_cls in possible: + if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: + failure_reasons.append( + f" {kernel_cls.__name__} disabled by environment variable" + ) + continue + + is_supported, reason = kernel_cls.is_supported() + if not is_supported: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + can_implement, reason = kernel_cls.can_implement(config) + if not can_implement: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + if kernel_cls is EmulationNvFp4LinearKernel and failure_reasons: + logger.warning_once( + "NVFP4 linear falling back to the slow and unoptimized " + "emulation backend as no optimized backend is available " + "(unavailable reasons:\n - %s\n). " + "In case you expect one of these backends to be used, " + "please verify your environment.", + "\n - ".join(failure_reasons), + ) + + logger.info_once("Using %s for NVFP4 GEMM", kernel_cls.__name__) + return kernel_cls(config) + + raise ValueError( + "Failed to find a kernel that can implement the " + "NVFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons) + ) + + def register_linear_kernel( kernel_class: type, platform: PlatformEnum, @@ -455,6 +700,14 @@ def register_linear_kernel( if platform not in _POSSIBLE_FP8_KERNELS: _POSSIBLE_FP8_KERNELS[platform] = [] _POSSIBLE_FP8_KERNELS[platform].append(kernel_class) + elif kernel_type == "mxfp8": + if platform not in _POSSIBLE_MXFP8_KERNELS: + _POSSIBLE_MXFP8_KERNELS[platform] = [] + _POSSIBLE_MXFP8_KERNELS[platform].append(kernel_class) + elif kernel_type == "nvfp4": + if platform not in _POSSIBLE_NVFP4_KERNELS: + _POSSIBLE_NVFP4_KERNELS[platform] = [] + _POSSIBLE_NVFP4_KERNELS[platform].append(kernel_class) else: raise ValueError(f"Unrecognized kernel type: {kernel_type}") @@ -462,14 +715,18 @@ def register_linear_kernel( __all__ = [ "init_fp8_linear_kernel", "init_int8_linear_kernel", + "init_nvfp4_linear_kernel", "choose_mp_linear_kernel", "register_linear_kernel", + "init_wfp8_a16_linear_kernel", "FP8ScaledMMLinearKernel", "Int8ScaledMMLinearKernel", "ScaledMMLinearKernel", "FP8ScaledMMLinearLayerConfig", "Int8ScaledMMLinearLayerConfig", "ScaledMMLinearLayerConfig", + "NvFp4LinearKernel", + "NvFp4LinearLayerConfig", "AiterInt8ScaledMMLinearKernel", "CPUInt8ScaledMMLinearKernel", "CutlassFP8ScaledMMLinearKernel", @@ -490,8 +747,23 @@ __all__ = [ "ExllamaLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", + "TritonW4A16LinearKernel", "XPUW4A8IntLinearKernel", "XPUwNa16LinearKernel", + "init_mxfp8_linear_kernel", + "Mxfp8LinearKernel", + "Mxfp8LinearLayerConfig", + "FlashInferCutlassMxfp8LinearKernel", + "MarlinMxfp8LinearKernel", + "XPUMxFp8LinearKernel", + "EmulationMxfp8LinearKernel", + "CutlassNvFp4LinearKernel", + "EmulationNvFp4LinearKernel", + "FbgemmNvFp4LinearKernel", + "FlashInferCutlassNvFp4LinearKernel", + "FlashInferTrtllmNvFp4LinearKernel", + "FlashInferCudnnNvFp4LinearKernel", + "MarlinNvFp4LinearKernel", "_KernelT", "DeepGemmFp8BlockScaledMMKernel", "FlashInferFp8DeepGEMMDynamicBlockScaledKernel", diff --git a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py index 6c144a5ec8a..4d659b36042 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/__init__.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/__init__.py @@ -29,6 +29,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( MPLinearKernel, MPLinearLayerConfig, ) +from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import ( + TritonW4A16LinearKernel, +) from vllm.model_executor.kernels.linear.mixed_precision.xpu import ( XPUW4A8IntLinearKernel, XPUwNa16LinearKernel, @@ -45,6 +48,7 @@ __all__ = [ "ExllamaLinearKernel", "MacheteLinearKernel", "MarlinLinearKernel", + "TritonW4A16LinearKernel", "XPUW4A8IntLinearKernel", "XPUwNa16LinearKernel", ] diff --git a/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py new file mode 100644 index 00000000000..5cc100b3e1e --- /dev/null +++ b/vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Triton-based W4A16 GEMM kernel for ROCm MI300. + +Implements fused int4-weight dequantization + fp16 GEMM in a single kernel, +using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]). +Plugs into the MPLinearKernel selection system and is preferred over +MarlinLinearKernel/ExllamaLinearKernel on ROCm. + +Weight layout expected by this kernel (post-process_weights_after_loading): + qweight: [K, N//8] int32 — rows=K (input), cols=N//8 (N is packed) + scales: [K//G, N] fp16/bf16 + qzeros: [K//G, N//8] int32 (optional; None for symmetric uint4b8) + +Checkpoint layout from compressed_tensors_wNa16 create_weights: + weight_packed: [N, K//8] int32 (output_dim=0, input_dim=1, packed_dim=1) + weight_scale: [N, K//G] fp16 (output_dim=0, input_dim=1) + weight_zero_point: [N//8, K//G] int32 (output_dim=0, packed_dim=0) +""" + +import torch + +from vllm.model_executor.layers.quantization.utils import replace_parameter +from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ +from vllm.platforms import current_platform +from vllm.scalar_type import scalar_types +from vllm.triton_utils import tl, triton + +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig + +TRITON_W4A16_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128, 256] +TRITON_W4A16_SUPPORTED_QUANT_TYPES = [ + scalar_types.uint4b8, # symmetric GPTQ (bias=8) + scalar_types.uint4, # asymmetric with explicit zeros +] + + +@triton.jit +def triton_w4a16_gemm_kernel( + # Pointers + a_ptr, # [M, K] fp16/bf16 activations + b_ptr, # [K, N//8] int32 packed 4-bit weights (N is the packed dim) + scales_ptr, # [K//G, N] fp16/bf16 scales + zeros_ptr, # [K//G, N//8] int32 packed zeros (unused when HAS_ZP=False) + c_ptr, # [M, N] fp16/bf16 output + # Dimensions + M, + N, + K, + # Strides + stride_am, + stride_ak, + stride_bk, + stride_bn, # stride in b along the packed N//8 dim + stride_cm, + stride_cn, + # Quantization parameters + group_size, + # Whether explicit zero points are provided + HAS_ZP: tl.constexpr, + # Zero bias used when HAS_ZP is False (e.g. 8 for uint4b8) + ZP_BIAS: tl.constexpr, + # Block sizes (tuned for MI300 wavefront=64) + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """ + Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N] + + B is stored as [K, N//8] int32 using GPTQ sequential packing: + each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28]. + + Dequant: w_fp = (w_int4 - zero) * scale + HAS_ZP=True: zero is loaded from zeros_ptr and unpacked + HAS_ZP=False: zero = ZP_BIAS constant (e.g. 8 for uint4b8 symmetric) + """ + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + # Row/col offsets for this tile + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + # b/zeros are stored with N packed: N//8 int32 columns per K row + offs_bn = pid_n * (BLOCK_N // 8) + tl.arange(0, BLOCK_N // 8) + + # GPTQ sequential shifts tiled across BLOCK_N: + # [0,4,8,...,28] repeating for every group of 8 N-values. + # Build 1D shifts_1d of length BLOCK_N: column j gets shift (j % 8) * 4. + shifts_row = tl.arange(0, 8) * 4 # [8] + shifts_1d_2d = tl.broadcast_to(shifts_row[None, :], (BLOCK_N // 8, 8)) + shifts_1d = tl.reshape(shifts_1d_2d, (BLOCK_N,)) # [BLOCK_N] + # Broadcast to [BLOCK_K, BLOCK_N] for weight unpacking + shifts = tl.broadcast_to(shifts_1d[None, :], (BLOCK_K, BLOCK_N)) + + # Scales column offsets: full N-width (one scale per output neuron) + offs_sn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + for k_start in range(0, tl.cdiv(K, BLOCK_K)): + offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K) + mask_k = offs_k < K + + # ---- Load activations A: [BLOCK_M, BLOCK_K] ---- + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + mask_a = (offs_m[:, None] < M) & mask_k[None, :] + a = tl.load(a_ptrs, mask=mask_a, other=0.0) + + # ---- Load packed weights B: [BLOCK_K, BLOCK_N//8] int32 ---- + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn + mask_b = mask_k[:, None] & (offs_bn[None, :] < N // 8) + b_packed = tl.load(b_ptrs, mask=mask_b, other=0) + + # ---- Unpack int4 weights → [BLOCK_K, BLOCK_N] ---- + # tl.interleave(x, x) doubles the last dim by interleaving. + # Starting from [BLOCK_K, BLOCK_N//8], three interleaves give + # [BLOCK_K, BLOCK_N], where each int32 is replicated 8 times. + b = tl.interleave(b_packed, b_packed) + b = tl.interleave(b, b) + b = tl.interleave(b, b) + # Extract the correct 4-bit nibble for each output column + b = (b >> shifts) & 0xF + + # ---- Compute scale/zero group row index ---- + g_idx = (k_start * BLOCK_K) // group_size + + # ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ---- + scale_offset = g_idx * N + offs_sn + scale_mask = offs_sn < N + scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0) + scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N)) + + # ---- Load / compute zeros ---- + if HAS_ZP: + # Load packed zeros row: [BLOCK_N//8] int32 + zero_offset = g_idx * (N // 8) + offs_bn + zero_mask = offs_bn < N // 8 + z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0) + # Unpack to [BLOCK_N] using same interleave+shift pattern + z = tl.interleave(z_packed, z_packed) + z = tl.interleave(z, z) + z = tl.interleave(z, z) + z = (z >> shifts_1d) & 0xF + z = tl.broadcast_to(z[None, :], (BLOCK_K, BLOCK_N)) + else: + z = tl.full((BLOCK_K, BLOCK_N), ZP_BIAS, dtype=tl.int32) + + # ---- Dequantize: (w - zero) * scale ---- + b_fp = (b - z).to(a.dtype) * scales + + # ---- Accumulate ---- + accumulator += tl.dot(a, b_fp, out_dtype=tl.float32) + + # ---- Store output C: [BLOCK_M, BLOCK_N] ---- + c = accumulator.to(c_ptr.type.element_ty) + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(c_ptrs, c, mask=mask_c) + + +def triton_w4a16_gemm( + a: torch.Tensor, # [M, K] fp16/bf16 + b_q: torch.Tensor, # [K, N//8] int32 + scales: torch.Tensor, # [K//G, N] fp16/bf16 + qzeros: torch.Tensor | None, # [K//G, N//8] int32, or None + group_size: int, + zp_bias: int = 8, # bias for uint4b8 when qzeros is None +) -> torch.Tensor: + """ + Fused W4A16 GEMM using GPTQ-packed int4 weights. + + Args: + a: Activation matrix [M, K], float16 or bfloat16. + b_q: Packed weight matrix [K, N//8], int32 (GPTQ sequential). + scales: Per-group scales [K//G, N], same dtype as a. + qzeros: Per-group packed zero points [K//G, N//8] int32, or None + for symmetric quantization (uses zp_bias instead). + group_size: Quantization group size (resolved from -1 to K by caller). + zp_bias: Constant zero used when qzeros is None (default 8 for uint4b8). + + Returns: + Output matrix [M, N], same dtype as a. + """ + assert a.is_contiguous(), "Activation matrix must be contiguous" + assert b_q.is_contiguous(), "Weight matrix must be contiguous" + assert scales.is_contiguous(), "Scales must be contiguous" + + M, K = a.shape + N = b_q.shape[1] * 8 + + assert b_q.shape == (K, N // 8), ( + f"b_q shape mismatch: {b_q.shape} vs ({K}, {N // 8})" + ) + assert scales.shape == (K // group_size, N), ( + f"scales shape mismatch: {scales.shape} vs ({K // group_size}, {N})" + ) + if qzeros is not None: + assert qzeros.shape == (K // group_size, N // 8), ( + f"qzeros shape mismatch: {qzeros.shape}" + ) + + c = torch.empty((M, N), dtype=a.dtype, device=a.device) + + has_zp = qzeros is not None + # Provide a dummy pointer when HAS_ZP=False (Triton requires a valid ptr) + zeros_ptr = qzeros if has_zp else b_q + + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1x + + if on_gfx1x(): + # Tuned for RDNA 3.5 (gfx1151, 40 CUs, 32-wide wavefronts). + if M <= 32: + BLOCK_M, BLOCK_N, BLOCK_K = 32, 32, 64 + elif M <= 64: + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 + else: + BLOCK_M, BLOCK_N, BLOCK_K = 128, 32, 64 + else: + # Tuned for MI300 (gfx942, 304 CUs, 64-wide wavefronts). + if M <= 32: + BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32 + elif M <= 64: + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 + else: + BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32 + else: + if M <= 32: + BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32 + elif M <= 64: + BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 + else: + BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32 + + # The kernel loads scales/zeros for a single group per BLOCK_K tile + # (one g_idx per iteration). If BLOCK_K > group_size, rows at the tail + # of the tile dequantize with the wrong group's scales, silently + # corrupting the output. Clamp BLOCK_K to group_size to keep one + # scale group per tile. + if group_size < BLOCK_K: + BLOCK_K = group_size + + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) + + triton_w4a16_gemm_kernel[grid]( + a, + b_q, + scales, + zeros_ptr, + c, + M, + N, + K, + a.stride(0), + a.stride(1), + b_q.stride(0), + b_q.stride(1), + c.stride(0), + c.stride(1), + group_size=group_size, + HAS_ZP=has_zp, + ZP_BIAS=zp_bias, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + ) + return c + + +class TritonW4A16LinearKernel(MPLinearKernel): + """ + Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer). + + Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric) + with grouped quantization. Weight tensors are transposed from the + compressed-tensors checkpoint layout to the kernel's [K, N//8] layout. + """ + + SUPPORTED_QUANT_TYPES = TRITON_W4A16_SUPPORTED_QUANT_TYPES + + @classmethod + def get_min_capability(cls) -> int: + # Triton handles capability checks itself + return 0 + + @classmethod + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: + if not current_platform.is_rocm(): + return False, "TritonW4A16LinearKernel only targets ROCm" + + if c.weight_type not in cls.SUPPORTED_QUANT_TYPES: + return ( + False, + f"Quant type {c.weight_type} not supported; " + f"supported: {cls.SUPPORTED_QUANT_TYPES}", + ) + + if c.act_type not in (torch.float16, torch.bfloat16): + return False, "Only float16/bfloat16 activations are supported" + + N = c.partition_weight_shape[1] + if N % 8 != 0: + return ( + False, + f"Output features ({N}) must be divisible by 8 " + "(8 int4 values packed per int32)", + ) + + if c.has_g_idx: + return ( + False, + "Activation reordering (g_idx) is not supported by " + "TritonW4A16LinearKernel", + ) + + gs = c.group_size + if ( + gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES + and gs != c.full_weight_shape[0] + ): + return ( + False, + f"Group size {gs} not supported; " + f"supported: {TRITON_W4A16_SUPPORTED_GROUP_SIZES} " + f"or full K ({c.full_weight_shape[0]})", + ) + + K = c.partition_weight_shape[0] + eff_gs = gs if gs != -1 else K + if K % eff_gs != 0: + return (False, f"Input features {K} not divisible by group size {eff_gs}") + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """ + Convert compressed-tensors checkpoint layout to kernel layout. + + Checkpoint (from compressed_tensors_wNa16.create_weights): + weight_packed: [N, K//8] int32 input_dim=1, output_dim=0, packed_dim=1 + weight_scale: [N, K//G] fp16 input_dim=1, output_dim=0 + weight_zero_point: [N//8, K//G] int32 output_dim=0, packed_dim=0 + + Kernel needs: + qweight: [K, N//8] int32 (transpose weight_packed) + scales: [K//G, N] fp16 (transpose weight_scale) + qzeros: [K//G, N//8] int32 (transpose weight_zero_point) + """ + + # ---- Transform qweight: [N, K//8] → [K//8, N] → back to [K, N//8] ---- + # permute_param_layout_(x, input_dim=0, output_dim=1) rearranges so that + # the input(K) dimension is at physical dim 0 and output(N) at dim 1. + # Checkpoint has input_dim=1, output_dim=0, packed_dim=1 (K is packed). + # After permute we get [K//8, N] (K packed at dim 0, N at dim 1). + # The kernel wants [K, N//8] (K at dim 0, N packed at dim 1), so we + # then transpose: [K//8, N].T = [N, K//8] — that's not right. + # + # Actually we need to change WHAT is packed: + # Original packing: K packed into K//8 (8 K-values per int32) + # Kernel packing: N packed into N//8 (8 N-values per int32) + # These require a full repack, not just a transpose. + # + # Simple approach: unpack → transpose the full [N, K] → repack as [K, N//8]. + # This is done CPU-side at load time (one-time cost). + def repack_w_q(x: BasevLLMParameter) -> BasevLLMParameter: + # x.data is [N, K//8] int32, K packed (GPTQ checkpoint format) + # Step 1: bring to [N, K//8] with output(N) at dim 0 + permute_param_layout_(x, input_dim=1, output_dim=0, packed_dim=1) + w = x.data # [N, K//8] int32 + + N_dim, K8 = w.shape + K_dim = K8 * 8 + # Step 2: unpack to [N, K] int32 (vectorized) + shifts = torch.arange(8, device=w.device, dtype=torch.int32) * 4 + w_unpacked = ((w.unsqueeze(-1) >> shifts) & 0xF).reshape(N_dim, K_dim) + # Step 3: transpose to [K, N] int32 + w_KN = w_unpacked.t().contiguous() + # Step 4: repack N into N//8 int32 values → [K, N//8] (vectorized) + N8 = N_dim // 8 + w_repacked = torch.sum( + (w_KN.view(K_dim, N8, 8) & 0xF) << shifts, + dim=2, + dtype=torch.int32, + ) + x.data = w_repacked.contiguous() + return x + + def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter: + # x.data is [N, K//G] fp16, bring to [K//G, N] + permute_param_layout_(x, input_dim=1, output_dim=0) + x.data = x.data.t().contiguous() + return x + + self._transform_param(layer, self.w_q_name, repack_w_q) + self._transform_param(layer, self.w_s_name, repack_w_s) + + if self.w_zp_name is not None: + zp = getattr(layer, self.w_zp_name, None) + if zp is not None: + # Checkpoint: [N//8, K//G] int32 (N packed at dim 0, K//G at dim 1) + # Kernel needs: [K//G, N//8] — just transpose + replace_parameter( + layer, + self.w_zp_name, + torch.nn.Parameter(zp.data.t().contiguous(), requires_grad=False), + ) + + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: + c = self.config + w_q, w_s, w_zp, _ = self._get_weight_params(layer) + + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) + + K = c.partition_weight_shape[0] + group_size = c.group_size if c.group_size != -1 else K + + # For symmetric types (uint4b8), use the scalar bias; no zeros tensor + zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0 + + output = triton_w4a16_gemm( + a=x_2d, + b_q=w_q, + scales=w_s, + qzeros=w_zp, + group_size=group_size, + zp_bias=zp_bias, + ) + + if bias is not None: + output.add_(bias) + + return output.reshape(out_shape) diff --git a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py new file mode 100644 index 00000000000..28c958d4fd4 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch + + +@dataclass +class Mxfp8LinearLayerConfig: + """Configuration for an MXFP8 linear layer. + + All MXFP8 layers share the same structure: FP8-E4M3 weights with + uint8 (E8M0) per-block scales at block size 32. + """ + + pass + + +class Mxfp8LinearKernel(ABC): + """Base class for MXFP8 quantized linear kernels. + + Each subclass implements a specific GEMM backend (FlashInfer CUTLASS, + Marlin, emulation). + """ + + def __init__(self, c: Mxfp8LinearLayerConfig) -> None: + assert self.can_implement(c)[0] + assert self.is_supported()[0] + self.config = c + + @classmethod + @abstractmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + raise NotImplementedError + + @classmethod + @abstractmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + raise NotImplementedError diff --git a/vllm/model_executor/kernels/linear/mxfp8/__init__.py b/vllm/model_executor/kernels/linear/mxfp8/__init__.py new file mode 100644 index 00000000000..507aedee14c --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearKernel, + Mxfp8LinearLayerConfig, +) + +__all__ = [ + "Mxfp8LinearKernel", + "Mxfp8LinearLayerConfig", +] diff --git a/vllm/model_executor/kernels/linear/mxfp8/emulation.py b/vllm/model_executor/kernels/linear/mxfp8/emulation.py new file mode 100644 index 00000000000..a7cc29be758 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/emulation.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch.nn.parameter import Parameter + +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + dequant_mxfp8_to_bf16, +) + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +class EmulationMxfp8LinearKernel(Mxfp8LinearKernel): + """Software emulation fallback for MXFP8 (dequant to BF16).""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight.data # [N, K] + N, K = weight.shape + scale_k = K // MXFP8_BLOCK_SIZE + + weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() + + layer.weight = Parameter(weight.contiguous(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + weight_scale = layer.weight_scale + if weight_scale.dtype != MXFP8_SCALE_DTYPE: + raise ValueError( + f"Emulation backend requires {MXFP8_SCALE_DTYPE} " + f"weight_scale dtype, got {weight_scale.dtype}." + ) + if weight_scale.ndim != 2: + raise ValueError( + f"Emulation backend requires 2D weight_scale, " + f"got {weight_scale.ndim}D. " + f"Ensure process_weights_after_loading was called." + ) + + weight_bf16 = dequant_mxfp8_to_bf16(layer.weight, weight_scale) + output = torch.nn.functional.linear(x, weight_bf16, bias) + return output.to(x.dtype) diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py new file mode 100644 index 00000000000..336da511ad8 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch.nn.parameter import Parameter + +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + mxfp8_e4m3_quantize, + swizzle_mxfp8_scale, +) +from vllm.platforms import current_platform +from vllm.utils import flashinfer as vllm_flashinfer + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel): + """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+).""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if current_platform.has_device_capability(100): + return True, None + return False, "requires >=sm_100 (Blackwell)" + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight = layer.weight.data # [N, K] + N, K = weight.shape + + scale_k = K // MXFP8_BLOCK_SIZE + weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous() + weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K) + + layer.weight = Parameter(weight.contiguous(), requires_grad=False) + layer.weight_scale = Parameter( + weight_scale_swizzled.contiguous(), requires_grad=False + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + weight = layer.weight + weight_scale = layer.weight_scale + out_dtype = x.dtype + N, K = weight.shape + + input_shape = x.shape + input_2d = x.view(-1, K) + M_orig = input_2d.shape[0] + + min_dim = 128 + + assert min_dim <= K, ( + f"mm_mxfp8 requires K >= {min_dim}, got K={K}. " + f"in_features is too small for mm_mxfp8." + ) + assert K % MXFP8_BLOCK_SIZE == 0, ( + f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}." + ) + assert min_dim <= N, ( + f"mm_mxfp8 requires N >= {min_dim}, got N={N}. " + f"out_features is too small for mm_mxfp8." + ) + + M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim + if M_padded != M_orig: + pad_rows = M_padded - M_orig + input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows)) + + input_mxfp8, input_scale = mxfp8_e4m3_quantize( + input_2d, is_sf_swizzled_layout=True + ) + + if not weight.is_contiguous(): + weight = weight.contiguous() + + output = vllm_flashinfer.mm_mxfp8( + input_mxfp8, + weight.t(), + input_scale, + weight_scale, + out_dtype=out_dtype, + backend="cutlass", + ) + + if M_padded != M_orig: + output = output[:M_orig, :] + + if bias is not None: + output = output + bias + + output_shape = (*input_shape[:-1], N) + return output.view(output_shape) diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py new file mode 100644 index 00000000000..bec54cd942e --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +class MarlinMxfp8LinearKernel(Mxfp8LinearKernel): + """MXFP8 W8A16 GEMM via Marlin (SM80+).""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + is_fp8_marlin_supported, + ) + + if is_fp8_marlin_supported(): + return True, None + return False, "Marlin FP8 not available" + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + prepare_mxfp8_layer_for_marlin, + ) + + prepare_mxfp8_layer_for_marlin(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + apply_mxfp8_marlin_linear, + ) + + return apply_mxfp8_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/xpu.py b/vllm/model_executor/kernels/linear/mxfp8/xpu.py new file mode 100644 index 00000000000..d64e175c295 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp8/xpu.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + xpu_mxfp8_quantize as quant_mxfp8, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig + + +class XPUMxFp8LinearKernel(Mxfp8LinearKernel): + """MXFP8 W8A8 GEMM on XPU.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if not current_platform.is_xpu(): + return False, "XPUMxFp8 only support on XPU" + return True, None + + @classmethod + def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu) + weight_scale = weight_scale.t().contiguous() + replace_parameter(layer, "weight", layer.weight.t()) + replace_parameter(layer, "weight_scale", weight_scale.data) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + out_dtype = x.dtype + x_fp8, x_scale = quant_mxfp8(x) + return torch.ops._xpu_C.fp8_gemm( + x_fp8, + layer.weight, + out_dtype, + x_scale, + layer.weight_scale, + bias, + ) diff --git a/vllm/model_executor/kernels/linear/nvfp4/__init__.py b/vllm/model_executor/kernels/linear/nvfp4/__init__.py new file mode 100644 index 00000000000..de72584057d --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.kernels.linear.nvfp4.base import ( + NvFp4LinearKernel, + NvFp4LinearLayerConfig, +) + +__all__ = [ + "NvFp4LinearKernel", + "NvFp4LinearLayerConfig", +] diff --git a/vllm/model_executor/kernels/linear/nvfp4/base.py b/vllm/model_executor/kernels/linear/nvfp4/base.py new file mode 100644 index 00000000000..24e0aa30892 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/base.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch + + +@dataclass +class NvFp4LinearLayerConfig: + """Configuration for an NVFP4 linear layer. + + All NVFP4 layers share the same structure: packed uint8 weights (2 FP4 values per + byte), FP8-E4M3 per-block weight scales (group size 16), and scalar global + scales for both weights and activations. + """ + + pass + + +class NvFp4LinearKernel(ABC): + """Base class for NVFP4 quantized linear kernels. + + Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc). + The kernel selection mechanism iterates over registered subclasses in + priority order,calling ``is_supported`` and ``can_implement`` to find the best + match for the current hardware. + """ + + def __init__(self, config: NvFp4LinearLayerConfig) -> None: + assert self.can_implement(config)[0] + assert self.is_supported()[0] + self.config = config + + @classmethod + @abstractmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + """Return whether this kernel can run on the current platform.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + """Return whether this kernel can handle *config*.""" + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Transform weights into the format required by this kernel. + + Called once after checkpoint weights have been loaded onto the + device. Implementations should repack / swizzle / pad weights + and scales in-place on *layer*. + """ + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the quantized GEMM.""" + raise NotImplementedError diff --git a/vllm/model_executor/kernels/linear/nvfp4/cutlass.py b/vllm/model_executor/kernels/linear/nvfp4/cutlass.py new file mode 100644 index 00000000000..0d0663dca17 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/cutlass.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm._custom_ops import ( + cutlass_scaled_fp4_mm, + scaled_fp4_quant, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + cutlass_fp4_supported, + pad_nvfp4_activation_for_cutlass, + pad_nvfp4_weight_for_cutlass, + slice_nvfp4_output, + swizzle_blockscale, +) + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + + +class CutlassNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM via the vLLM CUTLASS kernel.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if cutlass_fp4_supported(): + return True, None + return False, "CUTLASS FP4 kernels not available" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.weight_scale.data), requires_grad=False + ) + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( + layer.weight.data + ) + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) + layer.weights_padding_cols = weights_padding_cols + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="cutlass", + ) + + x_fp4 = pad_nvfp4_activation_for_cutlass( + x_fp4, getattr(layer, "weights_padding_cols", 0) + ) + + out = cutlass_scaled_fp4_mm( + x_fp4, + layer.weight, + x_blockscale, + layer.weight_scale, + layer.alpha, + output_dtype, + ) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) diff --git a/vllm/model_executor/kernels/linear/nvfp4/emulation.py b/vllm/model_executor/kernels/linear/nvfp4/emulation.py new file mode 100644 index 00000000000..2a55b317767 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/emulation.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + kE2M1ToFloat_handle, + run_nvfp4_emulations, +) + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + + +class EmulationNvFp4LinearKernel(NvFp4LinearKernel): + """Software emulation fallback for NVFP4 (dequant → BF16 matmul).""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + # Always available as a last-resort fallback. + return True, None + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Move the E2M1 lookup table to the device now, because + # `.to(device)` is not allowed during CUDA graph capture. + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + out = run_nvfp4_emulations( + x=x, + input_global_scale=layer.input_global_scale_inv, + weight=layer.weight, + weight_scale_swizzled=layer.weight_scale, + weight_global_scale=layer.weight_global_scale, + swizzle=False, + ) + if bias is not None: + out = out + bias + return out diff --git a/vllm/model_executor/kernels/linear/nvfp4/fbgemm.py b/vllm/model_executor/kernels/linear/nvfp4/fbgemm.py new file mode 100644 index 00000000000..fa30a75c47d --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/fbgemm.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm._custom_ops import scaled_fp4_quant +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + slice_nvfp4_output, + swizzle_blockscale, +) +from vllm.utils.import_utils import has_fbgemm_gpu + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + + +class FbgemmNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM via FBGEMM.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if has_fbgemm_gpu(): + return True, None + return False, "fbgemm_gpu required" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + swizzled = swizzle_blockscale(layer.weight_scale.data) + layer.weight_scale = torch.nn.Parameter( + swizzled.view(-1).view(torch.uint8), requires_grad=False + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + import fbgemm_gpu # noqa: F401 - registers torch.ops.fbgemm.* + + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="fbgemm", + ) + + out = torch.ops.fbgemm.f4f4bf16( + x_fp4, + layer.weight, + x_blockscale.view(-1).view(torch.uint8), + layer.weight_scale, + layer.alpha, + use_mx=False, + ).to(output_dtype) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py new file mode 100644 index 00000000000..399bc3dd278 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm._custom_ops import scaled_fp4_quant +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + pad_nvfp4_activation_for_cutlass, + pad_nvfp4_weight_for_cutlass, + slice_nvfp4_output, + swizzle_blockscale, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + + +class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM via FlashInfer's CUTLASS wrapper.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + cutlass_fp4_supported, + ) + + if ( + cutlass_fp4_supported() + and current_platform.has_device_capability(100) + and has_flashinfer() + ): + return True, None + return False, "FlashInfer + >=sm_100 required" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.weight_scale.data), requires_grad=False + ) + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( + layer.weight.data + ) + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) + layer.weights_padding_cols = weights_padding_cols + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="flashinfer-cutlass", + ) + + x_fp4 = pad_nvfp4_activation_for_cutlass( + x_fp4, getattr(layer, "weights_padding_cols", 0) + ) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_blockscale, + layer.weight_scale, + layer.alpha, + output_dtype, + backend="cutlass", + ) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) + + +class FlashInferTrtllmNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM via FlashInfer's TensorRT-LLM wrapper.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if has_flashinfer(): + return True, None + return False, "FlashInfer required" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a + + weight = layer.weight.data + weight_scale = layer.weight_scale.data + epilogue_tile_m = 128 + + layer.weight = torch.nn.Parameter( + shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m), + requires_grad=False, + ) + layer.weight_scale = torch.nn.Parameter( + shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m) + .reshape(weight_scale.shape) + .view(torch.float8_e4m3fn), + requires_grad=False, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="flashinfer-trtllm", + ) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_blockscale, + layer.weight_scale, + layer.alpha, + output_dtype, + backend="trtllm", + ) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) + + +class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM via FlashInfer's cuDNN wrapper.""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if has_flashinfer(): + return True, None + return False, "FlashInfer required" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # cuDNN uses the same swizzled + padded layout as CUTLASS + layer.weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.weight_scale.data), requires_grad=False + ) + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( + layer.weight.data + ) + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) + layer.weights_padding_cols = weights_padding_cols + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="flashinfer-cudnn", + ) + + x_fp4 = pad_nvfp4_activation_for_cutlass( + x_fp4, getattr(layer, "weights_padding_cols", 0) + ) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_blockscale, + layer.weight_scale, + layer.alpha, + output_dtype, + backend="cudnn", + ) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) diff --git a/vllm/model_executor/kernels/linear/nvfp4/marlin.py b/vllm/model_executor/kernels/linear/nvfp4/marlin.py new file mode 100644 index 00000000000..a05d6823c88 --- /dev/null +++ b/vllm/model_executor/kernels/linear/nvfp4/marlin.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( + apply_fp4_marlin_linear, + is_fp4_marlin_supported, + prepare_fp4_layer_for_marlin, +) + +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig + +logger = init_logger(__name__) + + +class MarlinNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 weight-only GEMM via Marlin (W4A16).""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if is_fp4_marlin_supported(): + return True, None + return False, "Marlin FP4 not available" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + logger.warning_once( + "Your GPU does not have native support for FP4 computation but " + "FP4 quantization is being used. Weight-only FP4 compression " + "will be used leveraging the Marlin kernel. This may degrade " + "performance for compute-heavy workloads." + ) + prepare_fp4_layer_for_marlin(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return apply_fp4_marlin_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + weight_global_scale=layer.weight_global_scale, + workspace=layer.workspace, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index b16ee169972..6d75a420e7e 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -9,6 +9,11 @@ from vllm.model_executor.kernels.linear import ( # noqa: E501 FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticChannelSym, + kFp8StaticTensorSym, +) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -23,6 +28,11 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): @classmethod def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: + if c.weight_quant_key not in {kFp8StaticChannelSym, kFp8StaticTensorSym}: + return ( + False, + "XPUFP8ScaledMM only support per-channel and per-tensor quantization", + ) if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}: return False, "XPUFP8ScaledMM only support FP8 weight dtype" return True, None @@ -35,6 +45,9 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): self.config = c self.layer_param_names = layer_param_names + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + replace_parameter(layer, "weight", layer.weight.data.t()) + def apply_weights( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 27cc211912b..a92e2f4ad18 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -25,6 +25,9 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, direct_register_custom_op, kv_cache_dtype_str_to_dtype, ) @@ -376,6 +379,10 @@ class Attention(nn.Module, AttentionLayerBase): # Initialize KV cache quantization attributes _init_kv_cache_quant(self, quant_config, prefix) + # Initialize TurboQuant buffers (Pi, S, centroids) if tq cache dtype + if kv_cache_dtype.startswith("turboquant_"): + self._init_turboquant_buffers(kv_cache_dtype, head_size, prefix) + # for attn backends supporting query quantization self.query_quant = None if ( @@ -394,6 +401,67 @@ class Attention(nn.Module, AttentionLayerBase): else GroupShape.PER_TENSOR, ) + def _init_turboquant_buffers( + self, cache_dtype: str, head_size: int, prefix: str + ) -> None: + """Initialize TurboQuant rotation/projection matrices and centroids.""" + from vllm.model_executor.layers.quantization.turboquant.centroids import ( + get_centroids, + ) + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + from vllm.model_executor.layers.quantization.turboquant.quantizer import ( + generate_wht_signs, + ) + + tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size) + + # Each layer needs a unique rotation matrix so quantization errors + # don't correlate across layers. Stride must exceed max head_dim to + # ensure non-overlapping RNG streams between adjacent layers. + _TQ_LAYER_SEED_STRIDE = 1337 + + from vllm.model_executor.models.utils import extract_layer_index + + layer_idx = extract_layer_index(prefix) + seed = tq_config.seed + layer_idx * _TQ_LAYER_SEED_STRIDE + + self.register_buffer( + "_tq_signs", + generate_wht_signs(head_size, seed=seed), + ) + self.register_buffer( + "_tq_centroids", + get_centroids(head_size, tq_config.centroid_bits), + ) + self._tq_config = tq_config + + # Pre-allocate decode intermediate buffers so model.to(device) moves + # them to GPU *before* the memory profiler runs. Without this the + # profiler gives all free memory to KV cache blocks and the first + # decode OOMs when these buffers are lazily allocated. + _vllm_cfg = get_current_vllm_config() + B = _vllm_cfg.scheduler_config.max_num_seqs + Hq = self.num_heads + S = _vllm_cfg.attention_config.tq_max_kv_splits_for_cuda_graph + D = head_size + self.register_buffer( + "_tq_mid_o_buf", + torch.empty(B, Hq, S, D + 1, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_tq_output_buf", + torch.empty(B, Hq, D, dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_tq_lse_buf", + torch.empty(B, Hq, dtype=torch.float32), + persistent=False, + ) + def forward( self, query: torch.Tensor, @@ -414,7 +482,9 @@ class Attention(nn.Module, AttentionLayerBase): `vllm.forward_context.get_forward_context().attn_metadata`. """ if self.calculate_kv_scales: - torch.ops.vllm.maybe_calc_kv_scales(query, key, value, self.layer_name) + torch.ops.vllm.maybe_calc_kv_scales( + query, key, value, _encode_layer_name(self.layer_name) + ) output_dtype = query.dtype if self.query_quant is not None: # quantizing with a simple torch operation enables @@ -466,6 +536,7 @@ class Attention(nn.Module, AttentionLayerBase): ) else: # Skip this if sharing KV cache with an earlier attention layer. + encoded = _encode_layer_name(self.layer_name) if ( not self.attn_backend.forward_includes_kv_cache_update and self.kv_sharing_target_layer_name is None @@ -473,14 +544,14 @@ class Attention(nn.Module, AttentionLayerBase): and value is not None ): kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update( - key, value, self.layer_name + key, value, encoded ) torch.ops.vllm.unified_attention_with_output( query, key, value, output, - self.layer_name, + encoded, kv_cache_dummy_dep=kv_cache_dummy_dep, ) return output.view(-1, hidden_size) @@ -538,6 +609,23 @@ class Attention(nn.Module, AttentionLayerBase): kv_quant_mode=quant_mode, sliding_window=self.sliding_window, ) + elif self.kv_cache_dtype.startswith("turboquant_"): + from vllm.model_executor.layers.quantization.turboquant.config import ( + TurboQuantConfig, + ) + from vllm.v1.kv_cache_interface import TQFullAttentionSpec + + tq_config = TurboQuantConfig.from_cache_dtype( + self.kv_cache_dtype, self.head_size + ) + return TQFullAttentionSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=self.head_size, + dtype=self.kv_cache_torch_dtype, + tq_slot_size=tq_config.slot_size_aligned, + ) else: return FullAttentionSpec( block_size=block_size, @@ -553,8 +641,9 @@ def maybe_calc_kv_scales( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: + layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] @@ -570,7 +659,7 @@ def maybe_calc_kv_scales_fake( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: return @@ -622,12 +711,13 @@ def get_attention_context( def unified_kv_cache_update( key: torch.Tensor, value: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> torch.Tensor: """ Returns a dummy that is passed to unified_attention to signal a side effect and the data dependency between them to ensure torch.compile preserves ordering. """ + layer_name = _resolve_layer_name(layer_name) _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name) if layer_slot_mapping is not None: assert hasattr(attn_layer.impl, "do_kv_cache_update"), ( @@ -647,7 +737,7 @@ def unified_kv_cache_update( def unified_kv_cache_update_fake( key: torch.Tensor, value: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> torch.Tensor: return torch.empty(0, device=key.device, dtype=key.dtype) @@ -666,7 +756,7 @@ def unified_attention_with_output( key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, kv_cache_dummy_dep: torch.Tensor | None = None, @@ -675,6 +765,7 @@ def unified_attention_with_output( # that ensures torch.compile preserves ordering between KV cache update and # attention forward. del kv_cache_dummy_dep + layer_name = _resolve_layer_name(layer_name) attn_metadata, self, kv_cache, _ = get_attention_context(layer_name) self.impl.forward( @@ -695,7 +786,7 @@ def unified_attention_with_output_fake( key: torch.Tensor, value: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, kv_cache_dummy_dep: torch.Tensor | None = None, diff --git a/vllm/model_executor/layers/attention/kv_transfer_utils.py b/vllm/model_executor/layers/attention/kv_transfer_utils.py index 4afc5ccb165..1dcd445b5ee 100644 --- a/vllm/model_executor/layers/attention/kv_transfer_utils.py +++ b/vllm/model_executor/layers/attention/kv_transfer_utils.py @@ -9,6 +9,7 @@ from vllm.distributed.kv_transfer import ( has_kv_transfer_group, is_v1_kv_transfer_group, ) +from vllm.utils.torch_utils import _resolve_layer_name def maybe_transfer_kv_layer(func: Callable) -> Callable: @@ -38,7 +39,7 @@ def maybe_transfer_kv_layer(func: Callable) -> Callable: if not has_kv_transfer_group() or not is_v1_kv_transfer_group(): return func(*args, **kwargs) - layer_name: str = args[layer_name_index] + layer_name = _resolve_layer_name(args[layer_name_index]) # Extract attention context (metadata, layer, kv_cache, layer_slot_mapping) attn_metadata, _, kv_cache, _ = get_attention_context(layer_name) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 1d046b16e2a..cbbf5f3c3ca 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -240,6 +240,9 @@ from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory from vllm.utils.math_utils import cdiv, round_down from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, direct_register_custom_op, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, @@ -473,7 +476,12 @@ class MLAAttention(nn.Module, AttentionLayerBase): output_shape: torch.Size | None = None, ) -> torch.Tensor: if self.calculate_kv_scales: - torch.ops.vllm.maybe_calc_kv_scales(q, kv_c_normed, k_pe, self.layer_name) + torch.ops.vllm.maybe_calc_kv_scales( + q, + kv_c_normed, + k_pe, + _encode_layer_name(self.layer_name), + ) if self.use_direct_call: forward_context: ForwardContext = get_forward_context() @@ -505,10 +513,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): ) return output else: + encoded = _encode_layer_name(self.layer_name) kv_cache_dummy_dep = torch.ops.vllm.unified_mla_kv_cache_update( kv_c_normed, k_pe, - self.layer_name, + encoded, self.kv_cache_dtype, self._k_scale, ) @@ -518,7 +527,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): kv_c_normed, k_pe, output, - self.layer_name, + encoded, kv_cache_dummy_dep=kv_cache_dummy_dep, ) return output @@ -900,7 +909,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): def unified_mla_kv_cache_update( kv_c_normed: torch.Tensor, k_pe: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, kv_cache_dtype: str, k_scale: torch.Tensor, ) -> torch.Tensor: @@ -908,6 +917,7 @@ def unified_mla_kv_cache_update( Returns a dummy that is passed to unified_attention to signal a side effect and the data dependency between them to ensure torch.compile preserves ordering. """ + layer_name = _resolve_layer_name(layer_name) forward_context = get_forward_context() attn_layer = forward_context.no_compile_layers[layer_name] kv_cache = attn_layer.kv_cache @@ -939,7 +949,7 @@ def unified_mla_kv_cache_update( def unified_mla_kv_cache_update_fake( kv_c_normed: torch.Tensor, k_pe: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, kv_cache_dtype: str, k_scale: torch.Tensor, ) -> torch.Tensor: @@ -959,7 +969,7 @@ def unified_mla_attention_with_output( kv_c_normed: torch.Tensor, k_pe: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, kv_cache_dummy_dep: torch.Tensor | None = None, @@ -968,6 +978,7 @@ def unified_mla_attention_with_output( # that ensures torch.compile preserves ordering between KV cache update and # attention forward. del kv_cache_dummy_dep + layer_name = _resolve_layer_name(layer_name) attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name) layer.forward_impl( q, @@ -986,7 +997,7 @@ def unified_mla_attention_with_output_fake( kv_c_normed: torch.Tensor, k_pe: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, kv_cache_dummy_dep: torch.Tensor | None = None, @@ -1432,6 +1443,19 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): scope="local", ) return model_dtype + elif ( + is_quantized_kv_cache(vllm_config.cache_config.cache_dtype) + and backend_supports_prefill_query_quantization() + ): + logger.warning_once( + "FP8 KV cache is enabled but prefill queries are not " + "quantized to FP8. For long-context workloads (ISL >= 4K), " + "enabling FP8 prefill attention can significantly optimize " + "prefill latency. To enable, add: " + '--attention-config \'{"use_prefill_query_quantization"' + ": true}'", + scope="local", + ) return model_dtype diff --git a/vllm/model_executor/layers/attention/static_sink_attention.py b/vllm/model_executor/layers/attention/static_sink_attention.py index 263d873218f..8d199be0a57 100644 --- a/vllm/model_executor/layers/attention/static_sink_attention.py +++ b/vllm/model_executor/layers/attention/static_sink_attention.py @@ -10,7 +10,12 @@ from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.attention import Attention from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) from vllm.v1.attention.backend import ( AttentionBackend, AttentionMetadata, @@ -170,7 +175,9 @@ class StaticSinkAttention(Attention, CustomOp): ) if not self.sink_populated: self_kv_cache = self.kv_cache - torch.ops.vllm.maybe_populate_sink(self_kv_cache, self.layer_name) + torch.ops.vllm.maybe_populate_sink( + self_kv_cache, _encode_layer_name(self.layer_name) + ) return super().forward(query, key, value, output_shape) @@ -224,8 +231,9 @@ class StaticSinkAttention(Attention, CustomOp): def maybe_populate_sink( self_kv_cache: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: + layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] if self.sink_populated or self_kv_cache.numel() == 0: @@ -235,7 +243,7 @@ def maybe_populate_sink( def maybe_populate_sink_fake( self_kv_cache: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: return diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent.py b/vllm/model_executor/layers/fla/ops/fused_recurrent.py index 17b59b5bce7..920efa44417 100644 --- a/vllm/model_executor/layers/fla/ops/fused_recurrent.py +++ b/vllm/model_executor/layers/fla/ops/fused_recurrent.py @@ -106,12 +106,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 else: i_t = 0 - # Load state index and check for PAD_SLOT_ID (-1) + # Load state index and check for invalid entries state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( tl.int64 ) - # Skip if state index is invalid (PAD_SLOT_ID = -1) - if state_idx < 0: + # Skip if state index is invalid (NULL_BLOCK_ID=0) + if state_idx <= 0: return p_h0 = h0 + state_idx * stride_init_state_token else: @@ -150,12 +150,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( # keep the states for multi-query tokens if INPLACE_FINAL_STATE: - # Load state index and check for PAD_SLOT_ID (-1) + # Load state index and check for invalid entries final_state_idx = tl.load( ssm_state_indices + i_n * stride_indices_seq + i_t ).to(tl.int64) - # Only store if state index is valid (not PAD_SLOT_ID) - if final_state_idx >= 0: + # Only store if state index is valid (not NULL_BLOCK_ID=0) + if final_state_idx > 0: p_ht = ht + final_state_idx * stride_final_state_token p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) @@ -292,7 +292,8 @@ def fused_recurrent_gated_delta_rule_packed_decode_kernel( state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64) p_o = o + (i_n * HV + i_hv) * V + o_v - if state_idx < 0: + # Skip if state index is invalid (NULL_BLOCK_ID=0) + if state_idx <= 0: zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty) tl.store(p_o, zero, mask=mask_v) return diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 07ed185413f..7e0c7e05cab 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -106,12 +106,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel( i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 else: i_t = 0 - # Load state index and check for PAD_SLOT_ID (-1) + # Load state index and check for invalid entries state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( tl.int64 ) - # Skip if state index is invalid (PAD_SLOT_ID = -1) - if state_idx < 0: + # Skip if state index is invalid (NULL_BLOCK_ID=0) + if state_idx <= 0: return p_h0 = h0 + state_idx * stride_init_state_token else: @@ -155,12 +155,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel( # keep the states for multi-query tokens if INPLACE_FINAL_STATE: - # Load state index and check for PAD_SLOT_ID (-1) + # Load state index and check for invalid entries final_state_idx = tl.load( ssm_state_indices + i_n * stride_indices_seq + i_t ).to(tl.int64) - # Only store if state index is valid (not PAD_SLOT_ID) - if final_state_idx >= 0: + # Only store if state index is valid (not NULL_BLOCK_ID=0) + if final_state_idx > 0: p_ht = ht + final_state_idx * stride_final_state_token p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) diff --git a/vllm/model_executor/layers/fused_moe/__init__.py b/vllm/model_executor/layers/fused_moe/__init__.py index b342e0c6e91..926f0d1d015 100644 --- a/vllm/model_executor/layers/fused_moe/__init__.py +++ b/vllm/model_executor/layers/fused_moe/__init__.py @@ -33,9 +33,6 @@ from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) -from vllm.model_executor.layers.fused_moe.zero_expert_fused_moe import ( - ZeroExpertFusedMoE, -) from vllm.triton_utils import HAS_TRITON _config: dict[str, Any] | None = None @@ -68,7 +65,6 @@ __all__ = [ "GateLinear", "RoutingMethodType", "SharedFusedMoE", - "ZeroExpertFusedMoE", "activation_without_mul", "apply_moe_activation", "override_config", diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 534004e112f..62b2602928f 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -225,6 +225,14 @@ def maybe_make_prepare_finalize( elif moe.use_fi_nvl_one_sided_kernels: assert quant_config is not None + if quant_config.quant_dtype != "nvfp4": + raise ValueError( + "The 'flashinfer_nvlink_one_sided' all2all backend only " + "supports nvfp4 activation quantization, but got " + f"quant_dtype={quant_config.quant_dtype!r}. Use a different " + "all2all backend (e.g. 'flashinfer_nvlink_two_sided' or " + "'allgather_reducescatter') for non-nvfp4 models." + ) max_num_tokens = ( get_current_vllm_config().scheduler_config.max_num_batched_tokens ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 0c93dc6a76f..a3b941dfa45 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -6,8 +6,7 @@ from typing import Union import torch -import vllm.envs as envs -from vllm.config import ParallelConfig +from vllm.config import ParallelConfig, SchedulerConfig from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -937,15 +936,6 @@ class FusedMoEParallelConfig: all2all_backend: str # all2all backend for MoE communication enable_eplb: bool # whether to enable expert load balancing - @property - def use_dp_chunking(self) -> bool: - return ( - self.use_deepep_ll_kernels - or self.use_mori_kernels - or self.use_fi_nvl_two_sided_kernels - or self.use_nixl_ep_kernels - ) and envs.VLLM_ENABLE_MOE_DP_CHUNK - @property def is_sequence_parallel(self) -> bool: return self.sp_size > 1 @@ -1184,7 +1174,7 @@ class FusedMoEConfig: intermediate_size_per_partition_unpadded: int | None = None moe_backend: str = "auto" - max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE + max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP has_bias: bool = False is_act_and_mul: bool = True is_lora_enabled: bool = False diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..bcec61632e3 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "96": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..705ca33d594 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2 + }, + "1024": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json new file mode 100644 index 00000000000..9c2ebaddd83 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1536,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,dtype=fp8_w8a8,block_shape=[128,128].json @@ -0,0 +1,147 @@ +{ + "triton_version": "3.6.0", + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + } +} diff --git a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py similarity index 100% rename from vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py rename to vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 6f490f00be1..1f0258fb657 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -358,6 +358,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit if self.routing_method_type == RoutingMethodType.DeepSeekV3: router_logits = router_logits.to(torch.float32) + # Currently FI requires bfloat16 routing bias. + # https://github.com/flashinfer-ai/flashinfer/issues/2909 + if e_score_correction_bias is not None: + e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16) + is_mxfp8 = self.quant_config.block_shape == [1, 32] if is_mxfp8: fp8_quant_type = Fp8QuantizationType.MxFp8 diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 81b778c8f4a..fc30815f719 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -50,6 +50,9 @@ class TrtLlmNvFp4ExpertsBase: moe_config.intermediate_size_per_partition ) self.hidden_dim = moe_config.hidden_dim + self.hidden_dim_unpadded = ( + moe_config.hidden_dim_unpadded or moe_config.hidden_dim + ) self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank @@ -114,8 +117,12 @@ class TrtLlmNvFp4ExpertsBase: @staticmethod def _supports_shape(hidden_dim: int) -> bool: - """Requires hidden dim to be multiple of 512.""" - return hidden_dim % 512 == 0 + # Weights are zero-padded to 256-alignment at load time and the MoE + # runner pads activations via _maybe_pad_hidden_states, so any + # hidden_dim is accepted. + # NOTE: non-256-aligned dims will trigger a warning log and may + # cause performance degradation due to activation slicing. + return True @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -194,7 +201,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula import vllm.utils.flashinfer as fi_utils if fi_utils._is_fi_autotuning: - return hidden_states + return # Invoke kernel. flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe( @@ -324,6 +331,8 @@ class TrtLlmNvFp4ExpertsMonolithic( e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16) # Invoke kernel. + # NOTE: Activation padding and output + # truncation are handled by the MoE runner's return flashinfer.fused_moe.trtllm_fp4_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index 26409804c48..0d47b0f3174 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -130,7 +130,14 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): p.is_device_capability(90) or p.is_device_capability_family(100) or p.is_device_capability_family(110) - or p.is_device_capability_family(120) + or p.is_device_capability(120) + # NOTE: SM121 (DGX Spark) is excluded because the bf16 + # unquantized CUTLASS MoE GEMM in flashinfer <= 0.6.7 has no + # Relu2 template instantiation and throws "Invalid activation + # type" on Nemotron-H. Fixed upstream by + # https://github.com/flashinfer-ai/flashinfer/pull/2926 + # (merged 2026-04-01, not yet in a stable release); lift this + # restriction once flashinfer >= 0.6.8 is the minimum. ) and has_flashinfer_cutlass_fused_moe() ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index c4fc1fd2557..190a9cc3b5d 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -8,7 +8,6 @@ from typing import Literal, cast, get_args, overload import torch from torch.nn.parameter import UninitializedParameter -import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.parallel import ExpertPlacementStrategy @@ -19,7 +18,7 @@ from vllm.distributed import ( ) from vllm.distributed.eplb.eplb_state import EplbLayerState, EplbState from vllm.logger import init_logger -from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -214,8 +213,8 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str: # --8<-- [start:fused_moe] -@CustomOp.register("fused_moe") -class FusedMoE(CustomOp): +@PluggableLayer.register("fused_moe") +class FusedMoE(PluggableLayer): """FusedMoE layer for MoE models. This layer contains both MergedColumnParallel weights (gate_up_proj / @@ -275,6 +274,7 @@ class FusedMoE(CustomOp): gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, routed_input_transform: torch.nn.Module | None = None, + zero_expert_type: str | None = None, ): super().__init__() @@ -463,6 +463,8 @@ class FusedMoE(CustomOp): # TODO(bnell): once we can construct the MK at init time, we # can make this a value. indices_type_getter=lambda: self.quant_method.topk_indices_dtype, + zero_expert_type=zero_expert_type, + num_logical_experts=self.logical_num_experts, ) self.routing_method_type: RoutingMethodType = self.router.routing_method_type @@ -479,7 +481,7 @@ class FusedMoE(CustomOp): in_dtype=moe_in_dtype, moe_backend=vllm_config.kernel_config.moe_backend, router_logits_dtype=router_logits_dtype, - max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE, + max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens, has_bias=has_bias, is_act_and_mul=is_act_and_mul, is_lora_enabled=vllm_config.lora_config is not None, @@ -842,7 +844,10 @@ class FusedMoE(CustomOp): if shard_id == "w2": hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim) expert_data = self._narrow_expert_data_for_padding( - expert_data, loaded_weight, hidden_dim=hidden_dim + expert_data, + loaded_weight, + hidden_dim=hidden_dim, + shard_dim=shard_dim, ) expert_data.copy_(loaded_weight) elif shard_id in ("w1", "w3"): @@ -882,29 +887,33 @@ class FusedMoE(CustomOp): expert_data: torch.Tensor, loaded_weight: torch.Tensor, hidden_dim: int, + shard_dim: int | None = None, ) -> torch.Tensor: - """Narrow expert_data hidden dim to match loaded_weight for padded - hidden_size. + """Narrow expert_data to match loaded_weight for padded dimensions. When backends (e.g., DeepEP) round up hidden_size, weight parameters are larger than checkpoint weights. Narrow the padded hidden dimension - before copying. + before copying. Similarly, when padding occurs on the shard + (intermediate) dimension (e.g. for MXFP4 GEMM), narrow that dimension + as well. Args: expert_data: The (possibly padded) parameter tensor to narrow. loaded_weight: The checkpoint weight tensor with original size. hidden_dim: The dimension index corresponding to hidden_size. Must be non-negative. + shard_dim: The dimension index corresponding to the shard + (intermediate) dimension. Defaults to `None`. """ - if ( - loaded_weight.ndim > 0 - and 0 <= hidden_dim < expert_data.ndim - and hidden_dim < loaded_weight.ndim - and expert_data.shape[hidden_dim] > loaded_weight.shape[hidden_dim] - ): - expert_data = expert_data.narrow( - hidden_dim, 0, loaded_weight.shape[hidden_dim] - ) + dims = (hidden_dim,) if shard_dim is None else (hidden_dim, shard_dim) + if loaded_weight.ndim > 0: + for dim in dims: + if ( + 0 <= dim < expert_data.ndim + and dim < loaded_weight.ndim + and expert_data.shape[dim] > loaded_weight.shape[dim] + ): + expert_data = expert_data.narrow(dim, 0, loaded_weight.shape[dim]) return expert_data def _load_w13( @@ -946,7 +955,10 @@ class FusedMoE(CustomOp): expert_data = expert_data.narrow(shard_dim, shard_size, shard_size) hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim) expert_data = self._narrow_expert_data_for_padding( - expert_data, loaded_weight, hidden_dim=hidden_dim + expert_data, + loaded_weight, + hidden_dim=hidden_dim, + shard_dim=shard_dim, ) expert_data.copy_(loaded_weight) @@ -979,7 +991,10 @@ class FusedMoE(CustomOp): # w2, down_proj: Load into only logical weight of w2. hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim) expert_data = self._narrow_expert_data_for_padding( - expert_data, loaded_weight, hidden_dim=hidden_dim + expert_data, + loaded_weight, + hidden_dim=hidden_dim, + shard_dim=shard_dim, ) expert_data.copy_(loaded_weight) @@ -1063,7 +1078,7 @@ class FusedMoE(CustomOp): expert_id: int, return_success: bool = False, ) -> bool | None: - if self.quant_config and self.quant_config.get_name() == "mxfp4": + if self.quant_config and self.quant_config.get_name() == "gpt_oss_mxfp4": # (FIXME) for gpt-oss all experts are combined if "bias" in weight_name: dim1 = loaded_weight.shape[1] @@ -1520,7 +1535,7 @@ class FusedMoE(CustomOp): """ return self.runner.maybe_all_reduce_tensor_model_parallel(final_hidden_states) - def forward_native( + def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -1536,13 +1551,6 @@ class FusedMoE(CustomOp): self._expert_map if not self.rocm_aiter_fmoe_enabled else self.expert_mask ) - def forward_cuda( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - return self.forward_native(hidden_states, router_logits) - @classmethod def make_expert_params_mapping( cls, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index d4a0817e0be..917d474fc9b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -101,7 +101,7 @@ def backend_to_kernel_cls( return [FlashInferExperts] elif backend == Mxfp4MoeBackend.TRITON: - from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( + from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 OAITritonExperts, OAITritonMxfp4ExpertsMonolithic, ) @@ -110,7 +110,7 @@ def backend_to_kernel_cls( return [OAITritonMxfp4ExpertsMonolithic, OAITritonExperts] elif backend == Mxfp4MoeBackend.TRITON_UNFUSED: - from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( + from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 UnfusedOAITritonExperts, ) @@ -194,7 +194,7 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: return None -def select_mxfp4_moe_backend( +def select_gpt_oss_mxfp4_moe_backend( config: FusedMoEConfig, ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]: """ @@ -400,7 +400,7 @@ def mxfp4_round_up_hidden_size_and_intermediate_size( return hidden_size, intermediate_size -def convert_to_mxfp4_moe_kernel_format( +def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend: Mxfp4MoeBackend, layer: torch.nn.Module, w13_weight: torch.Tensor, @@ -426,7 +426,10 @@ def convert_to_mxfp4_moe_kernel_format( sf_block_size = 32 # mxfp4 block size - if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN): + if mxfp4_backend in ( + Mxfp4MoeBackend.MARLIN, + Mxfp4MoeBackend.BATCHED_MARLIN, + ): from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( prepare_moe_mxfp4_layer_for_marlin, ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index b061b3d38b8..5b93b3d5c6e 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -176,11 +176,27 @@ class RoutedExpertsCapturer: end_loc = topk_ids.shape[0] token_num_per_dp = topk_ids.shape[0] else: # multi dp - token_num_per_dp = ctx.dp_metadata.num_tokens_across_dp_cpu[self.dp_rank] - cumsum = torch.cumsum(ctx.dp_metadata.num_tokens_across_dp_cpu, dim=0) - assert cumsum[-1] == topk_ids.shape[0] - end_loc = cumsum[self.dp_rank] - start_loc = end_loc - token_num_per_dp + num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu + token_num_per_dp = int(num_tokens_dp[self.dp_rank].item()) + total = int(num_tokens_dp.sum().item()) + n = topk_ids.shape[0] + + if n == total: + # Naive dispatch: all DP ranks' tokens concatenated before routing. + cumsum = torch.cumsum(num_tokens_dp, dim=0) + end_loc = int(cumsum[self.dp_rank].item()) + start_loc = end_loc - token_num_per_dp + elif n == token_num_per_dp: + # Modular-kernel path: DP combine happens inside quant_method.apply; + # select_experts only sees this rank's tokens. + start_loc = 0 + end_loc = token_num_per_dp + else: + raise AssertionError( + "RoutedExpertsCapturer: unexpected topk_ids batch dim " + f"{n} (expected {total} or {token_num_per_dp} " + f"for dp_rank={self.dp_rank})" + ) if layer_id >= self._device_buffer.shape[1]: return diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index 11027e894be..42d418d7e53 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import ( RoutingSimulatorRouter, ) +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) EMPTY_EPLB_STATE: EplbLayerState = EplbLayerState() @@ -49,6 +52,9 @@ def create_fused_moe_router( # eplb parameters enable_eplb: bool = False, eplb_state: EplbLayerState = EMPTY_EPLB_STATE, + # zero expert parameters + zero_expert_type: str | None = None, + num_logical_experts: int | None = None, ) -> FusedMoERouter: """ Factory function to create the appropriate FusedMoERouter subclass based on @@ -56,10 +62,11 @@ def create_fused_moe_router( The selection logic follows this priority order: 1. RoutingSimulatorRouter - if VLLM_MOE_ROUTING_SIMULATION_STRATEGY env var is set - 2. GroupedTopKRouter - if use_grouped_topk is True - 3. CustomRoutingRouter - if custom_routing_function is not None - 4. FusedTopKBiasRouter - if e_score_correction_bias is not None - 5. FusedTopKRouter - default fallback + 2. ZeroExpertRouter - if zero_expert_type is not None + 3. GroupedTopKRouter - if use_grouped_topk is True + 4. CustomRoutingRouter - if custom_routing_function is not None + 5. FusedTopKBiasRouter - if e_score_correction_bias is not None + 6. FusedTopKRouter - default fallback Common arguments: top_k: Number of experts to select per token @@ -86,6 +93,12 @@ def create_fused_moe_router( enable_eplb: Whether EPLB is enabled eplb_state: EPLB (Expert Parallelism Load Balancing) state + Zero expert arguments: + zero_expert_type: Type of zero expert (e.g. identity). If not None, + creates a ZeroExpertRouter. + num_logical_experts: Number of real (non-zero) experts. Required when + zero_expert_type is not None. + Returns: An instance of the appropriate FusedMoERouter subclass """ @@ -100,6 +113,27 @@ def create_fused_moe_router( indices_type_getter=indices_type_getter, ) + if zero_expert_type is not None: + assert num_logical_experts is not None, ( + "num_logical_experts is required when zero_expert_type is set" + ) + assert e_score_correction_bias is not None, ( + "e_score_correction_bias is required when zero_expert_type is set" + ) + return ZeroExpertRouter( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + e_score_correction_bias=e_score_correction_bias, + num_logical_experts=num_logical_experts, + zero_expert_type=zero_expert_type, + scoring_func=scoring_func, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + enable_eplb=enable_eplb, + indices_type_getter=indices_type_getter, + ) + if use_grouped_topk: assert custom_routing_function is None if num_expert_group is None or topk_group is None: diff --git a/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py new file mode 100644 index 00000000000..c87070bc5ac --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/zero_expert_router.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable + +import torch + +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, + get_routing_method_type, +) +from vllm.model_executor.layers.fused_moe.fused_moe import ( + zero_experts_compute_triton, +) +from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) + + +class ZeroExpertRouter(BaseRouter): + """Router that handles zero expert computation as part of routing. + + Routes over all experts (real + zero) using full e_score_correction_bias. + Computes zero expert identity contributions as a side effect during routing. + Remaps zero expert IDs to real expert ID 0 (with weight 0) so downstream + MoE computation can ignore them. + """ + + def __init__( + self, + top_k: int, + global_num_experts: int, + eplb_state: EplbLayerState, + e_score_correction_bias: torch.Tensor, + num_logical_experts: int, + zero_expert_type: str, + scoring_func: str = "softmax", + renormalize: bool = False, + routed_scaling_factor: float = 1.0, + enable_eplb: bool = False, + indices_type_getter: Callable[[], torch.dtype | None] | None = None, + ): + super().__init__( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + enable_eplb=enable_eplb, + indices_type_getter=indices_type_getter, + ) + self.e_score_correction_bias = e_score_correction_bias + self.num_logical_experts = num_logical_experts + self.zero_expert_type = zero_expert_type + self.scoring_func = scoring_func + self.renormalize = renormalize + self.routed_scaling_factor = routed_scaling_factor + self._zero_expert_output: torch.Tensor | None = None + + @property + def routing_method_type(self) -> RoutingMethodType: + return get_routing_method_type( + scoring_func=self.scoring_func, + top_k=self.top_k, + renormalize=self.renormalize, + num_expert_group=None, + has_e_score_bias=True, + ) + + def _compute_routing( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + indices_type: torch.dtype | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute routing with full bias, compute zero expert output, + mask zero expert IDs.""" + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + e_score_correction_bias=self.e_score_correction_bias.data, + topk=self.top_k, + renormalize=self.renormalize, + scoring_func=self.scoring_func, + indices_type=indices_type, + ) + + if self.routed_scaling_factor != 1.0: + topk_weights *= self.routed_scaling_factor + + # Compute zero expert output using pre-EPLB topk_ids/weights. + # zero_experts_compute_triton modifies its inputs in-place, so + # pass clones. + self._zero_expert_output = zero_experts_compute_triton( + expert_indices=topk_ids.clone(), + expert_scales=topk_weights.clone(), + num_experts=self.num_logical_experts, + zero_expert_type=self.zero_expert_type, + hidden_states=hidden_states, + ) + + # Mask zero expert entries: remap zero expert IDs to 0 with weight 0 + # so downstream MoE computation ignores them. + zero_mask = topk_ids >= self.num_logical_experts + topk_ids[zero_mask] = 0 + topk_weights[zero_mask] = 0.0 + + return topk_weights, topk_ids + + @property + def zero_expert_output(self) -> torch.Tensor | None: + """Retrieve and clear the zero expert output.""" + output = self._zero_expert_output + self._zero_expert_output = None + return output diff --git a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py deleted file mode 100644 index a8c75486d71..00000000000 --- a/vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py +++ /dev/null @@ -1,243 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -from vllm.forward_context import ( - get_forward_context, -) -from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( - FusedMoEMethodBase, -) -from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, -) -from vllm.utils.math_utils import cdiv -from vllm.v1.worker.ubatching import dbo_current_ubatch_id -from vllm.v1.worker.workspace import current_workspace_manager - - -class ChunkingMoERunner(MoERunnerBase): - """ - MoE runner wrapper that adds chunked processing to any MoERunnerBase. - - This runner wraps an inner MoERunnerBase and overrides _forward_impl to - process large batches by breaking them into smaller chunks. Each chunk - is delegated to the inner runner's _forward_impl, making chunking - composable with any runner implementation. - - All MoERunnerBase state (moe_config, router, quant_method, etc.) is - transparently delegated to the inner runner via __getattr__. - ChunkingMoERunner only owns chunking-specific state: the pre-allocated - workspace buffers and the reduce_results override. - - Key behaviors: - - Pre-allocates workspace tensors for CUDA graph compatibility - - Processes chunks via inner._forward_impl per chunk - - Never reduces results (reduce_results always returns False) - """ - - def __init__(self, inner: MoERunnerBase): - # Assert that _maybe_dispatch/_maybe_combine will be nops. - assert inner.moe_config.pcp_size == 1 - - # Skip MoERunnerBase.__init__ — all state is delegated to inner - # via __getattr__. Only chunking-specific state lives here. - self._inner = inner - - # Pre-allocated staging buffers. These need to exist ahead of time - # due to CUDA graph construction needing fixed buffer addresses. - self.batched_hidden_states, self.batched_router_logits = ( - self._init_dp_chunking() - ) - - def __getattr__(self, name): - # Delegate attribute access to the inner runner. This is only - # called when normal lookup (instance __dict__, class MRO) fails, - # so ChunkingMoERunner's own attributes and methods take priority. - return getattr(self._inner, name) - - @property - def shared_experts(self) -> SharedExperts | None: - return self._inner.shared_experts - - # TODO(bnell): temporary hack, do not call this method. - def _replace_quant_method(self, quant_method: FusedMoEMethodBase): - self._inner._replace_quant_method(quant_method) - self.quant_method = quant_method - - def is_internal_router(self) -> bool: - return self._inner.gate is not None - - # Reducing results when chunking is handled by the MK finalize operations - # when DP chunking is enabled.. - # This will be removed by #35949 - @property - def reduce_results(self) -> bool: - return False - - def _init_dp_chunking(self) -> list[torch.Tensor]: - states_shape: tuple[int, ...] - logits_shape: tuple[int, ...] - - moe = self.moe_config - - if self.enable_dbo: - states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts) - else: - states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim) - logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts) - - # Does this need some kind of profiling run check like modular_kernel.py? - return current_workspace_manager().get_simultaneous( - (states_shape, moe.in_dtype), - (logits_shape, moe.router_logits_dtype), - ) - - def _allocate_dp_chunking_outputs( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> tuple[torch.Tensor | None, torch.Tensor]: - # Assert the inputs are of the proper type and shape. - assert self.batched_hidden_states is not None - assert self.batched_router_logits is not None - - assert self.batched_hidden_states.dtype == hidden_states.dtype, ( - f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}" - ) - assert self.batched_router_logits.dtype == router_logits.dtype, ( - f"{self.batched_router_logits.dtype} == {router_logits.dtype}" - ) - - # Check size compatibility. - assert self.batched_hidden_states.size(-1) == hidden_states.size(-1) - assert self.batched_router_logits.size(-1) == router_logits.size(-1) - - final_fused_hidden_states = torch.empty_like(hidden_states) - if self.shared_experts is not None: - if shared_experts_input is not None: - final_shared_hidden_states = torch.empty_like(shared_experts_input) - else: - final_shared_hidden_states = torch.empty_like(hidden_states) - else: - final_shared_hidden_states = None - - return final_shared_hidden_states, final_fused_hidden_states - - def _slice_and_copy_input( - self, - out_slice: torch.Tensor, - orig: torch.Tensor | None, - start: int, - end: int, - ) -> torch.Tensor: - assert orig is not None - slice_size = end - start - orig_slice = orig[start:end, :] - if self.enable_dbo: - assert out_slice.dim() == 3 - batch_buffer_idx = dbo_current_ubatch_id() - out_slice = out_slice[batch_buffer_idx, :] - - assert out_slice.size(0) >= slice_size - out_slice = out_slice[:slice_size, :] - out_slice.copy_(orig_slice, non_blocking=True) - return out_slice - - def _forward_impl( - self, - layer: torch.nn.Module, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - shared_experts_input: torch.Tensor | None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - final_shared_hidden_states, final_fused_hidden_states = ( - self._allocate_dp_chunking_outputs( - hidden_states, router_logits, shared_experts_input - ) - ) - - ctx = get_forward_context() - # flashinfer_cutlass_kernels can handle: optional DP + TP/EP - max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu - moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens - - # If the input to the MoE is sequence parallel then divide by sp_size - # to find the maximum number of tokens for any individual dispatcher. - if self.moe_config.is_sequence_parallel: - max_tokens_across_dispatchers = cdiv( - max_tokens_across_dispatchers, self.moe_config.sp_size - ) - - num_tokens = hidden_states.size(0) - for chunk_idx, chunk_start_ in enumerate( - range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank) - ): - chunk_start = chunk_start_ - chunk_end = min( - chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers - ) - # clamp start and end - chunk_start = min(chunk_start, num_tokens - 1) - chunk_end = min(chunk_end, num_tokens) - chunk_sizes = ctx.dp_metadata.chunked_sizes( - self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx - ) - with chunk_sizes: - hidden_states_chunk = self._slice_and_copy_input( - self.batched_hidden_states, - hidden_states, - chunk_start, - chunk_end, - ) - - router_logits_chunk = self._slice_and_copy_input( - self.batched_router_logits, - router_logits, - chunk_start, - chunk_end, - ) - - shared_experts_input_chunk = ( - shared_experts_input[chunk_start:chunk_end, :] - if shared_experts_input is not None - else None - ) - - # Delegate per-chunk computation to the inner runner. - chunk_result = self._inner._forward_impl( - layer=layer, - hidden_states=hidden_states_chunk, - router_logits=router_logits_chunk, - shared_experts_input=shared_experts_input_chunk, - ) - - # Store outputs - # TODO(bnell): document when chunk_start >= num_tokens - if chunk_start < num_tokens: - if self.shared_experts is not None: - assert isinstance(chunk_result, tuple) - shared_output_chunk, hidden_states_chunk = chunk_result - final_fused_hidden_states[chunk_start:chunk_end, :].copy_( - hidden_states_chunk, non_blocking=True - ) - assert shared_output_chunk is not None - assert final_shared_hidden_states is not None - final_shared_hidden_states[chunk_start:chunk_end, :].copy_( - shared_output_chunk, non_blocking=True - ) - else: - assert isinstance(chunk_result, torch.Tensor) - final_fused_hidden_states[chunk_start:chunk_end, :].copy_( - chunk_result, non_blocking=True - ) - - if self.shared_experts is None: - return final_fused_hidden_states - else: - assert final_shared_hidden_states is not None - return (final_shared_hidden_states, final_fused_hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py index 481e787e279..692d45d3460 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py @@ -25,6 +25,9 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) +from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( + ZeroExpertRouter, +) from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( SharedExperts, @@ -32,15 +35,15 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( ) from vllm.platforms import current_platform from vllm.utils.torch_utils import ( - HAS_OPAQUE_TYPE, - ModuleName, + _USE_LAYERNAME, + LayerName, direct_register_custom_op, ) def get_layer_from_name(layer_name: str) -> torch.nn.Module: forward_context: ForwardContext = get_forward_context() - if layer_name == "from_forward_context": + if not _USE_LAYERNAME and layer_name == "from_forward_context": all_moe_layers = forward_context.all_moe_layers assert all_moe_layers is not None moe_layer_index = forward_context.moe_layer_index @@ -55,21 +58,21 @@ def get_layer_from_name(layer_name: str) -> torch.nn.Module: return forward_context.no_compile_layers[layer_name] -# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object; +# On torch >= 2.11, layer_name is a hoisted LayerName opaque object; # on older versions it remains a plain str. if TYPE_CHECKING: from typing import TypeAlias - _layer_name_type: TypeAlias = str | ModuleName + _layer_name_type: TypeAlias = str | LayerName else: - _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str + _layer_name_type = LayerName if _USE_LAYERNAME else str @torch.compiler.assume_constant_result -def _resolve_layer_name(layer_name: str | ModuleName) -> str: +def _resolve_layer_name(layer_name: str | LayerName) -> str: from torch._library.fake_class_registry import FakeScriptObject - if isinstance(layer_name, ModuleName): + if isinstance(layer_name, LayerName): return layer_name.value elif isinstance(layer_name, FakeScriptObject): return layer_name.real_obj.value @@ -331,9 +334,9 @@ class MoERunnerBase(MoERunner): assert len(trunc_sizes) == 1 return func(states, trunc_sizes[0]) - def _encode_layer_name(self) -> str | ModuleName: - if HAS_OPAQUE_TYPE: - return ModuleName(self.layer_name) + def _encode_layer_name(self) -> str | LayerName: + if _USE_LAYERNAME: + return LayerName(self.layer_name) # Can be unavailable or None in unittests if ( is_forward_context_available() @@ -443,6 +446,19 @@ class MoERunnerBase(MoERunner): if self._shared_experts is not None: self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input) + def _maybe_add_zero_expert_output( + self, + result: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if isinstance(self.router, ZeroExpertRouter): + zero_expert_output = self.router.zero_expert_output + assert zero_expert_output is not None + if isinstance(result, tuple): + result = (result[0], result[1] + zero_expert_output) + else: + result = result + zero_expert_output + return result + def forward( self, hidden_states: torch.Tensor, @@ -494,7 +510,9 @@ class MoERunnerBase(MoERunner): self._encode_layer_name(), ) - return self._maybe_reduce_output(fused_output, og_hidden_dims) + result = self._maybe_reduce_output(fused_output, og_hidden_dims) + + return self._maybe_add_zero_expert_output(result) def forward_dispatch( self, diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py index da5068fa091..2143fa3ce08 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py @@ -12,9 +12,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) -from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import ( - ChunkingMoERunner, -) from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import ( DefaultMoERunner, ) @@ -35,7 +32,7 @@ def create_moe_runner( reduce_results: bool, enable_dbo: bool, ) -> MoERunner: - runner = DefaultMoERunner( + return DefaultMoERunner( layer_name, moe_config, router, @@ -46,6 +43,3 @@ def create_moe_runner( reduce_results, enable_dbo, ) - if moe_config.moe_parallel_config.use_dp_chunking: - return ChunkingMoERunner(runner) - return runner diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index f5b07a6a51a..827a6e6bd3e 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -69,7 +69,6 @@ class SharedExperts: self._moe_config = moe_config self._quant_method = quant_method self._reduce_results = reduce_results - self._use_dp_chunking = moe_config.moe_parallel_config.use_dp_chunking # Allow disabling of the separate shared experts stream for # debug purposes. @@ -87,20 +86,6 @@ class SharedExperts: "Enabled separate cuda stream for MoE shared_experts", scope="local" ) - @property - def _use_external_experts(self) -> bool: - if self._use_dp_chunking: - return False - - # Disable shared expert overlap if: - # - we are using eplb with non-default backend, because of correctness issues - # - we are using flashinfer with DP, since there nothing to gain - backend = self._moe_config.moe_parallel_config.all2all_backend - return ( - self._moe_config.moe_parallel_config.enable_eplb - and backend != "allgather_reducescatter" - ) or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels - def _determine_shared_experts_order( self, hidden_states: torch.Tensor, @@ -110,7 +95,6 @@ class SharedExperts: should_run_shared_in_aux_stream = ( current_platform.is_cuda() - and not self._use_dp_chunking and self._stream is not None and hidden_states.shape[0] <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index c576b0a25c2..ce1e49bc4b0 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -163,8 +163,15 @@ def _int8_quantize( # activations apply per-token quantization. Otherwise, assume # activation tensor-wise fp8/int8 quantization, dynamic or static if block_shape is None: - assert per_act_token, "int8 quantization only supports block or channel-wise" - A, A_scale = per_token_quant_int8(A) + if per_act_token: + A, A_scale = per_token_quant_int8(A) + elif A_scale is not None: + # Static per-tensor: use the optimized CUDA kernel + A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale) + elif A_scale is None: + # Dynamic per-tensor: compute scale then quantize via kernel + A_scale = torch.clamp(A.abs().max() / 127.0, min=1e-10) + A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale) else: assert not per_act_token assert len(block_shape) == 2 diff --git a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py b/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py deleted file mode 100644 index 97d21767f4f..00000000000 --- a/vllm/model_executor/layers/fused_moe/zero_expert_fused_moe.py +++ /dev/null @@ -1,189 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from contextlib import contextmanager - -import torch -from torch import nn - -from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton -from vllm.model_executor.layers.fused_moe.layer import FusedMoE - - -class ZeroExpertFusedMoE(FusedMoE): - """ - A FusedMoE operation that also computes the results of zero experts. - Zero experts perform identity operations (scaled pass-through) instead - of full MLP computations. - - This class uses memoization to avoid redundant routing computation: - routing is computed once and reused for both zero expert computation - and the main FusedMoE forward pass. - """ - - def __init__( - self, - zero_expert_num: int, - zero_expert_type: str, - router: nn.Module, - **kwargs, - ): - # ZeroExpertFusedMoE manages its own custom_routing_function for memoization - assert ( - "custom_routing_function" not in kwargs - or kwargs.get("custom_routing_function") is None - ), ( - "ZeroExpertFusedMoE does not support external custom_routing_function. " - "It manages its own for routing memoization." - ) - - # Automatically slice router's e_score_correction_bias to only include - # real experts (not zero_experts) for the base FusedMoE. - # The full bias will be used temporarily in forward() for routing. - if hasattr(router, "e_score_correction_bias") and "num_experts" in kwargs: - num_real_experts = kwargs["num_experts"] - router_bias = router.e_score_correction_bias - user_bias = kwargs.get("e_score_correction_bias") - - # Use router's bias if: - # 1. User didn't provide bias, or - # 2. User provided full bias (same size as router) - if user_bias is None or user_bias.shape[0] == router_bias.shape[0]: - kwargs["e_score_correction_bias"] = router_bias[:num_real_experts] - - # FusedMoE no longer accepts zero_expert_num/zero_expert_type. - # We handle zero experts ourselves in forward(). - super().__init__(**kwargs) - # Store the actual zero_expert_num and zero_expert_type for our own use - self._actual_zero_expert_num = zero_expert_num - self._actual_zero_expert_type = zero_expert_type - self._router = router # Full router (includes zero experts) - - # Expose zero_expert_num and zero_expert_type as attributes for - # compatibility with quantization methods that check these attributes - self.zero_expert_num = 0 - self.zero_expert_type = None - - # Memoization state for routing results - self._memoized_topk_weights: torch.Tensor | None = None - self._memoized_topk_ids: torch.Tensor | None = None - - # Create custom_routing_function to reuse memoized routing results - def custom_routing_function(hidden_states, gating_output, topk, renormalize): - """Return memoized `topk_weights` and `topk_ids`.""" - if self._memoized_topk_weights is None or self._memoized_topk_ids is None: - raise RuntimeError( - "ZeroExpertFusedMoE: routing results not memoized. " - "Call select_experts first to compute routing." - ) - return self._memoized_topk_weights, self._memoized_topk_ids - - self.custom_routing_function = custom_routing_function - - @contextmanager - def _temporarily_set_attrs(self, **attrs): - """ - Temporarily set attributes using object.__setattr__ and restore them. - - This bypasses nn.Module.__setattr__ to avoid Dynamo tracing issues. - When PyTorch Dynamo traces the forward pass, it cannot handle - nn.Module.__setattr__ calls (which include parameter registration logic), - resulting in "Unsupported" errors. Using object.__setattr__ directly - sets the attribute without triggering nn.Module's custom __setattr__, - allowing Dynamo to trace the code successfully. - """ - originals = {key: getattr(self, key) for key in attrs} - try: - for key, value in attrs.items(): - object.__setattr__(self, key, value) - yield - finally: - for key, value in originals.items(): - object.__setattr__(self, key, value) - - def _compute_zero_expert_result( - self, - hidden_states: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - ) -> torch.Tensor | None: - """Compute zero expert results using pre-computed routing.""" - if ( - self._actual_zero_expert_num is None - or self._actual_zero_expert_num <= 0 - or self._actual_zero_expert_type is None - ): - return None - - return zero_experts_compute_triton( - expert_indices=topk_ids.clone(), - expert_scales=topk_weights.clone(), - num_experts=self.logical_num_experts, - zero_expert_type=self._actual_zero_expert_type, - hidden_states=hidden_states, - ) - - def forward( - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, # Full logits including zero experts - ) -> torch.Tensor: - """ - Forward pass with zero expert support and routing memoization. - - Args: - hidden_states: Input hidden states - router_logits: Full router logits (including zero experts) - - Returns: - Combined output from real experts and zero experts - """ - # Prepare temporary attribute overrides for routing computation - temp_attrs = { - "custom_routing_function": None, # Disable for first routing - } - if self._router is not None: - temp_attrs["e_score_correction_bias"] = self._router.e_score_correction_bias - - # Compute routing with temporary attributes - # Pass full router_logits (including zero experts) so that zero experts - # can be properly identified in topk_ids - with self._temporarily_set_attrs(**temp_attrs): - topk_weights, topk_ids = self.select_experts( - hidden_states=hidden_states, - router_logits=router_logits, # Full logits (includes zero experts) - ) - - # Compute zero expert result if needed - zero_expert_result = self._compute_zero_expert_result( - hidden_states=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - ) - - # Memoize routing results for reuse in super().forward() - self._memoized_topk_weights = topk_weights - self._memoized_topk_ids = topk_ids - - # Slice router_logits for real experts only - router_logits_sliced = router_logits[..., : self.logical_num_experts] - - # Compute real expert results (will reuse memoized routing via - # custom_routing_function) - # zero_expert_num is already 0, so FusedMoE won't handle zero experts - fused_out = super().forward( - hidden_states=hidden_states, - router_logits=router_logits_sliced, - ) - - # Combine results - # Both zero_expert_result and fused_out are computed from the same - # hidden_states, so they should be on the same device. - if zero_expert_result is not None: - fused_out = fused_out + zero_expert_result - - # Clear memoization after use - self._memoized_topk_weights = None - self._memoized_topk_ids = None - - return fused_out diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index dd2a61bc6a2..3541b970668 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -9,14 +9,14 @@ from vllm.distributed import ( tensor_model_parallel_all_gather, tensor_model_parallel_gather, ) -from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.platforms import current_platform # --8<-- [start:logits_processor] -@CustomOp.register("logits_processor") -class LogitsProcessor(CustomOp): +@PluggableLayer.register("logits_processor") +class LogitsProcessor(PluggableLayer): """Process logits and apply logits processors from sampling metadata. This layer does the following: diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index aec855d9aeb..3d875683d26 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -56,7 +56,12 @@ from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.triton_utils import tl, triton -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata @@ -568,7 +573,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): b, a, core_attn_out, - self.prefix, + _encode_layer_name(self.prefix), ) # ============================================================ @@ -702,19 +707,33 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): num_v_heads = self.num_v_heads // self.tp_size _, state_dtype = self.get_state_dtype() - # All kernels use BT = chunk_size (FLA_CHUNK_SIZE4), so a single pass with - # T = chunk_size is sufficient to populate every autotuner cache. + # All kernels use BT = chunk_size, so a single pass with T = chunk_size + # is sufficient to populate every autotuner cache. Mirror the real + # prefill path here: build q/k/v/g/beta via fused_post_conv_prep and + # then run chunk_gated_delta_rule with in-kernel L2 norm disabled. T = FLA_CHUNK_SIZE - q = torch.randn(1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype) - k = torch.randn(1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype) - v = torch.randn(1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype) - # NOTE: g and beta must have the same dtypes as during - # inference, so we construct them with the same function - # (fused_gdn_gating). dummy_a and dummy_b are throwaway - # inputs required by that function. + dummy_mixed_qkv = torch.randn( + T, mixed_qkv.shape[-1], device=device, dtype=dtype + ) dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype) dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype) - g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias) + q, k, v, g, beta = fused_post_conv_prep( + conv_output=dummy_mixed_qkv, + a=dummy_a, + b=dummy_b, + A_log=self.A_log, + dt_bias=self.dt_bias, + num_k_heads=num_k_heads, + head_k_dim=self.head_k_dim, + head_v_dim=self.head_v_dim, + apply_l2norm=True, + output_g_exp=False, + ) + q = q.unsqueeze(0) + k = k.unsqueeze(0) + v = v.unsqueeze(0) + g = g.unsqueeze(0) + beta = beta.unsqueeze(0) state = torch.zeros( 1, num_v_heads, @@ -735,7 +754,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): initial_state=state, output_final_state=True, cu_seqlens=cu_seqlens, - use_qk_l2norm_in_kernel=True, + use_qk_l2norm_in_kernel=False, ) except Exception: logger.warning( @@ -753,7 +772,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): self.prefix, ) finally: - del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens + del dummy_mixed_qkv, q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens torch.accelerator.empty_cache() @@ -1070,13 +1089,14 @@ def gdn_attention_core( b: torch.Tensor, a: torch.Tensor, core_attn_out: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: """ Custom op for the core attention computation. Only handles the convolution + recurrent attention part. Input/output projections are handled outside this op. """ + layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] self._forward_core( @@ -1092,7 +1112,7 @@ def gdn_attention_core_fake( b: torch.Tensor, a: torch.Tensor, core_attn_out: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: """Fake implementation for torch.compile.""" return diff --git a/vllm/model_executor/layers/mamba/lamport_workspace.py b/vllm/model_executor/layers/mamba/lamport_workspace.py new file mode 100644 index 00000000000..afae19c75fe --- /dev/null +++ b/vllm/model_executor/layers/mamba/lamport_workspace.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import array +import contextlib +import struct +import sys +import threading + +import torch + +try: + from cuda.bindings import runtime as cudart +except ImportError: + from cuda import cudart + +_ALIGN = 1 << 21 # 2 MiB — CUDA IPC allocation alignment + + +# --------------------------------------------------------------------------- +# CUDA helpers +# --------------------------------------------------------------------------- + + +def _check(error): + """Raise on CUDA runtime error.""" + success = getattr(cudart.cudaError_t, "cudaSuccess", None) or cudart.cudaError_t(0) + if error != success: + raise RuntimeError(f"CUDA runtime error: {error}") + + +def _cuda_malloc(size: int): + aligned = ((size + _ALIGN - 1) >> 21) << 21 + err, ptr = cudart.cudaMalloc(aligned) + _check(err) + return ptr, aligned + + +def _cuda_free(ptr: int): + if ptr: + _check(cudart.cudaFree(ptr)[0]) + + +def _cuda_memset_zero(ptr: int, size: int): + _check(cudart.cudaMemset(ptr, 0, size)[0]) + + +def _cuda_memcpy_d2d(dst: int, src: int, size: int): + _check( + cudart.cudaMemcpy( + dst, src, size, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice + )[0] + ) + + +# --------------------------------------------------------------------------- +# IPC buffer +# --------------------------------------------------------------------------- + + +class IpcBuffer: + """ + Allocates CUDA device memory and exchanges IPC handles with all ranks + so that every rank holds a valid device pointer to every other rank's buffer. + """ + + def __init__(self, rank: int, world_size: int, size: int, process_group=None): + self.rank = rank + self.world_size = world_size + self.peer_ptrs: list[int] = [0] * world_size + self.local_ptr: int = 0 + self._alive = False + + if size <= 0: + return + + self.local_ptr, _ = _cuda_malloc(size) + _cuda_memset_zero(self.local_ptr, size) + self._alive = True + + # --- exchange IPC handles via torch.distributed --- + err, local_handle = cudart.cudaIpcGetMemHandle(self.local_ptr) + _check(err) + + all_handles: list[bytes | None] = [None] * world_size + torch.distributed.all_gather_object( + all_handles, bytes(local_handle.reserved), group=process_group + ) + + for r in range(world_size): + if r == rank: + self.peer_ptrs[r] = self.local_ptr + else: + handle = cudart.cudaIpcMemHandle_t() + handle.reserved = all_handles[r] + err, ptr = cudart.cudaIpcOpenMemHandle( + handle, cudart.cudaIpcMemLazyEnablePeerAccess + ) + _check(err) + self.peer_ptrs[r] = ptr + + def serialize(self) -> list[int]: + """Return peer pointers as a list of int64 values (one per rank).""" + raw = b"" + for ptr in self.peer_ptrs: + raw += struct.pack("P", ptr) + return array.array("Q", raw).tolist() + + def cleanup(self): + if not self._alive: + return + self._alive = False + for r in range(self.world_size): + if self.peer_ptrs[r] == 0: + continue + if r == self.rank: + _cuda_free(self.peer_ptrs[r]) + else: + with contextlib.suppress(RuntimeError): + _check(cudart.cudaIpcCloseMemHandle(self.peer_ptrs[r])[0]) + self.peer_ptrs[r] = 0 + self.local_ptr = 0 + + def __del__(self): + if not sys.is_finalizing(): + self.cleanup() + + +# --------------------------------------------------------------------------- +# Lamport negative-zero initialization +# --------------------------------------------------------------------------- + + +def _lamport_fill_neg_zero(device_ptr: int, size_bytes: int): + """ + Fill device memory with IEEE-754 negative zero (-0.0f = 0x80000000). + This is the "slot empty" sentinel for the Lamport protocol: the kernel + spin-waits until a value is *not* negative zero. + """ + if size_bytes == 0 or device_ptr == 0: + return + n_floats = size_bytes // 4 + # torch preserves -0.0 in IEEE-754 + fill = torch.full((n_floats,), -0.0, dtype=torch.float32, device="cuda") + _cuda_memcpy_d2d(device_ptr, fill.data_ptr(), size_bytes) + del fill + + +# --------------------------------------------------------------------------- +# LamportWorkspace — the main class +# --------------------------------------------------------------------------- + + +class LamportWorkspace: + """ + Self-contained workspace for Lamport-based cross-GPU AllReduce. + + Parameters + ---------- + rank : int + Local rank (0-based). + world_size : int + Total number of ranks in the TP group. + comm_size : int + Size in bytes of *one* Lamport buffer slot. The total IPC allocation + per rank is ``3 * comm_size`` (triple-buffering). Must be large enough + to hold the per-slot data written by the kernel. Use + ``compute_comm_size_for_minimax()`` for a safe default. + process_group : optional + ``torch.distributed`` process group for IPC handle exchange. + ``None`` uses the default group. + """ + + def __init__(self, rank: int, world_size: int, comm_size: int, process_group=None): + assert world_size >= 2, "Lamport workspace requires at least 2 ranks" + assert comm_size > 0, "comm_size must be positive" + + self.rank = rank + self.world_size = world_size + self.comm_size = comm_size + + # 1) Lamport triple-buffer (the only IPC memory the kernel reads/writes) + lamport_total = 3 * comm_size + self._lamport = IpcBuffer(rank, world_size, lamport_total, process_group) + _lamport_fill_neg_zero(self._lamport.local_ptr, lamport_total) + + # 2) flag_buffer on device: int32[3] = {counter, unused, lamport_flag} + # counter — used for block-level sync inside the kernel + # unused — reserved (index 1) + # lamport_flag — triple-buffer rotation index (0 → 1 → 2 → 0 …) + self._flag_buf = torch.zeros(3, dtype=torch.int32, device="cuda") + + # 3) layout_buffer on device: int64[2] = {clear_size, comm_size} + # clear_size — bytes to clear from *previous* slot (set by kernel) + # comm_size — size of one triple-buffer slot + self._layout_buf = torch.tensor( + [0, comm_size], dtype=torch.int64, device="cuda" + ) + + # 4) Assemble device-side void* pointer array + N = world_size + ptrs: list[int] = [] + ptrs += [0] * N # [0 .. N-1] ipc_buffers (placeholder) + ptrs += [0] * N # [N .. 2N-1] ipc_barriers (placeholder) + ptrs += self._lamport.serialize() # [2N .. 3N-1] lamport peer ptrs + ptrs.append(self._flag_buf.data_ptr()) # [3N] flag_buffer + ptrs.append(self._layout_buf.data_ptr()) # [3N+1] layout_buffer + + self._workspace = torch.tensor(ptrs, dtype=torch.int64, device="cuda") + + @property + def workspace(self) -> torch.Tensor: + """Device tensor (int64) that can be passed to the kernel + as ``void** workspace``.""" + return self._workspace + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def compute_comm_size_for_minimax( + max_tokens: int, + world_size: int, + fused_qk: bool = True, + ) -> int: + """ + Return a safe ``comm_size`` (in bytes) for MiniMaxReduceRMSKernel. + + The kernel stores per-token variance scalars in the Lamport buffer: + - single-matrix path: ``world_size × max_tokens × 4`` bytes per slot + - fused Q+K path: ``world_size × 2 × ceil(max_tokens/4) × 16`` bytes per slot + + The returned value is rounded up to 2 MiB alignment. + """ + if fused_qk: + groups = (max_tokens + 3) // 4 + slot_bytes = world_size * 2 * groups * 16 # 16 = sizeof(float4) + else: + slot_bytes = world_size * max_tokens * 4 # 4 = sizeof(float) + return ((slot_bytes + _ALIGN - 1) >> 21) << 21 + + def cleanup(self): + if hasattr(self, "_lamport"): + self._lamport.cleanup() + + def __del__(self): + if not sys.is_finalizing(): + self.cleanup() + + def __repr__(self): + return ( + f"LamportWorkspace(rank={self.rank}, world_size={self.world_size}, " + f"comm_size={self.comm_size})" + ) + + +# --------------------------------------------------------------------------- +# Cached convenience function (mirrors TRT-LLM's get_allreduce_workspace) +# --------------------------------------------------------------------------- + +_cache_lock = threading.Lock() +_workspace_cache: dict = {} + + +def get_allreduce_workspace( + rank: int, + world_size: int, + comm_size: int | None = None, + max_tokens: int = 16384, + process_group=None, +) -> torch.Tensor: + """ + Return a cached workspace tensor for the given (rank, world_size) pair. + + On first call the workspace is allocated and IPC handles are exchanged; + subsequent calls with the same arguments return the cached tensor. + + Parameters + ---------- + rank, world_size : int + TP rank and TP size. + comm_size : int, optional + Explicit slot size in bytes. If ``None``, computed automatically + from ``max_tokens`` and ``world_size`` (fused Q+K path). + max_tokens : int + Maximum number of tokens per batch (used when ``comm_size is None``). + process_group : optional + ``torch.distributed`` process group. + """ + if comm_size is None: + comm_size = LamportWorkspace.compute_comm_size_for_minimax( + max_tokens, world_size, fused_qk=True + ) + pg_id = id(process_group) if process_group is not None else 0 + key = (rank, world_size, comm_size, pg_id) + with _cache_lock: + if key not in _workspace_cache: + ws = LamportWorkspace(rank, world_size, comm_size, process_group) + _workspace_cache[key] = ws + return _workspace_cache[key].workspace diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index fd83d4b8322..4509a095628 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -30,13 +30,16 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) -from vllm.model_executor.layers.mamba.ops.mamba_ssm import ( - selective_scan_fn, - selective_state_update, -) +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_scan_fn +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata @@ -228,7 +231,7 @@ class MambaMixer(MambaBase, PluggableLayer): torch.ops.vllm.mamba_mixer( hidden_states, output, - self.prefix, + _encode_layer_name(self.prefix), ) def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): @@ -426,14 +429,12 @@ class MambaMixer(MambaBase, PluggableLayer): B_d, C_d, self.D, - gate_d.transpose(0, 1), time_proj_bias, + z=gate_d.transpose(0, 1), dt_softplus=True, state_batch_indices=state_indices_tensor_d_input, dst_state_batch_indices=state_indices_tensor_d_output, out=scan_outputs_d, - enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding, - cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds, ) scan_outputs_d = scan_outputs_d.transpose(0, 1) @@ -515,8 +516,9 @@ def split_batch_to_prefill_and_decode( def mamba_mixer( hidden_states: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: + layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] self.forward_impl(hidden_states=hidden_states, output=output) @@ -525,7 +527,7 @@ def mamba_mixer( def mamba_mixer_fake( hidden_states: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: return diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 01ea3fdca57..0518bde2f42 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -31,10 +31,10 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_update, ) from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated -from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import ( LoaderFunction, @@ -44,7 +44,12 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata @@ -536,7 +541,7 @@ class MambaMixer2(MambaBase, PluggableLayer): torch.ops.vllm.mamba_mixer2( projected_states, ssm_output, - self.prefix, + _encode_layer_name(self.prefix), ) # 4. gated MLP @@ -885,8 +890,7 @@ class MambaMixer2(MambaBase, PluggableLayer): B_d, C_d, D_d, - z=None, - dt_bias=dt_bias, + dt_bias, dt_softplus=True, state_batch_indices=state_indices_tensor_d_input, dst_state_batch_indices=state_indices_tensor_d_output, @@ -894,8 +898,6 @@ class MambaMixer2(MambaBase, PluggableLayer): num_accepted_tokens=num_accepted_tokens, cu_seqlens=query_start_loc_d, is_blackwell=self.is_blackwell, - enable_stochastic_rounding=self.cache_config.enable_mamba_cache_stochastic_rounding, - cache_philox_rounds=self.cache_config.mamba_cache_philox_rounds, ) def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: @@ -927,8 +929,9 @@ class MambaMixer2(MambaBase, PluggableLayer): def mamba_mixer2( projected_states: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: + layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] self.conv_ssm_forward(projected_states=projected_states, output=output) @@ -937,7 +940,7 @@ def mamba_mixer2( def mamba_mixer2_fake( projected_states: torch.Tensor, output: torch.Tensor, - layer_name: str, + layer_name: LayerNameType, ) -> None: return diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index c4a0ef385d1..e3c8ba8312f 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -323,9 +323,9 @@ def selective_state_update( A, B, C, - D=None, + D, + dt_bias, z=None, - dt_bias=None, dt_softplus=False, state_batch_indices=None, dst_state_batch_indices=None, @@ -374,11 +374,11 @@ def selective_state_update( B = B.unsqueeze(1) if C.dim() == 2: C = C.unsqueeze(1) - if D is not None and D.dim() == 1: + if D.dim() == 1: D = D.unsqueeze(0) if z is not None and z.dim() == 2: z = z.unsqueeze(1) - if dt_bias is not None and dt_bias.dim() == 1: + if dt_bias.dim() == 1: dt_bias = dt_bias.unsqueeze(0) if out.dim() == 2: out = out.unsqueeze(1) @@ -410,12 +410,10 @@ def selective_state_update( assert nheads % ngroups == 0, "nheads must be divisible by ngroups" assert B.shape == (batch, ngroups, dstate) assert C.shape == B.shape - if D is not None: - assert D.shape == (nheads, dim) + assert D.shape == (nheads, dim) if z is not None: assert z.shape == x.shape - if dt_bias is not None: - assert dt_bias.shape == (nheads, dim) + assert dt_bias.shape == (nheads, dim) if state_batch_indices is not None: assert state_batch_indices.shape[0] >= N assert state_batch_indices.shape[1] >= max_seqlen @@ -506,7 +504,8 @@ def selective_state_update( dt.stride(0), dt.stride(1), dt.stride(2), - *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else 0, + dt_bias.stride(0), + dt_bias.stride(1), A.stride(0), A.stride(1), A.stride(2), @@ -516,7 +515,8 @@ def selective_state_update( C.stride(0), C.stride(1), C.stride(2), - *(D.stride(0), D.stride(1)) if D is not None else 0, + D.stride(0), + D.stride(1), z_strides[0], z_strides[1], z_strides[2], diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py new file mode 100644 index 00000000000..8a86b1a068b --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Dispatch module for Mamba selective state update (SSU) backends. + +Provides a unified `selective_state_update` function that dispatches to +either the Triton or FlashInfer backend based on the configured +`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM. +""" + +from abc import ABC, abstractmethod + +import torch + +from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.logger import init_logger +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +logger = init_logger(__name__) + + +class MambaSSUBackend(ABC): + """Abstract base class for Mamba SSU backends.""" + + def __init__(self, mamba_config: MambaConfig): + self._mamba_config = mamba_config + + @property + @abstractmethod + def name(self) -> str: ... + + @abstractmethod + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + dst_state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + is_blackwell: bool = False, + ) -> None: ... + + +class TritonSSUBackend(MambaSSUBackend): + """Triton-based SSU backend (vLLM's default).""" + + def __init__(self, mamba_config: MambaConfig): + super().__init__(mamba_config) + from vllm.model_executor.layers.mamba.ops.mamba_ssm import ( + selective_state_update as _triton_selective_state_update, + ) + + self._kernel = _triton_selective_state_update + + @property + def name(self) -> str: + return "triton" + + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + dst_state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + is_blackwell: bool = False, + ) -> None: + self._kernel( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + null_block_id=null_block_id, + out=out, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=cu_seqlens, + is_blackwell=is_blackwell, + enable_stochastic_rounding=self._mamba_config.enable_stochastic_rounding, + cache_philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds, + ) + + +class FlashInferSSUBackend(MambaSSUBackend): + """FlashInfer-based SSU backend.""" + + def __init__(self, mamba_config: MambaConfig): + super().__init__(mamba_config) + try: + from flashinfer.mamba import selective_state_update as _fi_ssu + except ImportError as e: + raise ImportError( + "FlashInfer is required for the flashinfer Mamba SSU backend. " + "Please install flashinfer (>= 0.6.4): " + "pip install flashinfer-python" + ) from e + self._kernel = _fi_ssu + + @property + def name(self) -> str: + return "flashinfer" + + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + dst_state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + is_blackwell: bool = False, + ) -> None: + rand_seed = ( + torch.randint(0, 2**32, (1,), device=state.device) + if self._mamba_config.enable_stochastic_rounding + else None + ) + + self._kernel( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=num_accepted_tokens, + cache_steps=state_batch_indices.size(-1) + if cu_seqlens is not None and state_batch_indices is not None + else 0, + pad_slot_id=null_block_id, + out=out, + rand_seed=rand_seed, + philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10, + ) + + +_BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { + MambaBackendEnum.TRITON: TritonSSUBackend, + MambaBackendEnum.FLASHINFER: FlashInferSSUBackend, +} + +_mamba_ssu_backend: MambaSSUBackend | None = None + + +def initialize_mamba_ssu_backend(mamba_config: MambaConfig) -> None: + """Initialize the global Mamba SSU backend. + + Args: + mamba_config: Mamba configuration. + """ + global _mamba_ssu_backend + + backend = mamba_config.backend + if backend not in _BACKEND_REGISTRY: + raise ValueError( + f"Unknown Mamba SSU backend: {backend}. " + f"Valid options: {list(_BACKEND_REGISTRY.keys())}" + ) + + _mamba_ssu_backend = _BACKEND_REGISTRY[backend](mamba_config) + logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name) + + +def get_mamba_ssu_backend() -> MambaSSUBackend: + """Get the current Mamba SSU backend. Raises if not initialized.""" + if _mamba_ssu_backend is None: + raise RuntimeError( + "Mamba SSU backend has not been initialized. " + "Call initialize_mamba_ssu_backend() first." + ) + return _mamba_ssu_backend + + +def selective_state_update( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + dst_state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + is_blackwell: bool = False, +) -> None: + """Unified dispatch for Mamba selective state update. + + Delegates to the initialized backend (Triton or FlashInfer). + """ + get_mamba_ssu_backend()( + state, + x, + dt, + A, + B, + C, + D, + dt_bias, + z=z, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + null_block_id=null_block_id, + out=out, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=cu_seqlens, + is_blackwell=is_blackwell, + ) diff --git a/vllm/model_executor/layers/pooler/seqwise/heads.py b/vllm/model_executor/layers/pooler/seqwise/heads.py index 31a96122392..2424d841075 100644 --- a/vllm/model_executor/layers/pooler/seqwise/heads.py +++ b/vllm/model_executor/layers/pooler/seqwise/heads.py @@ -103,14 +103,16 @@ class ClassifierPoolerHead(SequencePoolerHead): def __init__( self, classifier: ClassifierFn | None = None, - logit_bias: float | None = None, + logit_mean: float | None = None, + logit_sigma: float | None = None, head_dtype: torch.dtype | str | None = None, activation: ActivationFn | None = None, ) -> None: super().__init__() self.classifier = classifier - self.logit_bias = logit_bias + self.logit_mean = logit_mean + self.logit_sigma = logit_sigma self.head_dtype = head_dtype self.activation = activation @@ -138,8 +140,11 @@ class ClassifierPoolerHead(SequencePoolerHead): logits = pooled_data # logits shape: [batchsize, num_labels] - if self.logit_bias is not None: - logits -= self.logit_bias + # Affine score calibration: activation((logit - mean) / sigma) + if self.logit_mean is not None: + logits = logits - self.logit_mean + if self.logit_sigma is not None: + logits = logits / self.logit_sigma if self.activation is not None: flags = [p.use_activation for p in pooling_params] diff --git a/vllm/model_executor/layers/pooler/seqwise/poolers.py b/vllm/model_executor/layers/pooler/seqwise/poolers.py index f46834a7c3f..74fa4cdbbe4 100644 --- a/vllm/model_executor/layers/pooler/seqwise/poolers.py +++ b/vllm/model_executor/layers/pooler/seqwise/poolers.py @@ -118,7 +118,8 @@ def pooler_for_classify( head = ClassifierPoolerHead( head_dtype=model_config.head_dtype, classifier=classifier, - logit_bias=model_config.pooler_config.logit_bias, + logit_mean=model_config.pooler_config.logit_mean, + logit_sigma=model_config.pooler_config.logit_sigma, activation=resolve_classifier_act_fn( model_config, static_num_labels=True, act_fn=act_fn ), diff --git a/vllm/model_executor/layers/pooler/tokwise/heads.py b/vllm/model_executor/layers/pooler/tokwise/heads.py index 80c5c831fa0..0377a86755a 100644 --- a/vllm/model_executor/layers/pooler/tokwise/heads.py +++ b/vllm/model_executor/layers/pooler/tokwise/heads.py @@ -92,14 +92,16 @@ class TokenClassifierPoolerHead(TokenPoolerHead): def __init__( self, classifier: ClassifierFn | None = None, - logit_bias: float | None = None, + logit_mean: float | None = None, + logit_sigma: float | None = None, head_dtype: torch.dtype | str | None = None, activation: ActivationFn | None = None, ) -> None: super().__init__() self.classifier = classifier - self.logit_bias = logit_bias + self.logit_mean = logit_mean + self.logit_sigma = logit_sigma self.head_dtype = head_dtype self.activation = activation @@ -125,8 +127,11 @@ class TokenClassifierPoolerHead(TokenPoolerHead): logits = pooled_data # logits shape: [n_token, num_labels] - if self.logit_bias is not None: - logits -= self.logit_bias + # Affine score calibration: activation((logit - mean) / sigma) + if self.logit_mean is not None: + logits = logits - self.logit_mean + if self.logit_sigma is not None: + logits = logits / self.logit_sigma if self.activation is not None and pooling_param.use_activation: logits = self.activation(logits) diff --git a/vllm/model_executor/layers/pooler/tokwise/methods.py b/vllm/model_executor/layers/pooler/tokwise/methods.py index f242d215d7b..9ee6e8527c9 100644 --- a/vllm/model_executor/layers/pooler/tokwise/methods.py +++ b/vllm/model_executor/layers/pooler/tokwise/methods.py @@ -100,7 +100,7 @@ class StepPool(AllPool): ): # for unfinished chunked prefill if data is None: - pass + pooled_data.append(None) else: step_tag_id = pooling_param.step_tag_id returned_token_ids = pooling_param.returned_token_ids diff --git a/vllm/model_executor/layers/pooler/tokwise/poolers.py b/vllm/model_executor/layers/pooler/tokwise/poolers.py index c56970fcaba..6462a5056c5 100644 --- a/vllm/model_executor/layers/pooler/tokwise/poolers.py +++ b/vllm/model_executor/layers/pooler/tokwise/poolers.py @@ -58,7 +58,7 @@ class TokenPooler(Pooler): def __init__( self, pooling: TokenPoolingMethod | TokenPoolingFn, - head: TokenPoolerHead | TokenPoolingHeadFn, + head: TokenPoolerHead | TokenPoolingHeadFn | None = None, ) -> None: super().__init__() @@ -89,7 +89,8 @@ class TokenPooler(Pooler): pooling_metadata: PoolingMetadata, ) -> TokenPoolerOutput: pooled_data = self.pooling(hidden_states, pooling_metadata) - pooled_data = self.head(pooled_data, pooling_metadata) + if self.head is not None: + pooled_data = self.head(pooled_data, pooling_metadata) return pooled_data @@ -126,7 +127,8 @@ def pooler_for_token_classify( head = TokenClassifierPoolerHead( head_dtype=model_config.head_dtype, classifier=classifier, - logit_bias=model_config.pooler_config.logit_bias, + logit_mean=model_config.pooler_config.logit_mean, + logit_sigma=model_config.pooler_config.logit_sigma, activation=resolve_classifier_act_fn( model_config, static_num_labels=False, act_fn=act_fn ), diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 1ac0f9ee9cc..6313db78a82 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -30,6 +30,7 @@ QuantizationMethods = Literal[ "torchao", "inc", "mxfp4", + "gpt_oss_mxfp4", "mxfp8", "cpu_awq", "online", @@ -133,7 +134,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: ModelOptNvFp4Config, ) from .moe_wna16 import MoeWNA16Config - from .mxfp4 import Mxfp4Config + from .mxfp4 import GptOssMxfp4Config, Mxfp4Config from .mxfp8 import Mxfp8Config from .online.base import OnlineQuantizationConfig from .torchao import TorchAOConfig @@ -160,6 +161,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "auto-round": INCConfig, "inc": INCConfig, "mxfp4": Mxfp4Config, + "gpt_oss_mxfp4": GptOssMxfp4Config, "mxfp8": Mxfp8Config, "cpu_awq": CPUAWQConfig, "online": OnlineQuantizationConfig, diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index be3001a7fa1..cfad1f86faa 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -232,7 +232,7 @@ class AWQMarlinConfig(QuantizationConfig): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> "QuantizationMethods | None": # Skip override to marlin kernels, as they are not # batch invariant diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index eedc62f7d4d..344ddd8abd2 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -110,13 +110,22 @@ class QuantizationConfig(ABC): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, + hf_quant_cfg: dict[str, Any], + user_quant: str | None, + hf_config: Any = None, ) -> QuantizationMethods | None: """ Detects if this quantization method can support a given checkpoint format by overriding the user specified quantization method -- this method should only be overwritten by subclasses in exceptional - circumstances + circumstances. + + Args: + hf_quant_cfg: The checkpoint's quantization config dict. + user_quant: The user-specified quantization method string. + hf_config: The HuggingFace model config object (e.g. for + model_type checks). May be None if not available. """ return None diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 6ca65cdb188..8d16a143b10 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -49,6 +49,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsW4A16Mxfp4, CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8, + CompressedTensorsW8A8Mxfp8, CompressedTensorsW8A16Fp8, CompressedTensorsWNA16, ) @@ -403,6 +404,27 @@ class CompressedTensorsConfig(QuantizationConfig): and is_symmetric ) + @staticmethod + def _is_mxfp8(quant_args: QuantizationArgs) -> bool: + if quant_args is None: + return False + + is_group_quant = quant_args.strategy == QuantizationStrategy.GROUP.value + is_symmetric = quant_args.symmetric + is_group_size_32 = quant_args.group_size == 32 + is_float_type = quant_args.type == QuantizationType.FLOAT + is_8_bits = quant_args.num_bits == 8 + is_mxfp8_scale_dtype = quant_args.scale_dtype == torch.uint8 + + return ( + is_group_quant + and is_float_type + and is_8_bits + and is_group_size_32 + and is_symmetric + and is_mxfp8_scale_dtype + ) + @staticmethod def _is_static_tensor_w8a8( weight_quant: QuantizationArgs, input_quant: QuantizationArgs @@ -606,6 +628,9 @@ class CompressedTensorsConfig(QuantizationConfig): if self._is_mxfp4(weight_quant): return CompressedTensorsW4A16Mxfp4() + if self._is_mxfp8(weight_quant): + return CompressedTensorsW8A8Mxfp8() + if self._is_fp8_w4a8_sm90(weight_quant, input_quant): return CompressedTensorsW4A8Fp8( num_bits=weight_quant.num_bits, @@ -1098,6 +1123,17 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod): layer._v_scale = layer.v_scale layer._q_scale = layer.q_scale + # Set the _float variants that the attention backend uses. + def _to_scalar(tensor: torch.Tensor) -> float: + # For n_scales > 1 (e.g., ATTN_HEAD strategy), take max + if tensor.numel() > 1: + return tensor.max().item() + return tensor.item() + + layer._k_scale_float = _to_scalar(layer.k_scale) + layer._v_scale_float = _to_scalar(layer.v_scale) + layer._q_scale_float = _to_scalar(layer.q_scale) + # Discard all placeholders. del layer.k_scale del layer.v_scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 9ee8df9daba..f25b8af1d6b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -68,6 +68,13 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase): return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config) + if quant_config._is_mxfp8(weight_quant): + from .compressed_tensors_moe_w8a8_mxfp8 import ( + CompressedTensorsW8A8Mxfp8MoEMethod, + ) + + return CompressedTensorsW8A8Mxfp8MoEMethod(layer.moe_config) + if quant_config._is_wNa16_group_channel(weight_quant, input_quant): # group_size=None means channelwise group_size = weight_quant.group_size or -1 diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py new file mode 100644 index 00000000000..02e946b1b61 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoeWeightScaleSupported, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + make_fp8_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import ( # noqa: E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs + + +class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): + """Compressed-tensors MoE method for pre-quantized MXFP8 (W8A8) checkpoints. + + Loads FP8 (E4M3) weights with E8M0 uint8 per-group scales (group_size=32) + from checkpoint. Activations are dynamically quantized to MXFP8 at runtime. + Supports FlashInfer TRT-LLM and Marlin backends (auto-selected). + """ + + def __init__(self, moe: FusedMoEConfig): + super().__init__(moe) + self.weight_block_size = [1, MXFP8_BLOCK_SIZE] + self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.params_dtype = params_dtype + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + hidden_size, + dtype=MXFP8_VALUE_DTYPE, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=MXFP8_VALUE_DTYPE, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + w13_num_shards * intermediate_size_per_partition, + hidden_size // MXFP8_BLOCK_SIZE, + dtype=MXFP8_SCALE_DTYPE, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.GROUP.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + w2_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // MXFP8_BLOCK_SIZE, + dtype=MXFP8_SCALE_DTYPE, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + layer.w13_input_scale = None + layer.w2_input_scale = None + + def process_weights_after_loading(self, layer: FusedMoE) -> None: + layer.weight_block_size = self.weight_block_size + + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=layer.w13_weight, + w2=layer.w2_weight, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_input_scale=layer.w13_input_scale, + w2_input_scale=layer.w2_input_scale, + ) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + if self.moe_quant_config is not None: + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + return make_fp8_moe_quant_config( + fp8_backend=self.fp8_backend, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + block_shape=self.weight_block_size, + ) + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel " + "initialization logic. This function should not be called." + ) + + def apply_monolithic( + self, + layer: FusedMoE, + x: torch.Tensor, + router_logits: torch.Tensor, + ) -> torch.Tensor: + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py index c9dd98dfd4e..457794eb0a0 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py @@ -9,6 +9,7 @@ from .compressed_tensors_w4a16_mxfp4 import CompressedTensorsW4A16Mxfp4 from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4 from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8 from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8 +from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8 from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8 from .compressed_tensors_wNa16 import WNA16_SUPPORTED_BITS, CompressedTensorsWNA16 @@ -28,4 +29,5 @@ __all__ = [ "CompressedTensorsW4A4Fp4", "CompressedTensorsW4A8Int", "CompressedTensorsW4A8Fp8", + "CompressedTensorsW8A8Mxfp8", ] diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py index fff7387260e..c818f334589 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py @@ -6,15 +6,10 @@ import torch from torch.nn.parameter import Parameter from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) -from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( - NvFp4LinearBackend, - apply_nvfp4_linear, - convert_to_nvfp4_linear_kernel_format, - select_nvfp4_linear_backend, -) from vllm.model_executor.parameter import ( GroupQuantScaleParameter, ModelWeightParameter, @@ -29,13 +24,9 @@ __all__ = ["CompressedTensorsW4A4Fp4"] class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): def __init__(self): - self.backend = select_nvfp4_linear_backend() + self.kernel = init_nvfp4_linear_kernel() self.group_size = 16 - self.swizzle = None - if self.backend == NvFp4LinearBackend.EMULATION: - self.swizzle = False - @classmethod def get_min_capability(cls) -> int: return 75 @@ -130,7 +121,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): ) # Convert layer to NVFP4 linear kernel format - convert_to_nvfp4_linear_kernel_format(self.backend, layer) + self.kernel.process_weights_after_loading(layer) def apply_weights( self, @@ -138,10 +129,4 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - return apply_nvfp4_linear( - backend=self.backend, - layer=layer, - x=x, - bias=bias, - swizzle=self.swizzle, - ) + return self.kernel.apply_weights(layer=layer, x=x, bias=bias) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py index 7bffc3218b4..42b35a420ca 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py @@ -6,45 +6,49 @@ from collections.abc import Callable import torch from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy +from vllm.config import get_current_vllm_config +from vllm.model_executor.kernels.linear import ( + init_wfp8_a16_linear_kernel, +) from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) +from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( + STRATEGY_TO_PARAMETER_TYPE, + STRATEGY_TO_WEIGHT_QUANT_KEY, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_scale_parameter, create_fp8_weight_parameter, - process_fp8_weight_block_strategy, validate_fp8_block_shape, ) -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( - apply_fp8_marlin_linear, - prepare_fp8_layer_for_marlin, +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8DynamicTensorSym, + kFp8StaticTensorSym, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( convert_to_channelwise, ) -from vllm.model_executor.parameter import ( - BlockQuantScaleParameter, - ChannelQuantScaleParameter, - PerTensorScaleParameter, -) +from vllm.model_executor.parameter import PerTensorScaleParameter from vllm.model_executor.utils import replace_parameter __all__ = ["CompressedTensorsW8A16Fp8"] -strategy_to_parameter_type = { - QuantizationStrategy.BLOCK: BlockQuantScaleParameter, - QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter, - QuantizationStrategy.TENSOR: PerTensorScaleParameter, -} - class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool): self.weight_quant = weight_quant self.strategy = weight_quant.strategy + self.out_dtype = torch.get_default_dtype() + self.input_dtype = get_current_vllm_config().model_config.dtype self.is_static_input_scheme = is_static_input_scheme self.weight_block_size = self.weight_quant.block_structure + self.weight_quant_key = STRATEGY_TO_WEIGHT_QUANT_KEY[self.strategy] + self.activation_quant_key = ( + kFp8StaticTensorSym if is_static_input_scheme else kFp8DynamicTensorSym + ) + @classmethod def get_min_capability(cls) -> int: # turing and up @@ -89,7 +93,7 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): # WEIGHT SCALE weight_scale = create_fp8_scale_parameter( - strategy_to_parameter_type[self.strategy], + STRATEGY_TO_PARAMETER_TYPE[self.strategy], output_partition_sizes, input_size_per_partition, layer.weight_block_size, @@ -105,32 +109,36 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): ) layer.register_parameter("input_scale", input_scale) + self.linear_kernel = init_wfp8_a16_linear_kernel( + weight_quant_key=self.weight_quant_key, + activation_quant_key=self.activation_quant_key, + weight_shape=layer.weight.shape, + input_dtype=self.input_dtype, + out_dtype=self.out_dtype, + ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - weight = layer.weight - weight_scale = layer.weight_scale - size_k_first = True - # TODO(rob): refactor block quant into separate class. if self.strategy == QuantizationStrategy.BLOCK: assert self.is_static_input_scheme is False - size_k_first = False - weight, weight_scale = process_fp8_weight_block_strategy( - weight, weight_scale - ) + # MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block + # quant, while CT registers the scale as "weight_scale". + # Rename by deleting the old parameter and adding the new one so + # that prepare_fp8_layer_for_marlin (which prefers "weight_scale" + # over "weight_scale_inv") picks up "weight_scale_inv" correctly. + weight_scale_data = layer.weight_scale.data + del layer._parameters["weight_scale"] + replace_parameter(layer, "weight_scale_inv", weight_scale_data) else: - # Weights must be transposed for marlin - weight = weight.t() if self.strategy == QuantizationStrategy.TENSOR: - # If we have a fused module (QKV, MLP) with per tensor scales, - # we expand each scale to its shard's channels. - weight_scale = convert_to_channelwise( - weight_scale, layer.logical_widths + # For fused modules with per-tensor scales, expand each scale + # to its shard's channels. + replace_parameter( + layer, + "weight_scale", + convert_to_channelwise(layer.weight_scale, layer.logical_widths), ) - # Update layer with new values - replace_parameter(layer, "weight", weight.data) - replace_parameter(layer, "weight_scale", weight_scale.data) - - prepare_fp8_layer_for_marlin(layer, size_k_first=size_k_first) + self.linear_kernel.process_weights_after_loading(layer) def apply_weights( self, @@ -138,12 +146,4 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - return apply_fp8_marlin_linear( - input=x, - weight=layer.weight, - weight_scale=layer.weight_scale, - workspace=layer.workspace, - size_n=layer.output_size_per_partition, - size_k=layer.input_size_per_partition, - bias=bias, - ) + return self.linear_kernel.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py index c6b810eb967..3bf606ddb33 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py @@ -16,6 +16,9 @@ from vllm.model_executor.kernels.linear import ( from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) +from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( + STRATEGY_TO_PARAMETER_TYPE, +) from vllm.model_executor.layers.quantization.utils.fp8_utils import ( create_fp8_input_scale, create_fp8_scale_parameter, @@ -34,20 +37,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( cutlass_block_fp8_supported, ) -from vllm.model_executor.parameter import ( - BlockQuantScaleParameter, - ChannelQuantScaleParameter, - PerTensorScaleParameter, -) __all__ = ["CompressedTensorsW8A8Fp8"] -strategy_to_parameter_type = { - QuantizationStrategy.BLOCK: BlockQuantScaleParameter, - QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter, - QuantizationStrategy.TENSOR: PerTensorScaleParameter, -} - STATIC_QUANT = True DYNAMIC_QUANT = False activation_quant_key_mapping = { @@ -130,7 +122,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme): # WEIGHT SCALE weight_scale = create_fp8_scale_parameter( - strategy_to_parameter_type[self.strategy], + STRATEGY_TO_PARAMETER_TYPE[self.strategy], output_partition_sizes, input_size_per_partition, layer.weight_block_size, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py new file mode 100644 index 00000000000..5c511fc98d9 --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_mxfp8.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + +import torch + +from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, +) +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + MXFP8_VALUE_DTYPE, +) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + ModelWeightParameter, +) + +__all__ = ["CompressedTensorsW8A8Mxfp8"] + + +class CompressedTensorsW8A8Mxfp8(CompressedTensorsScheme): + """ + Compressed tensors scheme for MXFP8 quantization (W8A8). + + Loads pre-quantized MXFP8 weights from compressed-tensors checkpoints. + Activations are dynamically quantized to MXFP8 at runtime. + + MXFP8 format: + - 8-bit float weights (E4M3) stored as float8_e4m3fn + - Per-group E8M0 scales (uint8) with group_size=32 + - Activations dynamically quantized to MXFP8 during inference + """ + + def __init__(self): + self.kernel = init_mxfp8_linear_kernel() + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + def create_weights( + self, + layer: torch.nn.Module, + output_partition_sizes: list[int], + input_size_per_partition: int, + params_dtype: torch.dtype, + weight_loader: Callable, + **kwargs, + ): + output_size_per_partition = sum(output_partition_sizes) + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.params_dtype = params_dtype + + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=MXFP8_VALUE_DTYPE, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // MXFP8_BLOCK_SIZE, + dtype=MXFP8_SCALE_DTYPE, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.kernel.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py index f8809216911..04c64d9bd56 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/utils.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/utils.py @@ -6,8 +6,36 @@ from types import MappingProxyType import regex as re from compressed_tensors import CompressionFormat +from compressed_tensors.quantization import QuantizationStrategy from torch.nn import Module +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, +) +from vllm.model_executor.parameter import ( + BlockQuantScaleParameter, + ChannelQuantScaleParameter, + PerTensorScaleParameter, +) + +# Maps quantization strategy to the corresponding scale parameter type. +# Shared across compressed-tensor scheme classes (w8a16_fp8, w8a8_fp8, …). +STRATEGY_TO_PARAMETER_TYPE = { + QuantizationStrategy.BLOCK: BlockQuantScaleParameter, + QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter, + QuantizationStrategy.TENSOR: PerTensorScaleParameter, +} + +# Maps quantization strategy to the vLLM weight-quant key used for +# kernel selection. Shared across compressed-tensor scheme classes. +STRATEGY_TO_WEIGHT_QUANT_KEY = { + QuantizationStrategy.BLOCK: kFp8Static128BlockSym, + QuantizationStrategy.CHANNEL: kFp8StaticChannelSym, + QuantizationStrategy.TENSOR: kFp8StaticTensorSym, +} + def is_activation_quantization_format(format: str) -> bool: _ACTIVATION_QUANTIZATION_FORMATS = [ diff --git a/vllm/model_executor/layers/quantization/cpu_wna16.py b/vllm/model_executor/layers/quantization/cpu_wna16.py index 8ec569042d7..aea1067ff26 100644 --- a/vllm/model_executor/layers/quantization/cpu_wna16.py +++ b/vllm/model_executor/layers/quantization/cpu_wna16.py @@ -104,7 +104,7 @@ class CPUAWQConfig(QuantizationConfig): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> "QuantizationMethods | None": quant_method = hf_quant_cfg.get("quant_method", "").lower() if current_platform.is_cpu() and (quant_method == "awq"): diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index dfb09d57361..d7920462e61 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -517,10 +517,10 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod): # TODO: remove this check once the following RFC is resolved. # https://github.com/vllm-project/vllm/issues/33314 - # This check is required because Mxfp8OnlineLinearMethod inherits from - # Fp8OnlineLinearMethod but only calls super().create_weights(), so we must - # skip the fp8_linear kernel creation. - if hasattr(self, "mxfp8_linear"): + # Subclasses (e.g. Mxfp8OnlineLinearMethod) only need the weight + # registration above and manage their own kernel, so skip fp8_linear + # kernel creation for them. + if type(self) is not Fp8OnlineLinearMethod: return self.fp8_linear = init_fp8_linear_kernel( diff --git a/vllm/model_executor/layers/quantization/gguf.py b/vllm/model_executor/layers/quantization/gguf.py index 2a72da26cc6..61eb6c912a1 100644 --- a/vllm/model_executor/layers/quantization/gguf.py +++ b/vllm/model_executor/layers/quantization/gguf.py @@ -84,7 +84,7 @@ class GGUFConfig(QuantizationConfig): @classmethod def override_quantization_method( - cls, hf_quant_cfg: dict[str, Any], user_quant: str | None + cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None ) -> "QuantizationMethods | None": # When user explicitly specifies --quantization gguf, override # whatever quantization method is in the HF model config (e.g. fp8). diff --git a/vllm/model_executor/layers/quantization/gptq_marlin.py b/vllm/model_executor/layers/quantization/gptq_marlin.py index ce0dc0f4e05..1ca551d6351 100644 --- a/vllm/model_executor/layers/quantization/gptq_marlin.py +++ b/vllm/model_executor/layers/quantization/gptq_marlin.py @@ -214,7 +214,7 @@ class GPTQMarlinConfig(QuantizationConfig): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: can_convert = cls.is_gptq_marlin_compatible(hf_quant_cfg) diff --git a/vllm/model_executor/layers/quantization/inc.py b/vllm/model_executor/layers/quantization/inc.py index 93be5b76130..4457555c076 100644 --- a/vllm/model_executor/layers/quantization/inc.py +++ b/vllm/model_executor/layers/quantization/inc.py @@ -414,6 +414,7 @@ class INCConfig(QuantizationConfig): def apply_xpu_w4a16_quant_layer(self, layer, prefix: str): weight_bits, group_size, sym = self.get_layer_config(layer, prefix) + if not self.check_quantized(weight_bits): if isinstance(layer, (LinearBase, ParallelLMHead)): return UnquantizedLinearMethod() @@ -437,6 +438,27 @@ class INCConfig(QuantizationConfig): ) return None + def apply_cpu_w4a16_quant_layer(self, layer, prefix: str): + weight_bits, group_size, sym = self.get_layer_config(layer, prefix) + if not self.check_quantized(weight_bits): + if isinstance(layer, (LinearBase, ParallelLMHead)): + return UnquantizedLinearMethod() + else: + return None + + if weight_bits != 4: + raise NotImplementedError( + f"INC on CPU only supports 4-bit quantization, " + f"got weight_bits={weight_bits}." + ) + if not sym: + raise NotImplementedError( + "INC W4A16 on CPU only supports symmetric quantization for now." + ) + if isinstance(layer, (LinearBase, ParallelLMHead)): + return self.apply_gptq_quant_layer(layer, prefix) + return None + def get_quant_method(self, layer: torch.nn.Module, prefix: str): if prefix and self.extra_config: for layer_name in self.extra_config: @@ -446,14 +468,24 @@ class INCConfig(QuantizationConfig): return UnquantizedLinearMethod() if current_platform.is_xpu(): return self.apply_xpu_w4a16_quant_layer(layer, prefix) - if "gptq" in self.packing_format or "gptq" in self.backend: + is_gptq = "gptq" in self.packing_format or "gptq" in self.backend + if current_platform.is_cpu() and is_gptq: + return self.apply_cpu_w4a16_quant_layer(layer, prefix) + if is_gptq: return self.apply_gptq_quant_layer(layer, prefix) if "awq" in self.packing_format or "awq" in self.backend: return self.apply_awq_quant_layer(layer, prefix) + raise NotImplementedError( + f"Unsupported quantization configuration for layer '{prefix}'. " + f"Platform: CPU={current_platform.is_cpu()}. " + f"Platform: XPU={current_platform.is_xpu()}. " + f"Format: {self.packing_format}, Backend: {self.backend}." + ) + @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> "QuantizationMethods | None": """Override the `auto-round` method to `inc`.""" is_auto_round_format = hf_quant_cfg.get("quant_method", None) == "auto-round" diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index ad188c665e9..0b8ad0cbc1e 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -10,7 +10,11 @@ from torch.nn.parameter import Parameter import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.config import get_current_vllm_config from vllm.logger import init_logger -from vllm.model_executor.kernels.linear import init_fp8_linear_kernel +from vllm.model_executor.kernels.linear import ( + init_fp8_linear_kernel, + init_mxfp8_linear_kernel, + init_nvfp4_linear_kernel, +) from vllm.model_executor.layers.attention import Attention, MLAAttention from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -67,15 +71,8 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, MXFP8_SCALE_DTYPE, MXFP8_VALUE_DTYPE, - Mxfp8LinearOp, mxfp8_e4m3_quantize, ) -from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( - NvFp4LinearBackend, - apply_nvfp4_linear, - convert_to_nvfp4_linear_kernel_format, - select_nvfp4_linear_backend, -) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, create_fp8_quant_key, @@ -409,7 +406,7 @@ class ModelOptFp8Config(ModelOptQuantConfigBase): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) if algo is not None and algo == "FP8": @@ -1031,7 +1028,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) if algo is not None and ("NVFP4" in algo or "FP4" in algo): @@ -1090,11 +1087,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase): def __init__(self, quant_config: ModelOptNvFp4Config) -> None: self.quant_config = quant_config self.marlin_input_dtype = None - self.backend = select_nvfp4_linear_backend() - - self.swizzle = None - if self.backend == NvFp4LinearBackend.EMULATION: - self.swizzle = False + self.kernel = init_nvfp4_linear_kernel() def create_weights( self, @@ -1201,7 +1194,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase): ) # Convert layer to NVFP4 linear kernel format - convert_to_nvfp4_linear_kernel_format(self.backend, layer) + self.kernel.process_weights_after_loading(layer) def apply( self, @@ -1209,13 +1202,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - return apply_nvfp4_linear( - backend=self.backend, - layer=layer, - x=x, - bias=bias, - swizzle=self.swizzle, - ) + return self.kernel.apply_weights(layer=layer, x=x, bias=bias) class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): @@ -1538,7 +1525,7 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) if algo is not None and "MXFP8" in algo: @@ -1589,7 +1576,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase): "Dynamic quantization is not supported." ) - self.mxfp8_linear_op = Mxfp8LinearOp() + self.kernel = init_mxfp8_linear_kernel() def create_weights( self, @@ -1671,7 +1658,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase): f" got {layer.weight_scale.dtype}" ) - self.mxfp8_linear_op.process_weights(layer) + self.kernel.process_weights_after_loading(layer) def apply( self, @@ -1679,16 +1666,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - return self.mxfp8_linear_op.apply( - input=x, - weight=layer.weight, - weight_scale=layer.weight_scale, - out_dtype=x.dtype, - bias=bias, - workspace=getattr(layer, "workspace", None), - size_n=layer.output_size_per_partition, - size_k=layer.input_size_per_partition, - ) + return self.kernel.apply_weights(layer, x, bias) class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): @@ -2074,7 +2052,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: algo = cls._extract_modelopt_quant_algo(hf_quant_cfg) if algo is not None and algo == "MIXED_PRECISION": diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index a327ac17bbc..e5ef3f4c316 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -130,7 +130,7 @@ class MoeWNA16Config(QuantizationConfig): @classmethod def override_quantization_method( - cls, hf_quant_cfg, user_quant + cls, hf_quant_cfg, user_quant, hf_config=None ) -> QuantizationMethods | None: can_convert = cls.is_moe_wna16_compatible(hf_quant_cfg) if can_convert and user_quant == "moe_wna16": diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index adb191b0a0f..019bb45d65d 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -19,11 +19,11 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, - convert_to_mxfp4_moe_kernel_format, + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, mxfp4_round_up_hidden_size_and_intermediate_size, - select_mxfp4_moe_backend, + select_gpt_oss_mxfp4_moe_backend, ) from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod from vllm.model_executor.layers.quantization import QuantizationMethods @@ -38,6 +38,12 @@ logger = init_logger(__name__) class Mxfp4Config(QuantizationConfig): + """Canonical base config for MXFP4 quantization. + + Subclasses override get_name() and override_quantization_method() to + register themselves as the handler for a specific checkpoint format. + """ + def __init__(self, ignored_layers: list[str] | None = None): super().__init__() self.ignored_layers = ignored_layers @@ -62,6 +68,8 @@ class Mxfp4Config(QuantizationConfig): def get_config_filenames(cls) -> list[str]: return [] + # TODO (zyongye) This is only temporaty fallback. + # We should have `Mxfp4MoEMethod` after this migration is complete. def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": @@ -79,7 +87,7 @@ class Mxfp4Config(QuantizationConfig): ) return UnquantizedLinearMethod() elif isinstance(layer, FusedMoE): - return Mxfp4MoEMethod(layer.moe_config) + return GptOssMxfp4MoEMethod(layer.moe_config) elif isinstance(layer, Attention): logger.debug_once( "MXFP4 attention layer is not implemented. " @@ -93,13 +101,46 @@ class Mxfp4Config(QuantizationConfig): return True -class Mxfp4MoEMethod(FusedMoEMethodBase): +class GptOssMxfp4Config(Mxfp4Config): + """MXFP4 config for GPT-OSS checkpoints. + + Checkpoints carry ``"quant_method": "mxfp4"`` in their JSON config. + override_quantization_method() maps that to the canonical internal name + so that the rest of the loading path uses "gpt_oss_mxfp4" consistently. + """ + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "gpt_oss_mxfp4" + + @classmethod + def override_quantization_method( + cls, hf_quant_cfg, user_quant, hf_config=None + ) -> QuantizationMethods | None: + # Match both "mxfp4" (original checkpoint value) and "gpt_oss_mxfp4" + # (already normalized by verify_and_update_model_config) so that + # explicit --quantization mxfp4 from the user doesn't cause a mismatch. + if not ( + isinstance(hf_quant_cfg, dict) + and hf_quant_cfg.get("quant_method") in ("mxfp4", "gpt_oss_mxfp4") + ): + return None + # Require explicit confirmation that this is a GPT-OSS model. + # Do NOT fall back to returning the override when hf_config is None, + # as that would silently claim all mxfp4 checkpoints. + model_type = getattr(hf_config, "model_type", None) + if model_type != "gpt_oss": + return None + return "gpt_oss_mxfp4" + + +class GptOssMxfp4MoEMethod(FusedMoEMethodBase): """MXFP4 MoE quantization method.""" def __init__(self, moe: FusedMoEConfig): super().__init__(moe) - self.weight_dtype = "mxfp4" - self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) + self.weight_dtype = "gpt_oss_mxfp4" + self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe) self.max_capture_size = ( get_current_vllm_config().compilation_config.max_cudagraph_capture_size @@ -281,7 +322,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # Convert weights to kernel format w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( - convert_to_mxfp4_moe_kernel_format( + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend=self.mxfp4_backend, layer=layer, w13_weight=w13, diff --git a/vllm/model_executor/layers/quantization/mxfp8.py b/vllm/model_executor/layers/quantization/mxfp8.py index 6e0c1414385..5acf843f108 100644 --- a/vllm/model_executor/layers/quantization/mxfp8.py +++ b/vllm/model_executor/layers/quantization/mxfp8.py @@ -9,6 +9,7 @@ import torch from torch.nn import Module from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( FusedMoE, @@ -34,7 +35,6 @@ from vllm.model_executor.layers.quantization.fp8 import ( ) from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_BLOCK_SIZE, - Mxfp8LinearOp, mxfp8_e4m3_quantize, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -126,8 +126,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod): def __init__(self, quant_config: "Mxfp8Config"): self.quant_config = quant_config - self.out_dtype = torch.get_default_dtype() - self.mxfp8_linear = Mxfp8LinearOp() + self.kernel = init_mxfp8_linear_kernel() def create_weights( self, @@ -166,7 +165,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod): replace_parameter(layer, "weight", weight_fp8.data) replace_parameter(layer, "weight_scale", weight_scale.data) - self.mxfp8_linear.process_weights(layer) + self.kernel.process_weights_after_loading(layer) layer._already_called_process_weights_after_loading = True @@ -176,16 +175,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - return self.mxfp8_linear.apply( - input=x, - weight=layer.weight, - weight_scale=layer.weight_scale, - out_dtype=self.out_dtype, - bias=bias, - workspace=getattr(layer, "workspace", None), - size_n=layer.output_size_per_partition, - size_k=layer.input_size_per_partition, - ) + return self.kernel.apply_weights(layer, x, bias) class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod): diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index d0362cedcf2..33bd0cfc22e 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -389,6 +389,37 @@ class QuarkConfig(QuantizationConfig): return is_weight_mxfp4 and is_input_fp8 + def _is_dynamic_per_token_w8a8( + self, + weight_quant: dict[str, Any] | None, + input_quant: dict[str, Any] | None, + ) -> bool: + """Detect W8A8 INT8 with per-tensor or per-channel + weights and dynamic per-token input.""" + if weight_quant is None or input_quant is None: + return False + + is_int8_dtype = ( + weight_quant.get("dtype") == "int8" and input_quant.get("dtype") == "int8" + ) + + is_valid_weight_scheme = weight_quant.get("qscheme") in [ + "per_tensor", + "per_channel", + ] + is_per_token_input = input_quant.get("qscheme") == "per_channel" + + is_dynamic_input = input_quant.get("is_dynamic") is True + is_weight_symmetric = weight_quant.get("symmetric") is True + + return ( + is_int8_dtype + and is_valid_weight_scheme + and is_per_token_input + and is_dynamic_input + and is_weight_symmetric + ) + def _is_w_ocp_mx_a_x( self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None ) -> bool: @@ -556,6 +587,13 @@ class QuarkConfig(QuantizationConfig): ) if is_w4a8_supported: return QuarkW4A8_MXFP4_FP8(weight_config, input_config) + elif self._is_dynamic_per_token_w8a8(weight_config, input_config): + weight_qscheme = cast(str, weight_config.get("qscheme")) + return QuarkW8A8Int8( + qscheme=weight_qscheme, + is_static_input_scheme=False, + input_symmetric=input_config.get("symmetric"), + ) elif self._is_w_ocp_mx_a_x(weight_config, input_config): return QuarkOCP_MX( weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 3f7ddbfd756..2bab66709dd 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -30,11 +30,11 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_m from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, - convert_to_mxfp4_moe_kernel_format, + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, mxfp4_round_up_hidden_size_and_intermediate_size, - select_mxfp4_moe_backend, + select_gpt_oss_mxfp4_moe_backend, ) from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( prepare_fp8_moe_layer_for_marlin, @@ -109,6 +109,12 @@ class QuarkMoEMethod(FusedMoEMethodBase): return QuarkOCP_MX_MoEMethod( weight_config, input_config, module.moe_config ) + elif quant_config._is_static_tensor_w8a8( + weight_config, input_config + ) or quant_config._is_dynamic_per_token_w8a8(weight_config, input_config): + return QuarkW8A8Int8MoEMethod( + weight_config, input_config, module.moe_config + ) else: raise RuntimeError("Unsupported FusedMoe scheme") @@ -505,6 +511,282 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): ) +class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): + """Quark W8A8 INT8 MoE method.""" + + def __init__( + self, + weight_config: dict[str, Any], + input_config: dict[str, Any], + moe: FusedMoEConfig, + ): + super().__init__(moe) + self.weight_quant = weight_config + self.input_quant = input_config + self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor") + self.static_input_scales = not self.input_quant.get("is_dynamic", False) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.orig_dtype = params_dtype + layer.weight_block_size = None + params_dtype = torch.int8 + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + if self.weight_qscheme == "per_channel": + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + else: + # per-tensor: one scalar per expert + w13_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, 2, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + w2_weight_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + if self.static_input_scales: + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + # ZERO POINTS (loaded but discarded after loading; kernel uses symmetric) + w13_input_zero_point = torch.nn.Parameter( + torch.zeros(num_experts, 2, dtype=torch.int8), + requires_grad=False, + ) + layer.register_parameter("w13_input_zero_point", w13_input_zero_point) + set_weight_attrs(w13_input_zero_point, extra_weight_attrs) + + w2_input_zero_point = torch.nn.Parameter( + torch.zeros(num_experts, dtype=torch.int8), + requires_grad=False, + ) + layer.register_parameter("w2_input_zero_point", w2_input_zero_point) + set_weight_attrs(w2_input_zero_point, extra_weight_attrs) + + if self.weight_qscheme == "per_channel": + w13_weight_zero_point = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=torch.int8, + ), + requires_grad=False, + ) + w2_weight_zero_point = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=torch.int8), + requires_grad=False, + ) + else: + w13_weight_zero_point = torch.nn.Parameter( + torch.zeros(num_experts, 2, dtype=torch.int8), + requires_grad=False, + ) + w2_weight_zero_point = torch.nn.Parameter( + torch.zeros(num_experts, dtype=torch.int8), + requires_grad=False, + ) + layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point) + set_weight_attrs(w13_weight_zero_point, extra_weight_attrs) + layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point) + set_weight_attrs(w2_weight_zero_point, extra_weight_attrs) + + # BIAS + if self.has_bias: + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + else: + layer.w13_bias, layer.w2_bias = None, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Discard zero points (INT8 fused MoE kernel uses symmetric quant) + for attr in ( + "w13_input_zero_point", + "w2_input_zero_point", + "w13_weight_zero_point", + "w2_weight_zero_point", + ): + if hasattr(layer, attr): + delattr(layer, attr) + + # For static input scales, collapse per-expert scales to single max + if self.static_input_scales: + if layer.w13_input_scale is None or layer.w2_input_scale is None: + raise ValueError( + "QuantConfig has static quantization, but found " + "activation scales are None." + ) + if not all_close_1d(layer.w13_input_scale) or not all_close_1d( + layer.w2_input_scale + ): + logger.warning_once( + "Found input_scales that are not equal for " + "INT8 MoE layer. Using the maximum across experts " + "for each layer." + ) + layer.w13_input_scale = torch.nn.Parameter( + layer.w13_input_scale.max(), requires_grad=False + ) + layer.w2_input_scale = torch.nn.Parameter( + layer.w2_input_scale.max(), requires_grad=False + ) + + # For per-tensor weights, merge w1/w3 scales into single per-expert + if self.weight_qscheme == "per_tensor": + assert layer.w13_weight_scale is not None + shard_size = layer.intermediate_size_per_partition + max_w13_scales = layer.w13_weight_scale.max(dim=1).values + + for expert_id in range(layer.local_num_experts): + start = 0 + for shard_id in range(2): + dq_weight = per_tensor_dequantize( + layer.w13_weight[expert_id][start : start + shard_size, :], + layer.w13_weight_scale[expert_id][shard_id], + ) + layer.w13_weight[expert_id][start : start + shard_size, :], _, _ = ( + ops.scaled_int8_quant( + dq_weight, + scale=max_w13_scales[expert_id], + ) + ) + start += shard_size + + layer.w13_weight_scale = torch.nn.Parameter( + max_w13_scales, requires_grad=False + ) + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + is_dynamic = not self.static_input_scales + is_per_channel = self.weight_qscheme == "per_channel" + return FusedMoEQuantConfig.make( + torch.int8, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + a1_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), + per_act_token_quant=is_dynamic, + per_out_ch_quant=is_per_channel, + block_shape=None, + ) + + def apply( + self, + layer: FusedMoE, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + from vllm.model_executor.layers.fused_moe import fused_experts + + return fused_experts( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=not self.moe.disable_inplace, + activation=layer.activation, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + quant_config=self.moe_quant_config, + ) + + class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod): def __init__( self, @@ -713,7 +995,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): self.w2_precision_config = None if self.ocp_mx_scheme == "w_mxfp4": - self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) + self.mxfp4_backend, self.experts_cls = select_gpt_oss_mxfp4_moe_backend(moe) elif self.ocp_mx_scheme.startswith("w_mxfp4"): # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes. self.mxfp4_backend = Mxfp4MoeBackend.NONE @@ -1018,7 +1300,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod): # Convert weights to kernel format w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = ( - convert_to_mxfp4_moe_kernel_format( + convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( mxfp4_backend=self.mxfp4_backend, layer=layer, w13_weight=w13, @@ -1220,9 +1502,9 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): layer.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False) # FIXME warp need to be adjusted based on batch size - # only apply to batched mode + # only apply to batched mode if self.moe.use_ep: - num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8 + num_warps = 4 if self.moe.max_num_tokens <= 512 else 8 else: num_warps = 8 @@ -1309,7 +1591,7 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet." ) - from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501 + from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 triton_kernel_moe_forward, ) diff --git a/vllm/model_executor/layers/quantization/turboquant/__init__.py b/vllm/model_executor/layers/quantization/turboquant/__init__.py new file mode 100644 index 00000000000..10ee032c9ec --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant: Near-optimal KV-cache quantization for vLLM. + +PolarQuant compression: random rotation + per-coordinate Lloyd-Max +scalar quantization for keys, uniform quantization for values. + +Reference: "TurboQuant: Online Vector Quantization with Near-optimal +Distortion Rate" (ICLR 2026), Zandieh et al. +""" + +from vllm.model_executor.layers.quantization.turboquant.config import TurboQuantConfig + +__all__ = ["TurboQuantConfig"] diff --git a/vllm/model_executor/layers/quantization/turboquant/centroids.py b/vllm/model_executor/layers/quantization/turboquant/centroids.py new file mode 100644 index 00000000000..490265747c5 --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/centroids.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lloyd-Max optimal scalar quantizer for TurboQuant. + +After rotating a d-dimensional unit vector by a random orthogonal matrix, +each coordinate approximately follows N(0, 1/d) for d >= 64. +We solve the Lloyd-Max conditions to find optimal centroids. + +Based on: turboquant-pytorch/lloyd_max.py (Zandieh et al.) +""" + +import math +from functools import lru_cache + +import torch + + +def _gaussian_pdf(x: float, sigma2: float) -> float: + return (1.0 / math.sqrt(2 * math.pi * sigma2)) * math.exp(-x * x / (2 * sigma2)) + + +def _trapz(f, a: float, b: float, n: int = 200) -> float: + """Trapezoidal numerical integration (replaces scipy.integrate.quad).""" + h = (b - a) / n + result = 0.5 * (f(a) + f(b)) + for i in range(1, n): + result += f(a + i * h) + return result * h + + +def solve_lloyd_max( + d: int, + bits: int, + max_iter: int = 200, + tol: float = 1e-10, +) -> tuple[torch.Tensor, torch.Tensor]: + """Solve Lloyd-Max optimal quantizer for N(0, 1/d) distribution. + + Args: + d: Vector dimension (determines variance = 1/d). + bits: Number of quantization bits. + max_iter: Maximum Lloyd-Max iterations. + tol: Convergence tolerance. + + Returns: + centroids: Sorted tensor of 2^bits optimal centroids. + boundaries: Sorted tensor of 2^bits - 1 decision boundaries. + """ + n_levels = 2**bits + sigma2 = 1.0 / d + sigma = math.sqrt(sigma2) + + def pdf(x): + return _gaussian_pdf(x, sigma2) + + lo, hi = -3.5 * sigma, 3.5 * sigma + centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] + + for _ in range(max_iter): + boundaries = [ + (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1) + ] + edges = [lo * 3] + boundaries + [hi * 3] + new_centroids = [] + for i in range(n_levels): + a, b = edges[i], edges[i + 1] + num = _trapz(lambda x: x * pdf(x), a, b) + den = _trapz(pdf, a, b) + new_centroids.append(num / den if den > 1e-15 else centroids[i]) + + if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol: + break + centroids = new_centroids + + boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)] + return ( + torch.tensor(centroids, dtype=torch.float32), + torch.tensor(boundaries, dtype=torch.float32), + ) + + +@lru_cache(maxsize=32) +def get_centroids(d: int, bits: int) -> torch.Tensor: + """Get precomputed Lloyd-Max centroids (cached).""" + centroids, _ = solve_lloyd_max(d, bits) + return centroids diff --git a/vllm/model_executor/layers/quantization/turboquant/config.py b/vllm/model_executor/layers/quantization/turboquant/config.py new file mode 100644 index 00000000000..289bed12077 --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/config.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant configuration.""" + +import math +from dataclasses import dataclass + +# Named TQ presets: each maps to frozen config parameters. +# key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys. +# value_quant_bits: 3-4 = uniform quantized values. +TQ_PRESETS: dict[str, dict] = { + "turboquant_k8v4": { + "key_quant_bits": 8, + "value_quant_bits": 4, + "norm_correction": False, + }, + "turboquant_4bit_nc": { + "key_quant_bits": 4, + "value_quant_bits": 4, + "norm_correction": True, + }, + "turboquant_k3v4_nc": { + "key_quant_bits": 3, + "value_quant_bits": 4, + "norm_correction": True, + }, + "turboquant_3bit_nc": { + "key_quant_bits": 3, + "value_quant_bits": 3, + "norm_correction": True, + }, +} + + +@dataclass +class TurboQuantConfig: + """Configuration for TurboQuant KV-cache quantization. + + Uses PolarQuant (WHT rotation + Lloyd-Max scalar quantization) for keys + and uniform quantization for values. QJL is intentionally omitted — + community consensus (5+ independent groups) found it hurts attention + quality by amplifying variance through softmax. + + Named presets (use via --kv-cache-dtype): + turboquant_k8v4: FP8 keys + 4-bit values, 2.6x, +1.17% PPL + turboquant_4bit_nc: 4-bit MSE keys + 4-bit values + NC, 3.8x, +2.71% + turboquant_k3v4_nc: 3-bit MSE keys + 4-bit values + NC, ~3.5x, +10.63% + turboquant_3bit_nc: 3-bit MSE keys + 3-bit values + NC, 4.9x, +20.59% + + Args: + head_dim: Attention head dimension (e.g. 64, 96, 128). + key_quant_bits: Bits for key quantization. 8 = FP8 keys (no + rotation/MSE). 3-4 = Lloyd-Max MSE quantized keys. + value_quant_bits: Bits per value dimension for uniform quantization. + 3 = 8 levels, 4 = 16 levels (default). + seed: Base seed for deterministic random matrix generation. + Actual seed per layer = seed + layer_idx * 1337. + norm_correction: Re-normalize centroid vectors to unit norm before + inverse rotation during dequant. Fixes quantization-induced norm + distortion, improving PPL by ~0.8% at 4-bit. + """ + + head_dim: int = 128 + key_quant_bits: int = 3 # 3-4 = MSE keys, 8 = FP8 keys + value_quant_bits: int = 4 # 3-4 = uniform quantized values + seed: int = 42 + norm_correction: bool = False + + @property + def key_fp8(self) -> bool: + """Whether keys are stored as FP8 — no rotation/quantization needed.""" + return self.key_quant_bits == 8 + + @property + def mse_bits(self) -> int: + """MSE quantizer bit-width (determines centroid count: 2^mse_bits). + + For MSE key modes, equals key_quant_bits. + For FP8 key mode, falls back to value_quant_bits (centroids are still + needed for continuation-prefill dequant and decode kernel params). + """ + if self.key_fp8: + return self.value_quant_bits + return self.key_quant_bits + + @property + def key_mse_bits(self) -> int: + """MSE bits actually used for key quantization (0 if FP8 keys).""" + if self.key_fp8: + return 0 + return self.key_quant_bits + + @property + def centroid_bits(self) -> int: + """Bits for centroid generation — always non-zero.""" + return self.mse_bits + + @property + def n_centroids(self) -> int: + return 2**self.mse_bits + + @property + def key_packed_size(self) -> int: + """Packed bytes for a single KEY vector. + + FP8 mode (key_quant_bits=8): + head_dim bytes (1 byte per element, no overhead). + + TQ mode: + - MSE indices: ceil(head_dim * key_mse_bits / 8) bytes + - vec_norm: 2 bytes (float16) + """ + if self.key_fp8: + return self.head_dim # 1 byte per element + mse_bytes = math.ceil(self.head_dim * self.key_mse_bits / 8) + norm_bytes = 2 # vec_norm fp16 + return mse_bytes + norm_bytes + + @property + def effective_value_quant_bits(self) -> int: + """Actual bits used for value storage.""" + return self.value_quant_bits + + @property + def value_packed_size(self) -> int: + """Packed bytes for a single VALUE vector. + + Uniform quantization: ceil(head_dim * bits / 8) + 4 bytes (scale + zero fp16). + """ + data_bytes = math.ceil(self.head_dim * self.value_quant_bits / 8) + return data_bytes + 4 # +2 scale(fp16) +2 zero(fp16) + + @property + def slot_size(self) -> int: + """Total packed bytes per head per position (key + value combined). + + Layout: [key_packed | value_packed] + """ + return self.key_packed_size + self.value_packed_size + + @property + def slot_size_aligned(self) -> int: + """Slot size rounded up to next even number. + + Even-number is required so effective_head_size = slot_size_aligned // 2 + is integral. + """ + s = self.slot_size + return s + (s % 2) # round up to even + + @staticmethod + def get_boundary_skip_layers(num_layers: int, n: int = 2) -> list[str]: + """Get layer indices to skip TQ compression (boundary protection). + + Returns first N and last N layer indices as strings, suitable for + kv_cache_dtype_skip_layers. + """ + if n <= 0 or num_layers <= 0: + return [] + n = min(n, num_layers // 2) # don't skip more than half + first = list(range(n)) + last = list(range(num_layers - n, num_layers)) + # Deduplicate (if num_layers <= 2*n) + indices = sorted(set(first + last)) + return [str(i) for i in indices] + + @staticmethod + def from_cache_dtype(cache_dtype: str, head_dim: int) -> "TurboQuantConfig": + """Create config from a named preset. + + Valid presets: turboquant_k8v4, turboquant_4bit_nc, etc. + """ + if cache_dtype not in TQ_PRESETS: + valid = ", ".join(TQ_PRESETS.keys()) + raise ValueError( + f"Unknown TurboQuant cache dtype: {cache_dtype!r}. " + f"Valid presets: {valid}" + ) + preset = TQ_PRESETS[cache_dtype] + return TurboQuantConfig( + head_dim=head_dim, + key_quant_bits=preset["key_quant_bits"], + value_quant_bits=preset["value_quant_bits"], + norm_correction=preset["norm_correction"], + ) diff --git a/vllm/model_executor/layers/quantization/turboquant/quantizer.py b/vllm/model_executor/layers/quantization/turboquant/quantizer.py new file mode 100644 index 00000000000..aea63c52bac --- /dev/null +++ b/vllm/model_executor/layers/quantization/turboquant/quantizer.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TurboQuant quantizer utilities. + +Serving path uses generate_wht_signs() for WHT rotation sign buffers. +Triton kernels handle all quantization, packing, and dequantization on GPU. +""" + +import torch + +_CPU = torch.device("cpu") + + +def generate_wht_signs(d: int, seed: int, device: torch.device = _CPU) -> torch.Tensor: + """Generate deterministic random ±1 signs for WHT rotation. + + Used with Walsh-Hadamard Transform for per-layer rotation randomization. + Same seed derivation as QR (per-layer via seed + layer_idx * stride). + """ + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + bits = torch.randint(0, 2, (d,), generator=gen, device="cpu") + signs = bits.float() * 2 - 1 + return signs.to(device) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 397442aeced..ef0bf2bf7ac 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -10,6 +10,7 @@ import vllm.envs as envs from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( align_fp4_moe_weights_for_fi, + align_trtllm_fp4_moe_hidden_dim_for_fi, ) from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( swizzle_blockscale, @@ -341,6 +342,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels. if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: + w13, w13_scale, w2, w2_scale, padded_hidden = ( + align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) + ) + if layer.moe_config.hidden_dim_unpadded is None: + layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim + layer.moe_config.hidden_dim = padded_hidden + # Align weights for FI NVFP4 MoE kernels. min_alignment = 16 if is_gated else 128 w13, w13_scale, w2, w2_scale, padded_intermediate = ( diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 0e39dc881f2..32c7a772f3f 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -265,6 +265,48 @@ def align_fp4_moe_weights_for_fi( return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate +def align_trtllm_fp4_moe_hidden_dim_for_fi( + w13: torch.Tensor, + w13_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + min_alignment: int = 256, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + num_experts, gate_up_dim, packed_hidden_size = w13.shape + hidden_size = packed_hidden_size * 2 + padded_hidden_size = round_up(hidden_size, min_alignment) + + if padded_hidden_size == hidden_size: + return w13, w13_scale, w2, w2_scale, hidden_size + + logger.warning_once( + "Padding hidden size from %d to %d for TRTLLM NVFP4 MoE weights. " + "This requires activation slicing at runtime and may cause " + "performance degradation.", + hidden_size, + padded_hidden_size, + scope="local", + ) + + padded_w13 = w13.new_zeros((num_experts, gate_up_dim, padded_hidden_size // 2)) + padded_w13[:, :, :packed_hidden_size] = w13 + + padded_w13_scale = w13_scale.new_zeros( + (num_experts, gate_up_dim, padded_hidden_size // 16) + ) + padded_w13_scale[:, :, : w13_scale.shape[2]] = w13_scale + + padded_w2 = w2.new_zeros((num_experts, padded_hidden_size, w2.shape[2])) + padded_w2[:, : w2.shape[1], :] = w2 + + padded_w2_scale = w2_scale.new_zeros( + (num_experts, padded_hidden_size, w2_scale.shape[2]) + ) + padded_w2_scale[:, : w2_scale.shape[1], :] = w2_scale + + return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_hidden_size + + def align_fp8_moe_weights_for_fi( w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool, min_alignment: int = 16 ) -> tuple[torch.Tensor, torch.Tensor, int]: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py index 21c8aba1d56..51b7b29551d 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp4_utils.py @@ -162,3 +162,7 @@ try: quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4 except AttributeError as error: raise error + + +def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.vllm.xpu_mxfp4_quantize(x) diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index a3b2838e698..b9b7bd54273 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -1,52 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from enum import Enum - import torch -from torch.nn.parameter import Parameter -from vllm.logger import init_logger -from vllm.utils import flashinfer as vllm_flashinfer from vllm.utils.torch_utils import direct_register_custom_op -logger = init_logger(__name__) - - -class Mxfp8LinearBackend(Enum): - EMULATION = "emulation" - FLASHINFER_CUTLASS = "flashinfer-cutlass" - MARLIN = "marlin" - - # MXFP8 constants MXFP8_VALUE_DTYPE = torch.float8_e4m3fn MXFP8_SCALE_DTYPE = torch.uint8 MXFP8_BLOCK_SIZE = 32 -def select_mxfp8_linear_backend() -> Mxfp8LinearBackend: - """Select the best MXFP8 linear backend for the current device. - - - SM100+ (Blackwell): FLASHINFER_CUTLASS (native MXFP8 W8A8 GEMM) - - SM80+ (Ampere/Ada): MARLIN (MXFP8 W8A16 GEMM) - - Otherwise: EMULATION (dequant to BF16 fallback) - """ - from vllm.platforms import current_platform - - if current_platform.has_device_capability(100): - return Mxfp8LinearBackend.FLASHINFER_CUTLASS - - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( - is_fp8_marlin_supported, - ) - - if is_fp8_marlin_supported(): - return Mxfp8LinearBackend.MARLIN - - return Mxfp8LinearBackend.EMULATION - - def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor: """Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout.""" scaling_vector_size = MXFP8_BLOCK_SIZE # 32 for MXFP8 @@ -209,194 +173,3 @@ def xpu_mxfp8_quantize( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: return torch.ops.vllm.xpu_mxfp8_quantize(x, dtype) - - -class Mxfp8LinearOp: - def __init__(self): - self.backend = select_mxfp8_linear_backend() - logger.info_once("Using %s backend for MXFP8 GEMM", self.backend) - - def process_weights(self, layer: torch.nn.Module) -> None: - """Process MXFP8 weights after loading into backend-specific format.""" - if self.backend == Mxfp8LinearBackend.MARLIN: - self._process_weights_marlin(layer) - elif self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS: - self._process_weights_flashinfer_cutlass(layer) - else: - self._process_weights_emulation(layer) - - def _process_weights_emulation(self, layer: torch.nn.Module) -> None: - """Keep scales as 2D uint8 for dequant-to-BF16 emulation.""" - weight = layer.weight.data # [N, K] - N, K = weight.shape - scale_k = K // MXFP8_BLOCK_SIZE - - weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() - - layer.weight = Parameter(weight.contiguous(), requires_grad=False) - layer.weight_scale = Parameter(weight_scale, requires_grad=False) - - def _process_weights_flashinfer_cutlass(self, layer: torch.nn.Module) -> None: - """Swizzle scales to F8_128x4 layout for flashinfer CUTLASS.""" - weight = layer.weight.data # [N, K] - N, K = weight.shape - - scale_k = K // MXFP8_BLOCK_SIZE - weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous() - weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K) - - layer.weight = Parameter(weight.contiguous(), requires_grad=False) - layer.weight_scale = Parameter( - weight_scale_swizzled.contiguous(), requires_grad=False - ) - - def _process_weights_marlin(self, layer: torch.nn.Module) -> None: - """Repack MXFP8 weights and scales into Marlin kernel format.""" - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( - prepare_mxfp8_layer_for_marlin, - ) - - prepare_mxfp8_layer_for_marlin(layer) - - def _apply_emulation( - self, - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - out_dtype: torch.dtype, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - if weight_scale.dtype != MXFP8_SCALE_DTYPE: - raise ValueError( - f"TORCH backend requires {MXFP8_SCALE_DTYPE} weight_scale dtype, " - f"got {weight_scale.dtype}." - ) - if weight_scale.ndim != 2: - raise ValueError( - f"TORCH backend requires 2D weight_scale, got {weight_scale.ndim}D. " - f"Ensure process_weights_after_loading was called." - ) - - weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale) - - output = torch.nn.functional.linear(input, weight_bf16, bias) - return output.to(out_dtype) - - def _apply_flashinfer_cutlass( - self, - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - out_dtype: torch.dtype, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - N, K = weight.shape - - input_shape = input.shape - input_2d = input.view(-1, K) - M_orig = input_2d.shape[0] - - # Minimum dimension size for F8_128x4 block scaling layout - min_dim = 128 - - assert min_dim <= K, ( - f"mm_mxfp8 requires K >= {min_dim}, got K={K}. " - f"in_features is too small for mm_mxfp8." - ) - assert K % MXFP8_BLOCK_SIZE == 0, ( - f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}." - ) - assert min_dim <= N, ( - f"mm_mxfp8 requires N >= {min_dim}, got N={N}. " - f"out_features is too small for mm_mxfp8." - ) - - M_padded = ((M_orig + min_dim - 1) // min_dim) * min_dim - if M_padded != M_orig: - pad_rows = M_padded - M_orig - input_2d = torch.nn.functional.pad(input_2d, (0, 0, 0, pad_rows)) - - input_mxfp8, input_scale = mxfp8_e4m3_quantize( - input_2d, - is_sf_swizzled_layout=True, # Swizzled for best accuracy - ) - - if not weight.is_contiguous(): - weight = weight.contiguous() - - output = vllm_flashinfer.mm_mxfp8( - input_mxfp8, - weight.t(), - input_scale, - weight_scale, - out_dtype=out_dtype, - backend="cutlass", - ) - - if M_padded != M_orig: - output = output[:M_orig, :] - - if bias is not None: - output = output + bias - - output_shape = (*input_shape[:-1], N) - return output.view(output_shape) - - def _apply_marlin( - self, - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - out_dtype: torch.dtype, - bias: torch.Tensor | None = None, - *, - workspace: torch.Tensor, - size_n: int, - size_k: int, - ) -> torch.Tensor: - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( - apply_mxfp8_marlin_linear, - ) - - return apply_mxfp8_marlin_linear( - input=input, - weight=weight, - weight_scale=weight_scale, - workspace=workspace, - size_n=size_n, - size_k=size_k, - bias=bias, - ) - - def apply( - self, - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - out_dtype: torch.dtype, - bias: torch.Tensor | None = None, - *, - workspace: torch.Tensor | None = None, - size_n: int = 0, - size_k: int = 0, - ) -> torch.Tensor: - if self.backend == Mxfp8LinearBackend.EMULATION: - return self._apply_emulation(input, weight, weight_scale, out_dtype, bias) - - if self.backend == Mxfp8LinearBackend.MARLIN: - assert workspace is not None - return self._apply_marlin( - input, - weight, - weight_scale, - out_dtype, - bias, - workspace=workspace, - size_n=size_n, - size_k=size_k, - ) - - assert self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS - return self._apply_flashinfer_cutlass( - input, weight, weight_scale, out_dtype, bias - ) diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py index 12032274f18..539a28d4cb2 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py +++ b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py @@ -1,352 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from enum import Enum import torch -import vllm.envs as envs from vllm._custom_ops import ( - cutlass_scaled_fp4_mm, cutlass_scaled_mm_supports_fp4, - scaled_fp4_quant, -) -from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( - apply_fp4_marlin_linear, - is_fp4_marlin_supported, - prepare_fp4_layer_for_marlin, -) -from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( - kE2M1ToFloat_handle, - run_nvfp4_emulations, ) from vllm.platforms import current_platform -from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer -from vllm.utils.import_utils import has_fbgemm_gpu from vllm.utils.math_utils import round_up -logger = init_logger(__name__) - - -# NOTE: This is ordered by preferred backend. -# Example: if both are available, FLASHINFER_CUTLASS is preferred to VLLM_CUTLASS. -class NvFp4LinearBackend(Enum): - FLASHINFER_CUTLASS = "flashinfer-cutlass" - VLLM_CUTLASS = "cutlass" - MARLIN = "marlin" - FLASHINFER_TRTLLM = "flashinfer-trtllm" - FLASHINFER_CUDNN = "flashinfer-cudnn" - FBGEMM = "fbgemm" - EMULATION = "emulation" - - -NVFP4_LINEAR_BACKENDS = list(NvFp4LinearBackend) - - -def is_backend_supported(backend: NvFp4LinearBackend) -> tuple[bool, str | None]: - reason = None - supported = True - - if backend == NvFp4LinearBackend.FLASHINFER_CUTLASS: - # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both - # quantization and GEMM) were compiled for the current SM version. - # FlashInfer backends still rely on the vLLM quantization kernels, - # so we gate them on the same check. - supported = ( - cutlass_fp4_supported() - and current_platform.has_device_capability(100) - and has_flashinfer() - ) - - if not supported: - reason = "FlashInfer is required, >=sm_100 is required" - elif backend == NvFp4LinearBackend.VLLM_CUTLASS: - supported = cutlass_fp4_supported() - if not supported: - reason = "Cutlass is required" - elif backend == NvFp4LinearBackend.MARLIN: - supported = is_fp4_marlin_supported() - if not supported: - reason = "Marlin is required" - elif backend in [ - NvFp4LinearBackend.FLASHINFER_TRTLLM, - NvFp4LinearBackend.FLASHINFER_CUDNN, - ]: - supported = has_flashinfer() - if not supported: - reason = "FlashInfer is required" - elif backend == NvFp4LinearBackend.FBGEMM: - supported = has_fbgemm_gpu() - if not supported: - reason = "fbgemm_gpu is required" - elif backend == NvFp4LinearBackend.EMULATION: - # e.g. AMD Instinct does not support native NVFP4. - unsupported_reasons = {} - for other_backend in NVFP4_LINEAR_BACKENDS: - if other_backend == NvFp4LinearBackend.EMULATION: - continue - other_supported, other_reason = is_backend_supported(other_backend) - if not other_supported: - unsupported_reasons[other_backend] = other_reason - - if unsupported_reasons: - unsupported_reasons_str = "\n - ".join( - [f"{b.value}: {r}" for b, r in unsupported_reasons.items()] - ) - logger.warning_once( - f"NVFP4 linear falling back to the slow and unoptimized " - f"backend=NvFp4LinearBackend.EMULATION as no optimized backend is " - f"available (unavailable reasons:\n - {unsupported_reasons_str}\n). " - "In case you expect one of these backend to be used, " - "please verify your environment." - ) - - return supported, reason - - -def select_nvfp4_linear_backend() -> NvFp4LinearBackend: - """ - Select the best available NVFP4 GEMM backend based on environment - configuration and platform capabilities. - """ - if envs.VLLM_BATCH_INVARIANT: - logger.info_once( - "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the emulation " - "backend for deterministic execution." - ) - return NvFp4LinearBackend.EMULATION - - selected_backend: NvFp4LinearBackend | None = None - - if envs.VLLM_USE_FBGEMM: - try: - import fbgemm_gpu # noqa: F401 - except ImportError as exc: - raise ImportError( - "Backend fbgemm requires fbgemm.f4f4bf16 operator, " - "Please install with: pip install fbgemm-gpu-genai" - ) from exc - selected_backend = NvFp4LinearBackend.FBGEMM - elif envs.VLLM_USE_NVFP4_CT_EMULATIONS: - selected_backend = NvFp4LinearBackend.EMULATION - elif envs.VLLM_NVFP4_GEMM_BACKEND is None: - for backend in NVFP4_LINEAR_BACKENDS: - supported, reason = is_backend_supported(backend) - if supported: - selected_backend = backend - break - else: - selected_backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND) - - if selected_backend is None: - raise ValueError( - f"No NVFP4 GEMM backend selected, " - f"available backends: {NVFP4_LINEAR_BACKENDS}" - ) - - supported, reason = is_backend_supported(selected_backend) - - if not supported: - raise ValueError( - f"The selected backend={selected_backend} is not supported in current " - f"environment. Reason: {reason}. Current environment: " - f"{envs.VLLM_USE_FBGEMM=}, {envs.VLLM_USE_NVFP4_CT_EMULATIONS=}, " - f"{envs.VLLM_NVFP4_GEMM_BACKEND}." - ) - - logger.info_once(f"Using {selected_backend} for NVFP4 GEMM") - return selected_backend - - -def prepare_weights_for_nvfp4_flashinfer_trtllm( - weight: torch.Tensor, - weight_scale: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare weights and scales for FlashInfer TRTLLM FP4 GEMM.""" - from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a - - epilogue_tile_m = 128 - shuffled_weight = shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m) - shuffled_weight_scale = ( - shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m) - .reshape(weight_scale.shape) - .view(torch.float8_e4m3fn) - ) - - return shuffled_weight, shuffled_weight_scale - - -def prepare_weights_for_nvfp4_cutlass( - weight: torch.Tensor, - weight_scale: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, int]: - """ - Prepare weights and scales for CUTLASS/FlashInfer-CUTLASS FP4 GEMM. - This involves padding weights for alignment (K and N divisible by 32) - """ - swizzled_weight_scale = swizzle_blockscale(weight_scale) - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(weight) - return padded_weight, swizzled_weight_scale, weights_padding_cols - - -def prepare_weights_for_nvfp4_fbgemm( - weight: torch.Tensor, - weight_scale: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - """Prepare weights and scales for FBGEMM FP4 GEMM.""" - swizzled_weight_scale = swizzle_blockscale(weight_scale) - swizzled_weight_scale = swizzled_weight_scale.view(-1).view(torch.uint8) - return weight, swizzled_weight_scale - - -def convert_to_nvfp4_linear_kernel_format( - backend: NvFp4LinearBackend, - layer: torch.nn.Module, -) -> None: - """Convert layer to NVFP4 linear kernel format.""" - - assert layer.weight_scale.dtype == torch.float8_e4m3fn, ( - "Weight Block scale must be represented as FP8-E4M3" - ) - - # Default to no padding - layer.weights_padding_cols = 0 - - if backend == NvFp4LinearBackend.MARLIN: - logger.warning_once( - "Your GPU does not have native support for FP4 computation but " - "FP4 quantization is being used. Weight-only FP4 compression " - "will be used leveraging the Marlin kernel. This may degrade " - "performance for compute-heavy workloads." - ) - prepare_fp4_layer_for_marlin(layer) - elif backend == NvFp4LinearBackend.FLASHINFER_TRTLLM: - weight, weight_scale = prepare_weights_for_nvfp4_flashinfer_trtllm( - layer.weight.data, layer.weight_scale.data - ) - layer.weight = torch.nn.Parameter(weight, requires_grad=False) - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) - elif backend == NvFp4LinearBackend.FBGEMM: - weight, weight_scale = prepare_weights_for_nvfp4_fbgemm( - layer.weight.data, layer.weight_scale.data - ) - layer.weight = torch.nn.Parameter(weight, requires_grad=False) - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) - elif backend in ( - NvFp4LinearBackend.VLLM_CUTLASS, - NvFp4LinearBackend.FLASHINFER_CUTLASS, - NvFp4LinearBackend.FLASHINFER_CUDNN, - ): - weight, weight_scale, weights_padding_cols = prepare_weights_for_nvfp4_cutlass( - layer.weight.data, layer.weight_scale.data - ) - layer.weight = torch.nn.Parameter(weight, requires_grad=False) - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) - layer.weights_padding_cols = weights_padding_cols - elif backend == NvFp4LinearBackend.EMULATION: - # We can not call `.to(device)` during cuda graph capture - do it here instead. - # (operation not permitted when stream is capturing) - kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device) - - -def apply_nvfp4_linear( - backend: NvFp4LinearBackend, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - swizzle: bool | None = None, -) -> torch.Tensor: - """ - Apply NVFP4 linear transformation using the specified backend. - """ - weight = layer.weight - weight_scale = layer.weight_scale - weight_global_scale = layer.weight_global_scale - input_global_scale_inv = layer.input_global_scale_inv - alpha = layer.alpha - output_size = layer.output_size_per_partition - input_size = layer.input_size_per_partition - output_dtype = x.dtype - output_shape = [*x.shape[:-1], output_size] - - if backend == NvFp4LinearBackend.MARLIN: - return apply_fp4_marlin_linear( - input=x, - weight=weight, - weight_scale=weight_scale, - weight_global_scale=weight_global_scale, - workspace=layer.workspace, - size_n=output_size, - size_k=input_size, - bias=bias, - ) - elif backend == NvFp4LinearBackend.EMULATION: - x_2d = x.reshape(-1, x.shape[-1]) - out = run_nvfp4_emulations( - x=x_2d, - input_global_scale=input_global_scale_inv, - weight=weight, - weight_scale_swizzled=weight_scale, - weight_global_scale=weight_global_scale, - swizzle=swizzle, - ) - out = out[:, :output_size] - if bias is not None: - out = out + bias - return out.view(*output_shape) - - # Quantize BF16 or FP16 to (FP4 and interleaved block scale) - x_fp4, x_blockscale = scaled_fp4_quant( - x, input_global_scale_inv, is_sf_swizzled_layout=True, backend=backend.value - ) - - # Validate dtypes - assert x_fp4.dtype == torch.uint8 - assert weight.dtype == torch.uint8 - assert x_blockscale.dtype == torch.float8_e4m3fn - # weight_scale is fp8 for most backends, but uint8 for fbgemm - assert weight_scale.dtype in (torch.float8_e4m3fn, torch.uint8) - assert alpha.dtype == torch.float32 - - # Pad activations to match weight K-dimension padding - weights_padding_cols = getattr(layer, "weights_padding_cols", 0) - x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols) - - # Prepare args for the matmul - mm_args = ( - x_fp4, - weight, - x_blockscale, - weight_scale, - alpha, - output_dtype, - ) - - # Call the appropriate backend - if backend.value.startswith("flashinfer-"): - backend_name = backend.value[len("flashinfer-") :] - out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name) - elif backend == NvFp4LinearBackend.FBGEMM: - out = torch.ops.fbgemm.f4f4bf16( - x_fp4, - weight, - x_blockscale.view(-1).view(torch.uint8), - weight_scale, - alpha, - use_mx=False, - ).to(output_dtype) - else: - assert backend == NvFp4LinearBackend.VLLM_CUTLASS - out = cutlass_scaled_fp4_mm(*mm_args) - - # Slice output to remove N-dimension padding - out = slice_nvfp4_output(out, output_size) - - if bias is not None: - out = out + bias - - return out.view(*output_shape) - def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor: """ diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 1844b75561e..bdaa6af0945 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -11,7 +11,12 @@ from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp from vllm.platforms import current_platform from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits, has_deep_gemm -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) from vllm.v1.attention.backends.mla.indexer import ( DeepseekV32IndexerMetadata, ) @@ -30,7 +35,7 @@ RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024 def sparse_attn_indexer( hidden_states: torch.Tensor, - k_cache_prefix: str, + k_cache_prefix: LayerNameType, kv_cache: torch.Tensor, q_fp8: torch.Tensor, k: torch.Tensor, @@ -46,6 +51,7 @@ def sparse_attn_indexer( # careful! this will be None in dummy run attn_metadata = get_forward_context().attn_metadata fp8_dtype = current_platform.fp8_dtype() + k_cache_prefix = _resolve_layer_name(k_cache_prefix) # assert isinstance(attn_metadata, dict) if not isinstance(attn_metadata, dict): @@ -183,13 +189,15 @@ def sparse_attn_indexer( # TODO: move and optimize below logic with triton kernels batch_size = padded_q_fp8_decode_tokens.shape[0] next_n = padded_q_fp8_decode_tokens.shape[1] - assert batch_size == decode_metadata.seq_lens.shape[0] num_padded_tokens = batch_size * next_n + seq_lens = decode_metadata.seq_lens[:batch_size] + # seq_lens is (B, next_n) for native spec decode, (B,) otherwise. + # fp8_paged_mqa_logits and all topk kernels accept both shapes. logits = fp8_paged_mqa_logits( padded_q_fp8_decode_tokens, kv_cache, weights[:num_padded_tokens], - decode_metadata.seq_lens, + seq_lens, decode_metadata.block_table, decode_metadata.schedule_metadata, max_model_len=max_model_len, @@ -198,17 +206,6 @@ def sparse_attn_indexer( num_rows = logits.shape[0] topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if next_n == 1: - lengths = decode_metadata.seq_lens - else: - # (bs,) -> (bs, 1) + (next_n,) -> (bs, next_n) -> (bs * next_n,) - lengths = ( - decode_metadata.seq_lens.unsqueeze(1) - - next_n - + 1 - + decode_metadata.offsets - ).flatten() - if current_platform.is_cuda(): workspace_manager = current_workspace_manager() (topk_workspace,) = workspace_manager.get_simultaneous( @@ -216,7 +213,7 @@ def sparse_attn_indexer( ) torch.ops._C.persistent_topk( logits, - lengths, + seq_lens, topk_indices, topk_workspace, topk_tokens, @@ -227,7 +224,7 @@ def sparse_attn_indexer( ops.top_k_per_row_decode( logits, next_n, - decode_metadata.seq_lens, + seq_lens, topk_indices, num_rows, logits.stride(0), @@ -238,7 +235,7 @@ def sparse_attn_indexer( torch.ops._C.top_k_per_row_decode( logits, next_n, - decode_metadata.seq_lens, + seq_lens, topk_indices, num_rows, logits.stride(0), @@ -253,7 +250,7 @@ def sparse_attn_indexer( topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]), decode_lens, ) - topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = ( + topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = ( topk_indices ) @@ -262,7 +259,7 @@ def sparse_attn_indexer( def sparse_attn_indexer_fake( hidden_states: torch.Tensor, - k_cache_prefix: str, + k_cache_prefix: LayerNameType, kv_cache: torch.Tensor, q_fp8: torch.Tensor, k: torch.Tensor, @@ -351,7 +348,7 @@ class SparseAttnIndexer(CustomOp): ): return torch.ops.vllm.sparse_attn_indexer( hidden_states, - self.k_cache.prefix, + _encode_layer_name(self.k_cache.prefix), self.k_cache.kv_cache, q_fp8, k, @@ -375,7 +372,7 @@ class SparseAttnIndexer(CustomOp): if rocm_aiter_ops.is_enabled(): return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( hidden_states, - self.k_cache.prefix, + _encode_layer_name(self.k_cache.prefix), self.k_cache.kv_cache, q_fp8, k, diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index daaa86bed47..ddae01856da 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -8,13 +8,17 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter, UninitializedParameter +import vllm.envs as envs from vllm.distributed import ( divide, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) -from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.custom_op import PluggableLayer +from vllm.model_executor.layers.batch_invariant import ( + linear_batch_invariant, +) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, @@ -66,6 +70,8 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike(): + return linear_batch_invariant(x, layer.weight, bias) return dispatch_unquantized_gemm()(layer, x, layer.weight, bias) def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: @@ -182,8 +188,8 @@ def get_masked_input_and_mask( # --8<-- [start:vocab_parallel_embedding] -@CustomOp.register("vocab_parallel_embedding") -class VocabParallelEmbedding(CustomOp): +@PluggableLayer.register("vocab_parallel_embedding") +class VocabParallelEmbedding(PluggableLayer): """Embedding parallelized in the vocabulary dimension. Adapted from torch.nn.Embedding, note that we pad the vocabulary size to @@ -461,7 +467,7 @@ class VocabParallelEmbedding(CustomOp): param[: loaded_weight.shape[0]].data.copy_(loaded_weight) param[loaded_weight.shape[0] :].data.fill_(0) - def forward_native(self, input_): + def forward(self, input_): if self.tp_size > 1: # Build the mask. masked_input, input_mask = get_masked_input_and_mask( @@ -483,9 +489,6 @@ class VocabParallelEmbedding(CustomOp): output = tensor_model_parallel_all_reduce(output_parallel) return output - def forward_cuda(self, input_): - return self.forward_native(input_) - def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings_per_partition}" s += f", embedding_dim={self.embedding_dim}" @@ -496,7 +499,7 @@ class VocabParallelEmbedding(CustomOp): # --8<-- [start:parallel_lm_head] -@CustomOp.register("parallel_lm_head") +@PluggableLayer.register("parallel_lm_head") class ParallelLMHead(VocabParallelEmbedding): """Parallelized LM head. diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 8f370717d81..4ee9e90d741 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -175,7 +175,7 @@ _MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]() def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], str]: from vllm.model_executor.models.adapters import as_embedding_model, as_seq_cls_model - architectures = getattr(model_config.hf_config, "architectures", []) + architectures = getattr(model_config.hf_config, "architectures", None) or [] model_cls, arch = model_config.registry.resolve_model_cls( architectures, @@ -215,7 +215,7 @@ def get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], model_config.runner_type, model_config.trust_remote_code, model_config.model_impl, - tuple(getattr(model_config.hf_config, "architectures", [])), + tuple(getattr(model_config.hf_config, "architectures", None) or []), ) ) if key in _MODEL_ARCH_BY_HASH: diff --git a/vllm/model_executor/model_loader/weight_utils.py b/vllm/model_executor/model_loader/weight_utils.py index e3689dcd869..3b961e8e143 100644 --- a/vllm/model_executor/model_loader/weight_utils.py +++ b/vllm/model_executor/model_loader/weight_utils.py @@ -763,31 +763,24 @@ def np_cache_weights_iterator( yield name, torch.from_numpy(param) -def _checkpoints_fit_in_ram(files: list[str], threshold: float = 0.9) -> bool: - """Return True if total size of *files* fits within *threshold* of available RAM.""" +def _get_checkpoints_size_bytes(files: list[str]) -> int: + """Return the total size of the checkpoint files in bytes.""" if not files: - return True + return 0 + return sum(os.path.getsize(f) for f in files) + + +def _get_available_ram_bytes() -> int: + """Return the available RAM in bytes.""" import psutil - total_size = sum(os.path.getsize(f) for f in files) - available_ram = psutil.virtual_memory().available - fits = total_size <= threshold * available_ram - if not fits: - logger.warning( - "NFS detected but checkpoint total size (%.2f GiB) exceeds " - "%.0f%% of available RAM (%.2f GiB). Skipping prefetching checkpoints.", - total_size / (1024**3), - threshold * 100, - available_ram / (1024**3), - ) - return fits + return psutil.virtual_memory().available -def _is_nfs_path(files: list[str]) -> bool: - """Check whether the first file in *files* resides on an NFS - filesystem (Linux only).""" +def _get_fs_type(files: list[str]) -> str: + """Get the filesystem type of the first file in *files* (Linux only).""" if not files: - return False + return "" try: # Only the first file is checked — all checkpoint shards reside # in the same directory and therefore on the same filesystem. @@ -810,12 +803,11 @@ def _is_nfs_path(files: list[str]) -> bool: ) and len(mount_point) > len(best_mount): best_mount = mount_point best_fstype = fstype - return best_fstype in ("nfs", "nfs4") + return best_fstype except Exception: # /proc/mounts is Linux-specific; on other OSes (or if the read - # fails for any reason) we fall back to "not NFS" rather than - # crashing model loading. - return False + # fails for any reason) we fall back to an empty string. + return "" def _prefetch_checkpoint(file_path: str) -> None: @@ -901,11 +893,63 @@ def safetensors_weights_iterator( sorted_files = sorted(hf_weights_files, key=_natural_sort_key) - should_prefetch = safetensors_load_strategy == "prefetch" or ( - safetensors_load_strategy is None - and _is_nfs_path(sorted_files) - and _checkpoints_fit_in_ram(sorted_files) + fs_type = _get_fs_type(sorted_files) + is_net_fs = fs_type in ("nfs", "nfs4", "lustre") + total_bytes = _get_checkpoints_size_bytes(sorted_files) + avail_bytes = _get_available_ram_bytes() + ram_threshold_pct = 90 + fits_in_ram = total_bytes <= (ram_threshold_pct / 100.0) * avail_bytes + fs_name = fs_type.upper() if fs_type else "unknown" + + logger.info_once( + "Filesystem type for checkpoints: %s. Checkpoint size: %.2f GiB. " + "Available RAM: %.2f GiB.", + fs_name, + total_bytes / 1024**3, + avail_bytes / 1024**3, ) + + should_prefetch = safetensors_load_strategy == "prefetch" + if safetensors_load_strategy is None: + if is_net_fs and fits_in_ram: + should_prefetch = True + elif is_net_fs and not fits_in_ram: + logger.warning_once( + "Network filesystem (%s) detected but checkpoint total size " + "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB). " + "Skipping auto-prefetch.", + fs_name, + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + elif not is_net_fs and fits_in_ram: + logger.info_once( + "Auto-prefetch is disabled because the filesystem (%s) is not a " + "recognized network FS (NFS/Lustre). If you want to force " + "prefetching, start vLLM with --safetensors-load-strategy=prefetch.", + fs_name, + ) + elif not is_net_fs and not fits_in_ram: + logger.info_once( + "Auto-prefetch is disabled because the filesystem (%s) is not a " + "recognized network FS (NFS/Lustre) and the checkpoint size " + "(%.2f GiB) exceeds %d%% of available RAM (%.2f GiB).", + fs_name, + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + elif should_prefetch and not fits_in_ram: + logger.warning_once( + "safetensors_load_strategy='prefetch' was explicitly specified, but " + "checkpoint total size (%.2f GiB) exceeds %d%% of available RAM " + "(%.2f GiB). This may cause out-of-memory errors.", + total_bytes / 1024**3, + ram_threshold_pct, + avail_bytes / 1024**3, + ) + if should_prefetch: _prefetch_all_checkpoints(sorted_files) diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py index 467e8ab67bf..cccd849681f 100644 --- a/vllm/model_executor/models/adapters.py +++ b/vllm/model_executor/models/adapters.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import itertools from collections.abc import Iterable from contextlib import contextmanager from typing import TYPE_CHECKING, Any, TypeVar, cast @@ -181,7 +182,8 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T: seen_weights = list[tuple[str, torch.Tensor]]() for name, loaded_weight in weights: - seen_weights.append((name, loaded_weight)) + # Clone because the iterator may reuse the tensor buffer + seen_weights.append((name, loaded_weight.clone())) try: target_prefix = next( @@ -208,9 +210,11 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T: self._get_name(), ) + # Lazy chain so buffer-reusing weight iterators (e.g. + # runai_streamer) are consumed one tensor at a time. mapped_weights = ( (target_prefix + name, weight) - for name, weight in (*seen_weights, *weights) + for name, weight in itertools.chain(seen_weights, weights) ) def default_load_weights(weights): diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index 1cebea56a13..42206c11cb9 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -2007,6 +2007,7 @@ class CohereAsrForConditionalGeneration( supports_transcription_only = True supported_languages = ISO639_1_SUPPORTED_LANGS skip_warmup_audio_preprocessing = True + no_space_languages = {"ja", "zh"} @classmethod def validate_language(cls, language: str | None) -> str | None: diff --git a/vllm/model_executor/models/colmodernvbert.py b/vllm/model_executor/models/colmodernvbert.py index 1e8477e120e..b16bb4e221d 100644 --- a/vllm/model_executor/models/colmodernvbert.py +++ b/vllm/model_executor/models/colmodernvbert.py @@ -18,7 +18,6 @@ from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, @@ -358,70 +357,23 @@ class ColModernVBertForRetrieval( "model.text_model.layers.": "text_layers.", "model.text_model.embeddings.": "text_embeddings.", "model.text_model.final_norm.": "text_final_norm.", - "model.connector.modality_projection.": "connector.", + "model.connector.modality_projection.": "connector.proj.", "model.custom_text_proj.": "custom_text_proj.", - "model.vision_model.": "vision_model.vision_model.", + "model.vision_model.vision_model.": "vision_model.vision_model.", "model.": "", }, ) - # Checkpoint names for DecoupledEmbedding parts - _BASE_EMB = "model.text_model.embeddings.tok_embeddings.weight" - _EXTRA_EMB = ( - "model.text_model.embeddings.tok_embeddings.additional_embedding.weight" - ) - def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], ) -> set[str]: - # DecoupledEmbedding requires concatenating base + additional - # embedding tensors before loading, so we extract them first. - base_embedding_weight: torch.Tensor | None = None - additional_embedding_weight: torch.Tensor | None = None - remaining: list[tuple[str, torch.Tensor]] = [] - - for name, tensor in weights: - if name == self._BASE_EMB: - base_embedding_weight = tensor - elif name == self._EXTRA_EMB: - additional_embedding_weight = tensor - else: - remaining.append((name, tensor)) - - # Load all non-embedding weights via AutoWeightsLoader loader = AutoWeightsLoader(self) loaded_params = loader.load_weights( - remaining, + weights, mapper=self.hf_to_vllm_mapper, ) - # Concatenate and load DecoupledEmbedding weights - if base_embedding_weight is not None: - combined = base_embedding_weight - if additional_embedding_weight is not None: - combined = torch.cat( - [base_embedding_weight, additional_embedding_weight], - dim=0, - ) - param_name = "text_embeddings.tok_embeddings.weight" - params_dict = dict(self.named_parameters()) - if param_name in params_dict: - param = params_dict[param_name] - weight_loader = getattr( - param, - "weight_loader", - default_weight_loader, - ) - weight_loader(param, combined) - loaded_params.add(param_name) - elif additional_embedding_weight is not None: - raise ValueError( - "Found 'text_model.embeddings.tok_embeddings" - ".additional_embedding.weight' but not " - "'text_model.embeddings.tok_embeddings.weight'" - ) - # The pooler wraps ``custom_text_proj`` as its head projector. # Mark those params as loaded under the pooler path too. if hasattr(self, "pooler") and hasattr(self.pooler, "head"): diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 22d300a7ebf..521184e4c68 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -108,6 +108,23 @@ class Gemma4Config(VerifyAndUpdateConfig): class GptOssForCausalLMConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + quant_config = getattr(model_config.hf_config, "quantization_config", None) + if quant_config is not None and quant_config.get("quant_method") == "mxfp4": + model_config.hf_config.quantization_config["quant_method"] = "gpt_oss_mxfp4" + + hf_text_quant_config = getattr( + model_config.hf_text_config, "quantization_config", None + ) + if ( + hf_text_quant_config is not None + and hf_text_quant_config.get("quant_method") == "mxfp4" + ): + model_config.hf_text_config.quantization_config["quant_method"] = ( + "gpt_oss_mxfp4" + ) + @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: structured_outputs_config = vllm_config.structured_outputs_config @@ -192,6 +209,12 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig): pooler_config.use_activation = False +class JinaForRankingConfig(VerifyAndUpdateConfig): + @staticmethod + def verify_and_update_model_config(model_config: "ModelConfig") -> None: + model_config.hf_config.embedding_size = 512 + + class JinaRobertaModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -226,8 +249,8 @@ class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig): config = model_config.hf_config config.num_labels = 1 pooler_config = model_config.pooler_config - if pooler_config.logit_bias is None: - pooler_config.logit_bias = 2.65 + if pooler_config.logit_mean is None: + pooler_config.logit_mean = 2.65 class LlamaBidirectionalConfig(VerifyAndUpdateConfig): @@ -612,6 +635,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "GteNewForSequenceClassification": GteNewModelConfig, "GteNewModel": GteNewModelConfig, "JambaForSequenceClassification": JambaForSequenceClassificationConfig, + "JinaForRanking": JinaForRankingConfig, "JinaVLForRanking": JinaVLForSequenceClassificationConfig, "LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig, "LlamaBidirectionalModel": LlamaBidirectionalConfig, diff --git a/vllm/model_executor/models/conformer_encoder.py b/vllm/model_executor/models/conformer_encoder.py new file mode 100644 index 00000000000..0d2e3127019 --- /dev/null +++ b/vllm/model_executor/models/conformer_encoder.py @@ -0,0 +1,350 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shared Conformer encoder components for FireRedASR2 and FireRedLID. + +Both models use the same Conformer-based audio encoder architecture +(Conv2dSubsampling → RelPositionalEncoding → N × RelPosEmbConformerBlock). +This module factors out the common building blocks to avoid duplication. +""" + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.linear import ReplicatedLinear + + +class Conv2dSubsampling(nn.Module): + def __init__(self, idim: int, d_model: int, out_channels: int = 32): + super().__init__() + self.conv = nn.Sequential( + nn.Conv2d(1, out_channels, 3, 2), + nn.ReLU(), + nn.Conv2d(out_channels, out_channels, 3, 2), + nn.ReLU(), + ) + subsample_idim = ((idim - 1) // 2 - 1) // 2 + self.out = ReplicatedLinear( + input_size=out_channels * subsample_idim, + output_size=d_model, + bias=True, + ) + + self.subsampling = 4 + left_context = right_context = 3 # both exclude current frame + self.context = left_context + 1 + right_context # 7 + + def forward( + self, x: torch.Tensor, x_mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = x.unsqueeze(1) + x = self.conv(x) + N, C, T, D = x.size() + x, _ = self.out(x.transpose(1, 2).contiguous().view(N, T, C * D)) + mask = x_mask[:, :, :-2:2][:, :, :-2:2] + input_lengths = mask[:, -1, :].sum(dim=-1) + return x, input_lengths, mask + + +class Swish(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(x) + + +class RelPositionalEncoding(nn.Module): + def __init__(self, d_model: int, max_len: int = 5000): + super().__init__() + pe_positive = torch.zeros(max_len, d_model, requires_grad=False) + pe_negative = torch.zeros(max_len, d_model, requires_grad=False) + position = torch.arange(0, max_len).unsqueeze(1).float() + div_term = torch.exp( + torch.arange(0, d_model, 2).float() + * -(torch.log(torch.tensor(10000.0)).item() / d_model) + ) + pe_positive[:, 0::2] = torch.sin(position * div_term) + pe_positive[:, 1::2] = torch.cos(position * div_term) + pe_negative[:, 0::2] = torch.sin(-1 * position * div_term) + pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) + + pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0) + pe_negative = pe_negative[1:].unsqueeze(0) + self.pe = torch.cat([pe_positive, pe_negative], dim=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Tmax = 2 * max_len - 1 + Tmax, T = self.pe.size(1), x.size(1) + pos_emb = self.pe[:, Tmax // 2 - T + 1 : Tmax // 2 + T].clone().detach() + return pos_emb + + +class ConformerFeedForward(nn.Module): + def __init__(self, d_model: int): + super().__init__() + self.pre_layer_norm = nn.LayerNorm(d_model) + self.linear_expand = ReplicatedLinear( + input_size=d_model, + output_size=d_model * 4, + bias=True, + ) + self.nonlinear = Swish() + self.linear_project = ReplicatedLinear( + input_size=d_model * 4, + output_size=d_model, + bias=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.pre_layer_norm(x) + x, _ = self.linear_expand(x) + x = self.nonlinear(x) + x, _ = self.linear_project(x) + return x + residual + + +class EncoderMultiHeadAttention(nn.Module): + def __init__(self, n_head: int, d_model: int): + super().__init__() + assert d_model % n_head == 0 + self.n_head = n_head + self.d_k = d_model // n_head + self.d_v = self.d_k + + self.w_qs = ReplicatedLinear(d_model, n_head * self.d_k, bias=False) + self.w_ks = ReplicatedLinear(d_model, n_head * self.d_k, bias=False) + self.w_vs = ReplicatedLinear(d_model, n_head * self.d_v, bias=False) + + self.layer_norm_q = nn.LayerNorm(d_model) + self.layer_norm_k = nn.LayerNorm(d_model) + self.layer_norm_v = nn.LayerNorm(d_model) + + self.fc = ReplicatedLinear(n_head * self.d_v, d_model, bias=False) + + def forward_qkv( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + d_k, d_v, n_head = self.d_k, self.d_v, self.n_head + sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) + + q = self.layer_norm_q(q) + k = self.layer_norm_k(k) + v = self.layer_norm_v(v) + + q = self.w_qs(q)[0].view(sz_b, len_q, n_head, d_k) + k = self.w_ks(k)[0].view(sz_b, len_k, n_head, d_k) + v = self.w_vs(v)[0].view(sz_b, len_v, n_head, d_v) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + return q, k, v + + def forward_output( + self, + output: torch.Tensor, + residual: torch.Tensor, + sz_b: int, + len_q: int, + ) -> torch.Tensor: + output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1) + fc_out, _ = self.fc(output) + return fc_out + residual + + def forward_attention( + self, + attn: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if mask is not None: + mask = mask.unsqueeze(1) + mask = mask.eq(0) + attn = attn.masked_fill(mask, -float("inf")) + attn = torch.softmax(attn, dim=-1).masked_fill(mask, 0.0) + else: + attn = torch.softmax(attn, dim=-1) + output = torch.matmul(attn, v) + return output, attn + + +class RelPosMultiHeadAttention(EncoderMultiHeadAttention): + def __init__(self, n_head: int, d_model: int): + super().__init__(n_head, d_model) + d_k = d_model // n_head + self.scale = 1.0 / (d_k**0.5) + self.linear_pos = ReplicatedLinear(d_model, n_head * d_k, bias=False) + self.pos_bias_u = nn.Parameter(torch.empty([n_head, d_k])) + self.pos_bias_v = nn.Parameter(torch.empty([n_head, d_k])) + + def _rel_shift(self, x): + N, H, T1, T2 = x.size() + zero_pad = torch.zeros((N, H, T1, 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + x_padded = x_padded.view(N, H, T2 + 1, T1) + x = x_padded[:, :, 1:].view_as(x) + x = x[:, :, :, : x.size(-1) // 2 + 1] + return x + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + pos_emb: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + sz_b, len_q = q.size(0), q.size(1) + residual = q + q, k, v = self.forward_qkv(q, k, v) + + q = q.transpose(1, 2) + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb)[0].view(n_batch_pos, -1, self.n_head, self.d_k) + p = p.transpose(1, 2) + + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self._rel_shift(matrix_bd) + + attn_scores = matrix_ac + matrix_bd + attn_scores.mul_(self.scale) + + output, attn = self.forward_attention(attn_scores, v, mask=mask) + output = self.forward_output(output, residual, sz_b, len_q) + return output, attn + + +class ConformerConvolution(nn.Module): + def __init__(self, d_model: int, kernel_size: int = 33): + super().__init__() + assert kernel_size % 2 == 1 + self.pre_layer_norm = nn.LayerNorm(d_model) + self.pointwise_conv1 = nn.Conv1d( + d_model, d_model * 4, kernel_size=1, bias=False + ) + self.padding = (kernel_size - 1) // 2 + self.depthwise_conv = nn.Conv1d( + d_model * 2, + d_model * 2, + kernel_size, + stride=1, + padding=self.padding, + groups=d_model * 2, + bias=False, + ) + self.batch_norm = nn.LayerNorm(d_model * 2) + self.swish = Swish() + self.pointwise_conv2 = nn.Conv1d( + d_model * 2, d_model, kernel_size=1, bias=False + ) + + def forward( + self, x: torch.Tensor, mask: torch.Tensor | None = None + ) -> torch.Tensor: + residual = x + out = self.pre_layer_norm(x) + out = out.transpose(1, 2) + if mask is not None: + out.masked_fill_(mask.ne(1), 0.0) + out = self.pointwise_conv1(out) + out = F.glu(out, dim=1) + out = self.depthwise_conv(out) + out = out.transpose(1, 2) + out = self.swish(self.batch_norm(out)) + out = out.transpose(1, 2) + out = self.pointwise_conv2(out) + if mask is not None: + out.masked_fill_(mask.ne(1), 0.0) + out = out.transpose(1, 2) + return out + residual + + +class RelPosEmbConformerBlock(nn.Module): + def __init__(self, d_model: int, n_head: int, kernel_size: int = 33): + super().__init__() + self.ffn1 = ConformerFeedForward(d_model) + self.mhsa = RelPosMultiHeadAttention(n_head, d_model) + self.conv = ConformerConvolution(d_model, kernel_size) + self.ffn2 = ConformerFeedForward(d_model) + self.layer_norm = nn.LayerNorm(d_model) + + def forward( + self, + x: torch.Tensor, + pos_emb: torch.Tensor, + slf_attn_mask: torch.Tensor | None = None, + pad_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + out = 0.5 * x + 0.5 * self.ffn1(x) + out = self.mhsa(out, out, out, pos_emb, mask=slf_attn_mask)[0] + out = self.conv(out, pad_mask) + out = 0.5 * out + 0.5 * self.ffn2(out) + out = self.layer_norm(out) + return out + + +class ConformerEncoder(nn.Module): + """ + Conformer encoder shared by FireRedASR2 and FireRedLID. + """ + + def __init__( + self, + idim: int, + n_layers_enc: int, + n_head: int, + d_model: int, + kernel_size: int = 33, + pe_maxlen: int = 5000, + ): + super().__init__() + self.odim = d_model + + self.input_preprocessor = Conv2dSubsampling(idim, d_model) + self.positional_encoding = RelPositionalEncoding(d_model, max_len=pe_maxlen) + + self.layer_stack = nn.ModuleList() + for _ in range(n_layers_enc): + block = RelPosEmbConformerBlock(d_model, n_head, kernel_size) + self.layer_stack.append(block) + + def forward( + self, + padded_input: torch.Tensor, + input_lengths: torch.Tensor, + pad: bool = True, + ): + if pad: + padded_input = F.pad( + padded_input, + (0, 0, 0, self.input_preprocessor.context - 1), + "constant", + 0.0, + ) + src_mask = self.padding_position_is_0(padded_input, input_lengths) + + embed_output, input_lengths, src_mask = self.input_preprocessor( + padded_input, src_mask + ) + enc_output = embed_output + + pos_emb = self.positional_encoding(embed_output) + + for enc_layer in self.layer_stack: + enc_output = enc_layer( + enc_output, pos_emb, slf_attn_mask=src_mask, pad_mask=src_mask + ) + + return enc_output, input_lengths, src_mask + + def padding_position_is_0( + self, padded_input: torch.Tensor, input_lengths: torch.Tensor + ) -> torch.Tensor: + N, T = padded_input.size()[:2] + # Use broadcasting instead of a Python loop for efficiency. + positions = torch.arange(T, device=padded_input.device).unsqueeze(0) + mask = (positions < input_lengths.unsqueeze(1)).to(torch.uint8) + return mask.unsqueeze(1) diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index 126efb6f88e..a66ec7aa3e6 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -30,6 +30,7 @@ from .deepseek_v2 import ( DeepseekV2DecoderLayer, DeepseekV2MixtureOfExperts, DeepseekV2MoE, + _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) from .utils import maybe_prefix @@ -190,10 +191,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ) # Set MoE hyperparameters self.set_moe_parameters() - self.is_fp4_ckpt = ( - self.quant_config is not None - and self.quant_config.get_name() == "modelopt_fp4" - ) def set_moe_parameters(self): self.expert_weights = [] @@ -248,13 +245,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1), ] - if self.is_fp4_ckpt: - # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) - indexer_fused_mapping = [ - ("wk_weights_proj", "wk", 0), - ("wk_weights_proj", "weights_proj", 1), - ] - stacked_params_mapping.extend(indexer_fused_mapping) + # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) + indexer_fused_mapping = [ + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + stacked_params_mapping.extend(indexer_fused_mapping) expert_params_mapping = SharedFusedMoE.make_expert_params_mapping( self, @@ -271,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + _pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue @@ -281,6 +278,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts): rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) ) name = self._rewrite_spec_layer_name(spec_layer, name) + + if _try_load_fp8_indexer_wk( + name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params + ): + continue + for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 17ddd5edece..cd28fb0192f 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -66,6 +66,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.quantization.utils.fp8_utils import ( per_token_group_quant_fp8, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + scaled_dequantize, +) from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.sparse_attn_indexer import ( SparseAttnIndexer, @@ -628,10 +632,6 @@ class Indexer(nn.Module): self.vllm_config = vllm_config self.config = config self.quant_config = quant_config - self.is_fp4_ckpt = ( - self.quant_config is not None - and self.quant_config.get_name() == "modelopt_fp4" - ) # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] self.topk_tokens = config.index_topk self.n_head = config.index_n_heads # 64 @@ -646,36 +646,16 @@ class Indexer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.wq_b", ) - if self.is_fp4_ckpt: - # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. - # weights_proj does not get quantized, - # so we run both with quant_config=None - # wk may be upcasted from the default quant; - # experiments show fusion is always faster unless WK proj is in FP4, - # which is not the case for all known quants. - self.wk_weights_proj = MergedColumnParallelLinear( - hidden_size, - [self.head_dim, self.n_head], - bias=False, - quant_config=None, - disable_tp=True, - prefix=f"{prefix}.wk_weights_proj", - ) - else: - self.wk = ReplicatedLinear( - hidden_size, - self.head_dim, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.wk", - ) - self.weights_proj = ReplicatedLinear( - hidden_size, - self.n_head, - bias=False, - quant_config=None, - prefix=f"{prefix}.weights_proj", - ) + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # FP8 wk weights are upcasted to BF16 during loading to maintain fusion. + self.wk_weights_proj = MergedColumnParallelLinear( + hidden_size, + [self.head_dim, self.n_head], + bias=False, + quant_config=None, + disable_tp=True, + prefix=f"{prefix}.wk_weights_proj", + ) self.k_norm = LayerNorm(self.head_dim, eps=1e-6) self.softmax_scale = self.head_dim**-0.5 @@ -716,14 +696,10 @@ class Indexer(nn.Module): q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - if self.is_fp4_ckpt: - # Fused wk + weights_proj: one GEMM, then split - kw, _ = self.wk_weights_proj(hidden_states) - k = kw[:, : self.head_dim] - weights = kw[:, self.head_dim :] - else: - k, _ = self.wk(hidden_states) - weights, _ = self.weights_proj(hidden_states) + # Fused wk + weights_proj: one GEMM, then split + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + weights = kw[:, self.head_dim :] k = self.k_norm(k) k_pe, k_nope = torch.split( @@ -761,6 +737,46 @@ class Indexer(nn.Module): return self.indexer_op(hidden_states, q_fp8, k, weights) +def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params): + """ + We fuse the WK and weights_proj projections, but in some checkpoints WK is stored + in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16. + Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK + weights and scale, and when both are available, dequantizes to BF16 and stores into + the fused wk_weights_proj.weight parameter. + """ + if "indexer.wk." not in name or "wk_weights" in name: + return False # Weight is not an isolated WK weight for the indexer, ignore. + is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn + is_scale = "weight_scale_inv" in name + if not is_weight and not is_scale: + return False # WK is not in FP8 format, ignore. + # Buffer this tensor (weight or scale) until both have arrived. + layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer" + entry = buf.setdefault(layer_prefix, {}) + entry["weight" if is_weight else "scale"] = tensor + if "weight" not in entry or "scale" not in entry: + return True # still waiting for the other param + + # We have both weight and scale: dequantize FP8 to BF16. + weight_fp8, scale_inv = entry["weight"], entry["scale"] + del buf[layer_prefix] + block_size = weight_fp8.shape[1] // scale_inv.shape[1] + weight_bf16 = scaled_dequantize( + weight_fp8, + scale_inv, + group_shape=GroupShape(block_size, block_size), + out_dtype=torch.bfloat16, + ) + + # Load the dequantized weight into shard 0 of the fused buffer. + fused_name = f"{layer_prefix}.wk_weights_proj.weight" + param = params_dict[fused_name] + param.weight_loader(param, weight_bf16, 0) + loaded_params.add(fused_name) + return True + + def _min_latency_fused_qkv_a_proj_impl( input_: torch.Tensor, weight: torch.Tensor, @@ -1344,10 +1360,6 @@ class DeepseekV2ForCausalLM( quant_config = vllm_config.quant_config self.config = config self.quant_config = quant_config - self.is_fp4_ckpt = ( - self.quant_config is not None - and self.quant_config.get_name() == "modelopt_fp4" - ) qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0) qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0) @@ -1473,13 +1485,13 @@ class DeepseekV2ForCausalLM( ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ] - if self.is_fp4_ckpt: - # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) - indexer_fused_mapping = [ - ("wk_weights_proj", "wk", 0), - ("wk_weights_proj", "weights_proj", 1), - ] - stacked_params_mapping.extend(indexer_fused_mapping) + # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj) + _pending_wk_fp8: dict = {} # When WK is in FP8, we dequant to BF16 for fusion + indexer_fused_mapping = [ + ("wk_weights_proj", "wk", 0), + ("wk_weights_proj", "weights_proj", 1), + ] + stacked_params_mapping.extend(indexer_fused_mapping) if self.use_mha: stacked_params_mapping.extend(mha_params_mapping) @@ -1516,6 +1528,11 @@ class DeepseekV2ForCausalLM( rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name) ) + if _try_load_fp8_indexer_wk( + name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params + ): + continue + for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 08a4c4862ed..e7e71037cee 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -23,9 +23,8 @@ # limitations under the License. """Inference-only Ernie VL model compatible with HuggingFace weights.""" -import itertools import math -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from functools import partial from typing import Annotated, Any, Literal @@ -1401,131 +1400,62 @@ class Ernie4_5_VLMoeForConditionalGeneration( input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - kwargs = MultiModalFeatureSpec.gather_kwargs( - mm_features, - {"image_grid_thw", "video_grid_thw"}, - ) - image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])] - video_grid_thw = [item.tolist() for item in kwargs.get("video_grid_thw", [])] - - hf_config = self.config - image_token_id = hf_config.im_patch_id - video_start_token_id = hf_config.video_start_token_id - video_end_token_id = hf_config.video_end_token_id - spatial_conv_size = hf_config.spatial_conv_size - temporal_conv_size = hf_config.temporal_conv_size llm_pos_ids_list: list = [] + st = 0 - if image_grid_thw or video_grid_thw: - input_token_type: list[str] = [] - video_check_flg = False - for token in input_tokens: - if token == video_start_token_id: - video_check_flg = True - elif token == video_end_token_id: - video_check_flg = False + for ( + offset, + llm_grid_t, + llm_grid_h, + llm_grid_w, + ) in self.iter_mm_grid_thw(mm_features): + text_len = offset - st + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx + ) - if (token == image_token_id) and (video_check_flg is False): - input_token_type.append("image") - elif (token == image_token_id) and (video_check_flg is True): - input_token_type.append("video") - else: - input_token_type.append("text") + grid_indices = np.indices((llm_grid_t, llm_grid_h, llm_grid_w)).reshape( + 3, -1 + ) + llm_pos_ids_list.append(grid_indices + text_len + st_idx) + st = offset + llm_grid_t * llm_grid_h * llm_grid_w - input_type_group: list[tuple[str, int, int]] = [] - for key, group_iter in itertools.groupby( - enumerate(input_token_type), lambda x: x[1] - ): - group_list = list(group_iter) - start_index = group_list[0][0] - end_index = group_list[-1][0] + 1 - input_type_group.append((key, start_index, end_index)) + if st < len(input_tokens): + text_len = len(input_tokens) - st + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx + ) - video_frame_num = 1 - mm_data_idx = 0 - for modality_type, start_idx, end_idx in input_type_group: - st_idx = ( - llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - ) - if modality_type == "image": - t, h, w = image_grid_thw[mm_data_idx] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t, - h // spatial_conv_size, - w // spatial_conv_size, - ) - - t_index = ( - torch.arange(llm_grid_t) - .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - .flatten() - ) - h_index = ( - torch.arange(llm_grid_h) - .view(1, -1, 1) - .expand(llm_grid_t, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w) - .view(1, 1, -1) - .expand(llm_grid_t, llm_grid_h, -1) - .flatten() - ) - llm_pos_ids_list.append( - torch.stack([t_index, h_index, w_index]) + st_idx - ) - mm_data_idx += 1 - - elif modality_type == "video": - t, h, w = video_grid_thw[mm_data_idx] - llm_grid_t, llm_grid_h, llm_grid_w = ( - t // temporal_conv_size, - h // spatial_conv_size, - w // spatial_conv_size, - ) - - for t_idx in range(llm_grid_t): - t_index = ( - torch.tensor(t_idx) - .view(-1, 1) - .expand(-1, llm_grid_h * llm_grid_w) - .flatten() - ) - h_index = ( - torch.arange(llm_grid_h) - .view(1, -1, 1) - .expand(1, -1, llm_grid_w) - .flatten() - ) - w_index = ( - torch.arange(llm_grid_w) - .view(1, 1, -1) - .expand(1, llm_grid_h, -1) - .flatten() - ) - llm_pos_ids_list.append( - torch.stack([t_index, h_index, w_index]) + st_idx - ) - - mm_data_idx += 1 - video_frame_num += 1 - - else: - text_len = end_idx - start_idx - llm_pos_ids_list.append( - torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx - ) - video_frame_num = 1 - - else: - text_len = len(input_tokens) - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1)) - - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item() - return llm_positions, mrope_position_delta + return torch.from_numpy(llm_positions), mrope_position_delta + + def iter_mm_grid_thw( + self, mm_features: list[MultiModalFeatureSpec] + ) -> Iterator[tuple[int, int, int, int]]: + spatial_conv_size = self.config.spatial_conv_size + temporal_conv_size = self.config.temporal_conv_size + + for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset): + if mm_feature.data is None: + raise ValueError("M-RoPE calculation requires multimodal feature data") + + offset = mm_feature.mm_position.offset + if mm_feature.modality == "image": + t, h, w = mm_feature.data["image_grid_thw"].data.tolist() + yield offset, t, h // spatial_conv_size, w // spatial_conv_size + elif mm_feature.modality == "video": + t, h, w = mm_feature.data["video_grid_thw"].data.tolist() + yield ( + offset, + t // temporal_conv_size, + h // spatial_conv_size, + w // spatial_conv_size, + ) + else: + raise ValueError(f"Unsupported modality: {mm_feature.modality}") def _parse_and_validate_image_input( self, **kwargs: object diff --git a/vllm/model_executor/models/exaone4.py b/vllm/model_executor/models/exaone4.py index 485b145b9cd..04708de93d3 100644 --- a/vllm/model_executor/models/exaone4.py +++ b/vllm/model_executor/models/exaone4.py @@ -75,6 +75,7 @@ class Exaone4GatedMLP(nn.Module): reduce_results: bool = True, bias: bool = False, prefix: str = "", + use_data_parallel: bool = False, ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( @@ -83,6 +84,7 @@ class Exaone4GatedMLP(nn.Module): bias=bias, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", + disable_tp=use_data_parallel, ) self.down_proj = RowParallelLinear( input_size=intermediate_size, @@ -91,6 +93,7 @@ class Exaone4GatedMLP(nn.Module): quant_config=quant_config, reduce_results=reduce_results, prefix=f"{prefix}.down_proj", + disable_tp=use_data_parallel, ) if hidden_act != "silu": raise ValueError( diff --git a/vllm/model_executor/models/exaone4_5.py b/vllm/model_executor/models/exaone4_5.py new file mode 100644 index 00000000000..1eac43ccb0c --- /dev/null +++ b/vllm/model_executor/models/exaone4_5.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E501 + +# 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. +"""Inference-only EXAONE-4.5 model compatible with HuggingFace weights.""" + +from collections.abc import Callable, Iterable +from functools import partial + +import einops +import torch +import torch.nn as nn +from transformers.models.exaone4_5 import ( + Exaone4_5_Config, + Exaone4_5_ImageProcessor, + Exaone4_5_Processor, +) +from transformers.models.exaone4_5.configuration_exaone4_5 import Exaone4_5_VisionConfig + +from vllm.compilation.decorators import ( + should_torch_compile_mm_encoder, + support_torch_compile, +) +from vllm.config import VllmConfig +from vllm.distributed import parallel_state +from vllm.distributed import utils as dist_utils +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding.common import ( + ApplyRotaryEmb, +) +from vllm.model_executor.models.exaone4 import Exaone4GatedMLP as Exaone4_5_VisionMLP +from vllm.model_executor.models.qwen2_5_vl import ( + Qwen2_5_VisionTransformer, + Qwen2_5_VLForConditionalGeneration, + Qwen2VLProcessingInfo, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .qwen2_vl import Qwen2VLDummyInputsBuilder as Exaone4_5_DummyInputsBuilder +from .qwen2_vl import Qwen2VLMultiModalProcessor as Exaone4_5_MultiModalProcessor +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix + +logger = init_logger(__name__) + + +# === Vision Encoder === # + + +class EXAONE4_5_VisionAttention(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + num_kv_heads: int, + projection_size: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + super().__init__() + # Per attention head and per partition values. + self.tp_size = ( + 1 + if use_data_parallel + else parallel_state.get_tensor_model_parallel_world_size() + ) + self.tp_rank = parallel_state.get_tensor_model_parallel_rank() + self.hidden_size_per_attention_head = dist_utils.divide( + projection_size, num_heads + ) + self.num_attention_heads_per_partition = dist_utils.divide( + num_heads, self.tp_size + ) + + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + self.num_heads = num_heads // self.tp_size + self.num_kv_heads = max(1, num_kv_heads // self.tp_size) + + self.head_dim = embed_dim // num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + self.qkv = QKVParallelLinear( + hidden_size=embed_dim, + head_size=self.hidden_size_per_attention_head, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.qkv", + disable_tp=use_data_parallel, + ) + + self.proj = RowParallelLinear( + input_size=projection_size, + output_size=embed_dim, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.proj", + disable_tp=use_data_parallel, + ) + + self.attn = MMEncoderAttention( + num_heads=self.num_attention_heads_per_partition, + head_size=self.hidden_size_per_attention_head, + num_kv_heads=self.num_kv_heads, + scale=self.hidden_size_per_attention_head**-0.5, + prefix=f"{prefix}.attn", + ) + + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) + + def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]: + # qkv: [s, b, (h + 2*hk) * d] + s, b, _ = qkv.shape + h = self.num_heads + hk = self.num_kv_heads + d = self.head_dim + + qkv = qkv.view(s, b, h + 2 * hk, d) + + q = qkv[:, :, :h, :] + k = qkv[:, :, h : h + hk, :] + v = qkv[:, :, h + hk :, :] + + # [s, b, h, d] -> [b, s, h, d] + return ( + q.permute(1, 0, 2, 3).contiguous(), + k.permute(1, 0, 2, 3).contiguous(), + v.permute(1, 0, 2, 3).contiguous(), + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: int | None = None, + ) -> torch.Tensor: + # [s, b, c] --> [s, b, head * 3 * head_dim] + x, _ = self.qkv(x) + seq_len, batch_size, _ = x.shape + + q, k, v = self.split_qkv(x) + q = self.apply_rotary_emb( + q, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + ) + + k = self.apply_rotary_emb( + k, + rotary_pos_emb_cos, + rotary_pos_emb_sin, + ) + + context_layer = self.attn( + query=q, + key=k, + value=v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + + context_layer = einops.rearrange( + context_layer, "b s h d -> s b (h d)", b=batch_size + ).contiguous() + + output, _ = self.proj(context_layer) + return output + + +@support_torch_compile( + dynamic_arg_dims={ + "x": 0, + "cu_seqlens": 0, + "rotary_pos_emb_cos": 0, + "rotary_pos_emb_sin": 0, + }, + enable_if=should_torch_compile_mm_encoder, + is_encoder=True, +) +class Exaone4_5_VisionBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + num_kv_heads: int, + mlp_hidden_dim: int, + hidden_act: str = "silu", + norm_layer: Callable[[int], nn.Module] | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim) + self.attn = EXAONE4_5_VisionAttention( + embed_dim=dim, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + projection_size=dim, + quant_config=quant_config, + prefix=f"{prefix}.attn", + use_data_parallel=use_data_parallel, + ) + self.mlp = Exaone4_5_VisionMLP( + dim, + mlp_hidden_dim, + hidden_act=hidden_act, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + use_data_parallel=use_data_parallel, + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: int | None = None, # Only used for Flash Attention + seqlens: list[int] | None = None, # Only used for xFormers + ) -> torch.Tensor: + x_attn = self.attn( + self.norm1(x), + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + max_seqlen=max_seqlen, + ) + x_fused_norm, residual = self.norm2(x, residual=x_attn) + x = residual + self.mlp(x_fused_norm) + return x + + +class EXAONE4_5_VisionTransformer(Qwen2_5_VisionTransformer): + def __init__( + self, + vision_config: Exaone4_5_VisionConfig, + norm_eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + use_data_parallel: bool = False, + ) -> None: + super().__init__( + vision_config=vision_config, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=prefix, + ) + depth = vision_config.depth + self.num_kv_heads = vision_config.num_key_value_heads + + norm_layer = partial(RMSNorm, eps=norm_eps) + + self.blocks = nn.ModuleList( + [ + Exaone4_5_VisionBlock( + dim=self.hidden_size, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + mlp_hidden_dim=vision_config.intermediate_size, + hidden_act=vision_config.hidden_act, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{layer_idx}", + use_data_parallel=use_data_parallel, + ) + for layer_idx in range(depth) + ] + ) + + +class Exaone4_5_ProcessingInfo(Qwen2VLProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(Exaone4_5_Config) + + def get_hf_processor(self, **kwargs: object) -> Exaone4_5_Processor: + return self.ctx.get_hf_processor( + Exaone4_5_Processor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + + def get_image_processor(self, **kwargs: object) -> Exaone4_5_ImageProcessor: + return Exaone4_5_ImageProcessor(**kwargs) + + +@MULTIMODAL_REGISTRY.register_processor( + Exaone4_5_MultiModalProcessor, + info=Exaone4_5_ProcessingInfo, + dummy_inputs=Exaone4_5_DummyInputsBuilder, +) +class Exaone4_5_ForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + nn.Module.__init__(self) + + config: Exaone4_5_Config = vllm_config.model_config.hf_config + self.vllm_config = vllm_config + multimodal_config = vllm_config.model_config.multimodal_config + + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self.config = config + self.multimodal_config = multimodal_config + self.is_multimodal_pruning_enabled = ( + multimodal_config.is_multimodal_pruning_enabled() + ) + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = EXAONE4_5_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "visual"), + use_data_parallel=self.use_data_parallel, + ) + + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + hf_config=config.get_text_config(), + architectures=["Exaone4ForCausalLM"], + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["mtp."]), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "<|image_pad|>" + if modality.startswith("video"): + return "<|video_pad|>" + + raise ValueError("Only image or video modality is supported") diff --git a/vllm/model_executor/models/exaone4_5_mtp.py b/vllm/model_executor/models/exaone4_5_mtp.py new file mode 100644 index 00000000000..99bf724bdaa --- /dev/null +++ b/vllm/model_executor/models/exaone4_5_mtp.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only EXAONE-4_5 MTP model.""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.exaone4 import Exaone4DecoderLayer +from vllm.model_executor.models.exaone_moe_mtp import ( + ExaoneMoeMTP, + ExaoneMoeMultiTokenPredictor, +) + +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + _require_is_multimodal, +) +from .utils import ( + AutoWeightsLoader, + _merge_multimodal_embeddings, + maybe_prefix, +) + +logger = init_logger(__name__) + +KVCache = tuple[torch.Tensor, torch.Tensor] + + +@support_torch_compile +class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + nn.Module.__init__(self) + + model_config = vllm_config.model_config + quant_config = vllm_config.quant_config + lora_config = vllm_config.lora_config + config = model_config.hf_config + + self.config = config + lora_vocab = ( + (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) + if lora_config + else 0 + ) + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1) + + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + ) + + self.fc = ColumnParallelLinear( + self.config.hidden_size * 2, + self.config.hidden_size, + gather_output=True, + bias=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fc", + ) + self.layers = nn.ModuleList( + Exaone4DecoderLayer( + vllm_config.model_config.hf_config, + quant_config=quant_config, + prefix=f"{prefix}.layers.{idx}", + ) + for idx in range(self.num_mtp_layers) + ) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.pre_fc_norm_hidden = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.pre_fc_norm_embedding = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + +@support_torch_compile +class Exaone4_5_MTP(ExaoneMoeMTP, SupportsMultiModal): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + config = vllm_config.model_config.hf_config + self.vllm_config = vllm_config + self.quant_config = vllm_config.quant_config + + nn.Module.__init__(self) + self.config = config + self.model = Exaone4_5MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp") + ) + self.unpadded_vocab_size = config.vocab_size + self.lm_head = ParallelLMHead( + self.unpadded_vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_processor = LogitsProcessor( + self.unpadded_vocab_size, config.vocab_size + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.model.embed_input_ids, + is_multimodal=is_multimodal, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + is_multimodal = _require_is_multimodal(is_multimodal) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + shared_weight_names = ["embed_tokens", "lm_head"] + + def remap_weight_names(weights): + for name, weight in weights: + if name.startswith("mtp."): + name = name.replace("mtp.", "model.") + elif any(key in name for key in shared_weight_names): + if "embed_tokens" in name: + name = name.replace("language_model.", "") + else: + continue + yield name, weight + + loader = AutoWeightsLoader(self) + return loader.load_weights(remap_weight_names(weights)) diff --git a/vllm/model_executor/models/exaone_moe_mtp.py b/vllm/model_executor/models/exaone_moe_mtp.py index b3c71e6aef6..b3f8552aac5 100644 --- a/vllm/model_executor/models/exaone_moe_mtp.py +++ b/vllm/model_executor/models/exaone_moe_mtp.py @@ -184,11 +184,6 @@ class ExaoneMoeMTP(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config self.vllm_config = vllm_config - cache_config = vllm_config.cache_config - assert not cache_config.enable_prefix_caching, ( - "ExaoneMoeMTP currently does not support prefix caching" - ) - self.quant_config = vllm_config.quant_config super().__init__() diff --git a/vllm/model_executor/models/fireredasr2.py b/vllm/model_executor/models/fireredasr2.py index 217bb5b2d13..41b4318504f 100644 --- a/vllm/model_executor/models/fireredasr2.py +++ b/vllm/model_executor/models/fireredasr2.py @@ -6,7 +6,6 @@ from typing import Annotated, Literal, cast import numpy as np import torch -import torch.nn.functional as F from torch import nn from transformers import ( BatchFeature, @@ -45,6 +44,7 @@ from vllm.transformers_utils.processors.fireredasr2 import ( ) from vllm.utils.tensor_schema import TensorSchema, TensorShape +from .conformer_encoder import ConformerEncoder from .interfaces import ( MultiModalEmbeddings, SupportsMultiModal, @@ -84,352 +84,6 @@ class FireRedASR2AudioInputs(TensorSchema): ] -class Swish(nn.Module): - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x * torch.sigmoid(x) - - -class Conv2dSubsampling(nn.Module): - def __init__(self, idim: int, d_model: int, out_channels: int = 32): - super().__init__() - self.conv = nn.Sequential( - nn.Conv2d(1, out_channels, 3, 2), - nn.ReLU(), - nn.Conv2d(out_channels, out_channels, 3, 2), - nn.ReLU(), - ) - subsample_idim = ((idim - 1) // 2 - 1) // 2 - self.out = ReplicatedLinear( - input_size=out_channels * subsample_idim, - output_size=d_model, - bias=True, - ) - - self.subsampling = 4 - left_context = right_context = 3 # both exclude current frame - self.context = left_context + 1 + right_context # 7 - - def forward( - self, x: torch.Tensor, x_mask: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - x = x.unsqueeze(1) - x = self.conv(x) - N, C, T, D = x.size() - x, _ = self.out(x.transpose(1, 2).contiguous().view(N, T, C * D)) - mask = x_mask[:, :, :-2:2][:, :, :-2:2] - input_lengths = mask[:, -1, :].sum(dim=-1) - return x, input_lengths, mask - - -class RelPositionalEncoding(nn.Module): - def __init__(self, d_model: int, max_len: int = 5000): - super().__init__() - pe_positive = torch.zeros(max_len, d_model, requires_grad=False) - pe_negative = torch.zeros(max_len, d_model, requires_grad=False) - position = torch.arange(0, max_len).unsqueeze(1).float() - div_term = torch.exp( - torch.arange(0, d_model, 2).float() - * -(torch.log(torch.tensor(10000.0)).item() / d_model) - ) - pe_positive[:, 0::2] = torch.sin(position * div_term) - pe_positive[:, 1::2] = torch.cos(position * div_term) - pe_negative[:, 0::2] = torch.sin(-1 * position * div_term) - pe_negative[:, 1::2] = torch.cos(-1 * position * div_term) - - pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0) - pe_negative = pe_negative[1:].unsqueeze(0) - self.pe = torch.cat([pe_positive, pe_negative], dim=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # Tmax = 2 * max_len - 1 - Tmax, T = self.pe.size(1), x.size(1) - pos_emb = self.pe[:, Tmax // 2 - T + 1 : Tmax // 2 + T].clone().detach() - return pos_emb - - -class ConformerFeedForward(nn.Module): - def __init__(self, d_model: int): - super().__init__() - self.pre_layer_norm = nn.LayerNorm(d_model) - self.linear_expand = ReplicatedLinear( - input_size=d_model, - output_size=d_model * 4, - bias=True, - ) - self.nonlinear = Swish() - self.linear_project = ReplicatedLinear( - input_size=d_model * 4, - output_size=d_model, - bias=True, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = self.pre_layer_norm(x) - x, _ = self.linear_expand(x) - x = self.nonlinear(x) - x, _ = self.linear_project(x) - output = x + residual - return output - - -class EncoderMultiHeadAttention(nn.Module): - def __init__(self, n_head: int, d_model: int): - super().__init__() - assert d_model % n_head == 0 - self.n_head = n_head - self.d_k = d_model // n_head - self.d_v = self.d_k - - self.w_qs = ReplicatedLinear( - input_size=d_model, output_size=n_head * self.d_k, bias=False - ) - self.w_ks = ReplicatedLinear( - input_size=d_model, output_size=n_head * self.d_k, bias=False - ) - self.w_vs = ReplicatedLinear( - input_size=d_model, output_size=n_head * self.d_v, bias=False - ) - - self.layer_norm_q = nn.LayerNorm(d_model) - self.layer_norm_k = nn.LayerNorm(d_model) - self.layer_norm_v = nn.LayerNorm(d_model) - - self.fc = ReplicatedLinear( - input_size=n_head * self.d_v, output_size=d_model, bias=False - ) - - def forward_qkv( - self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - d_k, d_v, n_head = self.d_k, self.d_v, self.n_head - sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) - - q = self.layer_norm_q(q) - k = self.layer_norm_k(k) - v = self.layer_norm_v(v) - - q = self.w_qs(q)[0].view(sz_b, len_q, n_head, d_k) - k = self.w_ks(k)[0].view(sz_b, len_k, n_head, d_k) - v = self.w_vs(v)[0].view(sz_b, len_v, n_head, d_v) - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - return q, k, v - - def forward_output( - self, output: torch.Tensor, residual: torch.Tensor, sz_b: int, len_q: int - ) -> torch.Tensor: - output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1) - fc_out, _ = self.fc(output) - output = fc_out - output = output + residual - return output - - def forward_attention( - self, attn: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None - ) -> tuple[torch.Tensor, torch.Tensor]: - if mask is not None: - mask = mask.unsqueeze(1) - mask = mask.eq(0) - attn = attn.masked_fill(mask, -float("inf")) - attn = torch.softmax(attn, dim=-1).masked_fill(mask, 0.0) - else: - attn = torch.softmax(attn, dim=-1) - - d_attn = attn - output = torch.matmul(d_attn, v) - - return output, attn - - -class RelPosMultiHeadAttention(EncoderMultiHeadAttention): - def __init__(self, n_head: int, d_model: int): - super().__init__(n_head, d_model) - d_k = d_model // n_head - self.scale = 1.0 / (d_k**0.5) - self.linear_pos = ReplicatedLinear( - input_size=d_model, output_size=n_head * d_k, bias=False - ) - self.pos_bias_u = nn.Parameter(torch.empty([n_head, d_k])) - self.pos_bias_v = nn.Parameter(torch.empty([n_head, d_k])) - - def _rel_shift(self, x): - N, H, T1, T2 = x.size() - zero_pad = torch.zeros((N, H, T1, 1), device=x.device, dtype=x.dtype) - x_padded = torch.cat([zero_pad, x], dim=-1) - - x_padded = x_padded.view(N, H, T2 + 1, T1) - x = x_padded[:, :, 1:].view_as(x) - x = x[:, :, :, : x.size(-1) // 2 + 1] - return x - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - pos_emb: torch.Tensor, - mask: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - sz_b, len_q = q.size(0), q.size(1) - - residual = q - q, k, v = self.forward_qkv(q, k, v) - - q = q.transpose(1, 2) - n_batch_pos = pos_emb.size(0) - p = self.linear_pos(pos_emb)[0].view(n_batch_pos, -1, self.n_head, self.d_k) - p = p.transpose(1, 2) - - q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) - q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) - - matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) - - matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) - matrix_bd = self._rel_shift(matrix_bd) - - attn_scores = matrix_ac + matrix_bd - attn_scores.mul_(self.scale) - - output, attn = self.forward_attention(attn_scores, v, mask=mask) - - output = self.forward_output(output, residual, sz_b, len_q) - return output, attn - - -class ConformerConvolution(nn.Module): - def __init__(self, d_model: int, kernel_size: int = 33): - super().__init__() - assert kernel_size % 2 == 1 - self.pre_layer_norm = nn.LayerNorm(d_model) - self.pointwise_conv1 = nn.Conv1d( - d_model, d_model * 4, kernel_size=1, bias=False - ) - self.padding = (kernel_size - 1) // 2 - self.depthwise_conv = nn.Conv1d( - d_model * 2, - d_model * 2, - kernel_size, - stride=1, - padding=self.padding, - groups=d_model * 2, - bias=False, - ) - self.batch_norm = nn.LayerNorm(d_model * 2) - self.swish = Swish() - self.pointwise_conv2 = nn.Conv1d( - d_model * 2, d_model, kernel_size=1, bias=False - ) - - def forward( - self, x: torch.Tensor, mask: torch.Tensor | None = None - ) -> torch.Tensor: - residual = x - out = self.pre_layer_norm(x) - out = out.transpose(1, 2) - if mask is not None: - out.masked_fill_(mask.ne(1), 0.0) - out = self.pointwise_conv1(out) - out = F.glu(out, dim=1) - out = self.depthwise_conv(out) - - out = out.transpose(1, 2) - out = self.swish(self.batch_norm(out)) - out = out.transpose(1, 2) - - out = self.pointwise_conv2(out) - if mask is not None: - out.masked_fill_(mask.ne(1), 0.0) - out = out.transpose(1, 2) - return out + residual - - -class RelPosEmbConformerBlock(nn.Module): - def __init__(self, d_model, n_head, kernel_size=33): - super().__init__() - self.ffn1 = ConformerFeedForward(d_model) - self.mhsa = RelPosMultiHeadAttention(n_head, d_model) - self.conv = ConformerConvolution(d_model, kernel_size) - self.ffn2 = ConformerFeedForward(d_model) - self.layer_norm = nn.LayerNorm(d_model) - - def forward( - self, - x: torch.Tensor, - pos_emb: torch.Tensor, - slf_attn_mask: torch.Tensor | None = None, - pad_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - out = 0.5 * x + 0.5 * self.ffn1(x) - out = self.mhsa(out, out, out, pos_emb, mask=slf_attn_mask)[0] - out = self.conv(out, pad_mask) - out = 0.5 * out + 0.5 * self.ffn2(out) - out = self.layer_norm(out) - return out - - -class ConformerEncoder(nn.Module): - def __init__( - self, - idim: int, - n_layers_enc: int, - n_head: int, - d_model: int, - kernel_size: int = 33, - pe_maxlen: int = 5000, - ): - super().__init__() - self.odim = d_model - - self.input_preprocessor = Conv2dSubsampling(idim, d_model) - self.positional_encoding = RelPositionalEncoding(d_model) - - self.layer_stack = nn.ModuleList() - for _ in range(n_layers_enc): - block = RelPosEmbConformerBlock(d_model, n_head, kernel_size) - self.layer_stack.append(block) - - def forward( - self, padded_input: torch.Tensor, input_lengths: torch.Tensor, pad: bool = True - ): - if pad: - padded_input = F.pad( - padded_input, - (0, 0, 0, self.input_preprocessor.context - 1), - "constant", - 0.0, - ) - src_mask = self.padding_position_is_0(padded_input, input_lengths) - - embed_output, input_lengths, src_mask = self.input_preprocessor( - padded_input, src_mask - ) - enc_output = embed_output - - pos_emb = self.positional_encoding(embed_output) - - enc_outputs = [] - for enc_layer in self.layer_stack: - enc_output = enc_layer( - enc_output, pos_emb, slf_attn_mask=src_mask, pad_mask=src_mask - ) - enc_outputs.append(enc_output) - - return enc_output, input_lengths, src_mask - - def padding_position_is_0( - self, padded_input: torch.Tensor, input_lengths: torch.Tensor - ) -> torch.Tensor: - N, T = padded_input.size()[:2] - mask = torch.ones((N, T)).to(padded_input.device) - for i in range(N): - mask[i, input_lengths[i] :] = 0 - mask = mask.unsqueeze(dim=1) - return mask.to(torch.uint8) - - class FireRedASR2Adapter(nn.Module): def __init__(self, encoder_dim: int, llm_dim: int, downsample_rate: int = 2): super().__init__() diff --git a/vllm/model_executor/models/fireredlid.py b/vllm/model_executor/models/fireredlid.py new file mode 100644 index 00000000000..804ed2bc9fd --- /dev/null +++ b/vllm/model_executor/models/fireredlid.py @@ -0,0 +1,792 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +FireRedLID – Language Identification model adapted for vLLM. + +Architecture: ConformerEncoder + TransformerDecoder (6-layer cross-attn) +Vocabulary: 120 LID tokens (dict.txt) +Output: Up to 2 tokens (e.g. "en", "zh mandarin") + +This implementation follows the Whisper-style encoder-decoder pattern: + • Encoder processes audio features (Fbank + CMVN via FeatureExtractor) + • Decoder performs single-step autoregressive forward + • vLLM's generation loop handles beam search / sampling +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from typing import Annotated, Literal + +import numpy as np +import torch +from torch import nn +from transformers import BatchFeature + +from vllm.config import ModelConfig, VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.config.speech_to_text import SpeechToTextConfig +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.inputs import MultiModalDataDict, PromptType +from vllm.logger import init_logger +from vllm.model_executor.layers.attention import Attention, CrossAttention +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseProcessingInfo, + EncDecMultiModalProcessor, + PromptReplacement, + PromptUpdate, +) +from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from .conformer_encoder import ConformerEncoder +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsTranscription, +) +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, +) +from .whisper_utils import ISO639_1_SUPPORTED_LANGS + +logger = init_logger(__name__) + + +class FireRedLIDAudioInputs(TensorSchema): + """ + Dimensions: + - b: Batch size + - t: Time frames (variable across utterances) + - nmb: Number of mel bins (80) + """ + + input_features: Annotated[ + list[torch.Tensor] | None, + TensorShape("b", "t", "nmb", dynamic_dims={"t"}), + ] + speech_lengths: Annotated[ + list[torch.Tensor] | None, + TensorShape("b"), + ] + fake_token_lengths: Annotated[ + list[torch.Tensor] | None, + TensorShape("b"), + ] + + +FireRedLIDEncoder = ConformerEncoder + + +class FireRedLIDPositionalEmbedding(nn.Module): + """Absolute sinusoidal positional embedding indexed by `positions`.""" + + def __init__(self, d_model: int, max_len: int = 5000): + super().__init__() + assert d_model % 2 == 0 + pe = torch.zeros(max_len, d_model, requires_grad=False) + position = torch.arange(0, max_len).unsqueeze(1).float() + div_term = torch.exp( + torch.arange(0, d_model, 2).float() + * -(torch.log(torch.tensor(10000.0)).item() / d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + self.register_buffer("pe", pe, persistent=False) + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: + return self.pe[position_ids] + + +class FireRedLIDAttention(nn.Module): + """Base attention with shared QKV/FC projections for the LID decoder.""" + + def __init__( + self, + d_model: int, + n_head: int, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + assert n_head % tp_size == 0 + self.total_num_heads = n_head + self.num_heads = n_head // tp_size + self.num_kv_heads = max(1, n_head // tp_size) + self.head_dim = d_model // n_head + self.scaling = self.head_dim**-0.5 + + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.w_qs = ColumnParallelLinear( + d_model, + d_model, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.w_qs", + ) + self.w_ks = ColumnParallelLinear( + d_model, + d_model, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.w_ks", + ) + self.w_vs = ColumnParallelLinear( + d_model, + d_model, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.w_vs", + ) + self.fc = RowParallelLinear( + d_model, + d_model, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.fc", + ) + self._init_attn(cache_config, quant_config, prefix) + + def _init_attn(self, cache_config, quant_config, prefix: str) -> None: + raise NotImplementedError + + +class FireRedLIDSelfAttention(FireRedLIDAttention): + def _init_attn(self, cache_config, quant_config, prefix: str) -> None: + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + q, _ = self.w_qs(hidden_states) + k, _ = self.w_ks(hidden_states) + v, _ = self.w_vs(hidden_states) + attn_output = self.attn(q, k, v) + output, _ = self.fc(attn_output) + return output + + +class FireRedLIDCrossAttention(FireRedLIDAttention): + def _init_attn(self, cache_config, quant_config, prefix: str) -> None: + self.attn = CrossAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + q, _ = self.w_qs(hidden_states) + if encoder_hidden_states is not None: + k, _ = self.w_ks(encoder_hidden_states) + v, _ = self.w_vs(encoder_hidden_states) + else: + k = v = None + + attn_output = self.attn(q, k, v) + output, _ = self.fc(attn_output) + return output + + +class FireRedLIDFFN(nn.Module): + def __init__(self, d_model: int, d_ff: int): + super().__init__() + self.w_1 = ReplicatedLinear(d_model, d_ff, bias=True) + self.act = nn.GELU() + self.w_2 = ReplicatedLinear(d_ff, d_model, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.w_1(x) + x = self.act(x) + x, _ = self.w_2(x) + return x + + +class FireRedLIDDecoderLayer(nn.Module): + """vLLM-native decoder layer while preserving FireRedLID parameter names.""" + + def __init__( + self, + d_model: int, + n_head: int, + *, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + self.self_attn_norm = nn.LayerNorm(d_model) + self.self_attn = FireRedLIDSelfAttention( + d_model, + n_head, + vllm_config=vllm_config, + prefix=f"{prefix}.self_attn", + ) + + self.cross_attn_norm = nn.LayerNorm(d_model) + self.cross_attn = FireRedLIDCrossAttention( + d_model, + n_head, + vllm_config=vllm_config, + prefix=f"{prefix}.cross_attn", + ) + + self.mlp_norm = nn.LayerNorm(d_model) + self.mlp = FireRedLIDFFN(d_model, d_model * 4) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.self_attn_norm(hidden_states) + hidden_states = self.self_attn(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.cross_attn_norm(hidden_states) + hidden_states = self.cross_attn(hidden_states, encoder_hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = residual + self.mlp(hidden_states) + + return hidden_states + + +class FireRedLIDDecoder(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.pad_id = getattr(config, "pad_token_id", 2) + self.n_layers = getattr(config, "n_layers_lid_dec", 6) + self.d_model = getattr(config, "d_model", 1280) + self.scale = self.d_model**0.5 + + self.tgt_word_emb = nn.Embedding( + getattr(config, "vocab_size", 120), + self.d_model, + padding_idx=self.pad_id, + ) + self.positional_encoding = FireRedLIDPositionalEmbedding( + self.d_model, + max_len=getattr(config, "pe_maxlen", 5000), + ) + + self.layer_stack = nn.ModuleList( + [ + FireRedLIDDecoderLayer( + self.d_model, + getattr(config, "n_head", 20), + vllm_config=vllm_config, + prefix=f"{prefix}.layer_stack.{idx}", + ) + for idx in range(self.n_layers) + ] + ) + self.layer_norm_out = nn.LayerNorm(self.d_model) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + ) -> torch.Tensor: + hidden_states = self.tgt_word_emb(input_ids) * self.scale + hidden_states = hidden_states + self.positional_encoding(positions) + + for layer in self.layer_stack: + hidden_states = layer(hidden_states, encoder_hidden_states) + + hidden_states = self.layer_norm_out(hidden_states) + return hidden_states + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.tgt_word_emb(input_ids) + + +class FireRedLIDModel(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + + self.encoder = FireRedLIDEncoder( + idim=getattr(config, "idim", 80), + n_layers_enc=getattr(config, "n_layers_enc", 16), + n_head=getattr(config, "n_head", 20), + d_model=getattr(config, "d_model", 1280), + kernel_size=getattr(config, "kernel_size", 33), + pe_maxlen=getattr(config, "pe_maxlen", 5000), + ) + + self.decoder = FireRedLIDDecoder( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "decoder"), + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + encoder_outputs: list[torch.Tensor] | None = None, + ) -> torch.Tensor: + enc_states = ( + torch.cat(encoder_outputs, dim=0) + if encoder_outputs and len(encoder_outputs) > 0 + else None + ) + decoder_outputs = self.decoder( + input_ids=input_ids, + positions=positions, + encoder_hidden_states=enc_states, + ) + return decoder_outputs + + def get_encoder_outputs( + self, + speech: torch.Tensor | list[torch.Tensor], + speech_lengths: torch.Tensor | list[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the encoder and return padded outputs plus true sequence lengths.""" + enc_output, enc_lengths, _ = self.encoder(speech, speech_lengths) + return enc_output, enc_lengths + + +class FireRedLIDProcessingInfo(BaseProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config() + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + def get_feature_extractor(self, **kwargs): + hf_processor = self.get_hf_processor(**kwargs) + feature_extractor = hf_processor.feature_extractor + return feature_extractor + + def get_data_parser(self) -> MultiModalDataParser: + feature_extractor = self.get_feature_extractor() + return MultiModalDataParser( + target_sr=feature_extractor.sampling_rate, + target_channels=1, + ) + + @property + def skip_prompt_length_check(self) -> bool: + return True + + def get_num_audio_tokens(self) -> int: + # For encoder profiling – return a reasonable dummy length. + # This doesn't affect actual inference since encoder processes + # variable-length features. + return 1 + + +class FireRedLIDDummyInputsBuilder(BaseDummyInputsBuilder[FireRedLIDProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return "" + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + feature_extractor = self.info.get_feature_extractor() + sampling_rate = feature_extractor.sampling_rate + audio_len = feature_extractor.chunk_length * sampling_rate + num_audios = mm_counts.get("audio", 0) + audio_overrides = mm_options.get("audio") + return { + "audio": self._get_dummy_audios( + length=audio_len, + num_audios=num_audios, + overrides=audio_overrides, + ) + } + + +class FireRedLIDMultiModalProcessor( + EncDecMultiModalProcessor[FireRedLIDProcessingInfo] +): + def create_encoder_prompt( + self, + prompt: str | list[int], + mm_items: MultiModalDataItems, + ) -> str | list[int]: + # Dummy encoder prompt for profiling (encoder only processes audio). + return [0] + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + if mm_data: + feature_extractor = self.info.get_feature_extractor(**mm_kwargs) + mm_data = dict(audio=mm_data.pop("audios")) + mm_kwargs = dict( + **mm_kwargs, + sampling_rate=feature_extractor.sampling_rate, + ) + processed_outputs = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + if "labels" in processed_outputs: + processed_outputs["input_ids"] = processed_outputs.pop("labels") + return processed_outputs + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return dict( + input_features=MultiModalFieldConfig.batched("audio"), + speech_lengths=MultiModalFieldConfig.batched("audio"), + fake_token_lengths=MultiModalFieldConfig.batched("audio"), + ) + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + out_mm_data = out_mm_kwargs.get_data() + fake_token_lengths = out_mm_data.get("fake_token_lengths") + + if fake_token_lengths is None: + # Fallback to max encoder output length if not available + audio_output_lengths = [] + else: + assert isinstance(fake_token_lengths, torch.Tensor) + audio_output_lengths = fake_token_lengths.tolist() + + def get_replacement(item_idx: int): + if audio_output_lengths: + num_tokens = int(audio_output_lengths[item_idx]) + else: + num_tokens = self.info.get_num_audio_tokens() + return [0] * num_tokens + + return [ + PromptReplacement( + modality="audio", + target=[0], + replacement=get_replacement, + ) + ] + + +# FireRedLID supports a wider set of languages than Whisper's shared list. +# Only ISO 639-1 codes are listed; FireRedLID's dialect tokens (mandarin, +# xinan, wu, …) are output tokens but not valid language *request* codes. +_FIREREDLID_SUPPORTED_LANGUAGES: Mapping[str, str] = { + **ISO639_1_SUPPORTED_LANGS, + "am": "Amharic", + "as": "Assamese", + "ba": "Bashkir", + "bn": "Bengali", + "bo": "Tibetan", + "br": "Breton", + "eu": "Basque", + "fo": "Faroese", + "gu": "Gujarati", + "ha": "Hausa", + "haw": "Hawaiian", + "ht": "Haitian Creole", + "jw": "Javanese", + "ka": "Georgian", + "km": "Khmer", + "la": "Latin", + "lb": "Luxembourgish", + "ln": "Lingala", + "lo": "Lao", + "mg": "Malagasy", + "ml": "Malayalam", + "mn": "Mongolian", + "mt": "Maltese", + "my": "Myanmar", + "nn": "Nynorsk", + "oc": "Occitan", + "pa": "Panjabi", + "ps": "Pashto", + "sa": "Sanskrit", + "sd": "Sindhi", + "si": "Sinhala", + "sn": "Shona", + "so": "Somali", + "sq": "Albanian", + "su": "Sundanese", + "te": "Telugu", + "tg": "Tajik", + "tk": "Turkmen", + "tt": "Tatar", + "uz": "Uzbek", + "yi": "Yiddish", + "yo": "Yoruba", + "yue": "Cantonese", +} + + +@MULTIMODAL_REGISTRY.register_processor( + FireRedLIDMultiModalProcessor, + info=FireRedLIDProcessingInfo, + dummy_inputs=FireRedLIDDummyInputsBuilder, +) +class FireRedLIDForConditionalGeneration( + nn.Module, SupportsTranscription, SupportsMultiModal +): + # -- SupportsTranscription protocol attributes -- + supports_transcription_only = True + supported_languages = _FIREREDLID_SUPPORTED_LANGUAGES + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + "encoder.": "model.encoder.", + "lid_decoder.": "model.decoder.", + # Encoder FFN: nn.Sequential indices → named children + "net.0": "pre_layer_norm", + "net.1": "linear_expand", + "net.4": "linear_project", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + self.config = config + self.dtype = vllm_config.model_config.dtype + + with self._mark_composite_model( + vllm_config, + language_targets=FireRedLIDDecoder, + tower_targets={"audio": FireRedLIDEncoder}, + ): + self.model = FireRedLIDModel( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + self.proj_out = ParallelLMHead( + getattr(config, "vocab_size", 120), + getattr(config, "d_model", 1280), + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "proj_out"), + ) + self.proj_out = self.proj_out.tie_weights(self.model.decoder.tgt_word_emb) + + logit_scale = getattr(config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + getattr(config, "vocab_size", 120), + scale=logit_scale, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + encoder_outputs: list[torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + if encoder_outputs is None: + encoder_outputs = [] + decoder_outputs = self.model( + input_ids=input_ids, + positions=positions, + encoder_outputs=encoder_outputs, + ) + return decoder_outputs + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + """Run encoder on audio features and return per-item embeddings.""" + audio_input = self._parse_and_validate_audio_input(**kwargs) + + speech = audio_input["input_features"] + speech_lengths = audio_input["speech_lengths"] + if speech is None or speech_lengths is None: + return [] + + # When audio items have different time lengths, vLLM's + # MultiModalBatchedField._reduce_data returns a plain + # list[Tensor] instead of a stacked Tensor. The encoder + # expects a padded [B, Tmax, feat_dim] Tensor, so we + # normalise both speech and speech_lengths here. + if isinstance(speech, (list, tuple)): + # Each element: [Ti, feat_dim] (or [1, Ti, feat_dim]) + tensors = [ + s.squeeze(0) if s.dim() == 3 and s.size(0) == 1 else s for s in speech + ] + device = tensors[0].device + dtype = tensors[0].dtype + feat_dim = tensors[0].shape[-1] + lengths = torch.tensor( + [t.size(0) for t in tensors], + device=device, + dtype=torch.int32, + ) + t_max = int(lengths.max().item()) + # Pre-allocate zero-padded batch tensor + speech = torch.zeros( + (len(tensors), t_max, feat_dim), + device=device, + dtype=dtype, + ) + for i, t in enumerate(tensors): + speech[i, : t.size(0)] = t + speech_lengths = lengths + else: + # Already a batched Tensor [B, T, feat_dim] + if speech.dim() == 2: + speech = speech.unsqueeze(0) + + speech_lengths = torch.as_tensor( + speech_lengths, dtype=torch.int32, device=speech.device + ) + + enc_output, enc_lengths = self.model.get_encoder_outputs( + speech=speech, + speech_lengths=speech_lengths, + ) + + # vLLM expects one 2D tensor per multimodal item. Slice each batch entry + # by the true encoder length so cross-attention never sees padded frames. + return tuple( + enc_output[i, : max(0, int(enc_lengths[i].item()))] + for i in range(enc_output.size(0)) + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model.decoder.embed_input_ids(input_ids) + + def _parse_and_validate_audio_input( + self, **kwargs: object + ) -> FireRedLIDAudioInputs: + input_features = kwargs.pop("input_features", None) + speech_lengths = kwargs.pop("speech_lengths", None) + fake_token_lengths = kwargs.pop("fake_token_lengths", None) + return FireRedLIDAudioInputs( + input_features=input_features, + speech_lengths=speech_lengths, + fake_token_lengths=fake_token_lengths, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = self.logits_processor(self.proj_out, hidden_states) + return logits + + @classmethod + def validate_language(cls, language: str | None) -> str | None: + # FireRedLID is a language *identification* model – the caller does + # not need to specify a language up-front. Accept None silently. + if language is None: + return None + return super().validate_language(language) + + @classmethod + def get_generation_prompt( + cls, + audio: np.ndarray, + stt_config: SpeechToTextConfig, + model_config: ModelConfig, + language: str | None, + task_type: Literal["transcribe", "translate"], + request_prompt: str, + to_language: str | None, + ) -> PromptType: + """Build the prompt for the FireRedLID encoder-decoder model. + + The decoder receives a single token; the encoder processes + the raw audio waveform via the multimodal pipeline. + """ + prompt: PromptType = { + "encoder_prompt": { + "prompt": "", + "multi_modal_data": { + "audio": (audio, int(stt_config.sample_rate)), + }, + }, + "decoder_prompt": { + "prompt": "", + }, + } + return prompt + + @classmethod + def get_speech_to_text_config( + cls, + model_config: ModelConfig, + task_type: Literal["transcribe", "translate"], + ) -> SpeechToTextConfig: + processor = cached_processor_from_config(model_config) + return SpeechToTextConfig( + max_audio_clip_s=processor.feature_extractor.chunk_length, + sample_rate=processor.feature_extractor.sampling_rate, + # LID output is at most 2 tokens – no chunking needed. + min_energy_split_window_size=None, + ) + + @classmethod + def post_process_output(cls, text: str) -> str: + # Strip any leading/trailing whitespace from the raw LID output. + return text.strip() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=[ + # Position encoding buffers are rebuilt at init + "model.encoder.positional_encoding.pe", + "model.decoder.positional_encoding.pe", + # Tied output projection (shared with embedding) + "model.decoder.tgt_word_prj.weight", + "proj_out.", + ], + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py index 2e9fc681903..06189540090 100644 --- a/vllm/model_executor/models/gemma4.py +++ b/vllm/model_executor/models/gemma4.py @@ -60,9 +60,16 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.sequence import IntermediateTensors from vllm.v1.attention.backends.utils import KVSharingFastPrefillMetadata -from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP +from .interfaces import ( + EagleModelMixin, + MixtureOfExperts, + SupportsEagle3, + SupportsLoRA, + SupportsPP, +) from .utils import ( AutoWeightsLoader, + WeightsMapper, extract_layer_index, is_pp_missing_parameter, make_layers, @@ -838,7 +845,7 @@ class Gemma4CrossDecoderLayers(nn.Module): @support_torch_compile( enable_if=lambda vllm_config: not vllm_config.cache_config.kv_sharing_fast_prefill ) -class Gemma4Model(nn.Module): +class Gemma4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = _get_text_config(vllm_config.model_config.hf_config) @@ -1168,7 +1175,7 @@ class Gemma4Model(nn.Module): inputs_embeds: torch.Tensor | None = None, per_layer_inputs: torch.Tensor | None = None, **kwargs, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: if self.fast_prefill_enabled: hidden_states = self.fast_prefill_forward( input_ids, @@ -1204,6 +1211,7 @@ class Gemma4Model(nn.Module): residual = intermediate_tensors["residual"] per_layer_inputs = intermediate_tensors.get("per_layer_inputs") + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) for layer_idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer) ): @@ -1222,6 +1230,9 @@ class Gemma4Model(nn.Module): per_layer_input=layer_per_input, **kwargs, ) + self._maybe_add_hidden_state( + aux_hidden_states, layer_idx + 1, hidden_states, residual + ) if not get_pp_group().is_last_rank: return IntermediateTensors( { @@ -1236,6 +1247,9 @@ class Gemma4Model(nn.Module): hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1248,21 +1262,27 @@ class Gemma4Model(nn.Module): ("gate_up_proj", "up_proj", 1), ] - # MoE expert weight mapping: checkpoint 3D packed tensors are - # exploded in _weight_iterator to per-expert 2D weights like: + # MoE expert weight mapping: checkpoint can have either: + # 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D) + # 2. Already per-expert 2D weights (if quantized) + # Map to FusedMoE parameters: # moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13) # moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13) # moe.experts.{id}.down_proj → FusedMoE w2 - # We build the mapping directly since Gemma4 uses bare param - # names (no .weight suffix) unlike standard MoE checkpoints. + # + # Use prefix matching to handle both weights and + # quantization scale parameters. The param_name is a prefix ending + # in underscore, and weight_name ends with a dot, so that: + # "experts.0.gate_proj.weight_scale" -> "experts.w13_weight_scale" + # "experts.0.gate_proj.weight" -> "experts.w13_weight" num_experts = getattr(self.config, "num_experts", None) or 0 expert_params_mapping = [ # (param_name, weight_name, expert_id, shard_id) ( - "experts.w13_weight" + "experts.w13_" if proj_name in ["gate_proj", "up_proj"] - else "experts.w2_weight", - f"experts.{expert_id}.{proj_name}", + else "experts.w2_", + f"experts.{expert_id}.{proj_name}.", expert_id, shard_id, ) @@ -1322,9 +1342,21 @@ class Gemma4Model(nn.Module): expert_id, shard_id, ) in expert_params_mapping: - if weight_name not in name: + # Match both: + # - Bare weights: "experts.0.down_proj" (from 3D explosion) + # - With suffix: "experts.0.down_proj.weight_scale" (2D quantized) + # weight_name has trailing dot, so check with and without it + weight_name_base = weight_name.rstrip(".") + if weight_name in name: + # Has suffix (e.g., .weight_scale) + moe_name = name.replace(weight_name, param_name) + elif name.endswith(weight_name_base): + # Bare weight (no suffix) + moe_name = name.replace( + weight_name_base, param_name.rstrip("_") + "_weight" + ) + else: continue - moe_name = name.replace(weight_name, param_name) if moe_name not in params_dict: continue if is_pp_missing_parameter(moe_name, self): @@ -1334,15 +1366,12 @@ class Gemma4Model(nn.Module): # orientation for FusedMoE after _weight_iterator: # gate/up: [I, H] → w1/w3 expects [I, H] # down: [H, I] → w2 expects [H, I] - assert loaded_weight.dim() == 2, ( - f"Expected 2D expert weight for {weight_name}, " - f"got shape {loaded_weight.shape}" - ) + # Scales and other quantization params may be 1D or scalar. weight_loader = param.weight_loader weight_loader( param, loaded_weight, - weight_name + ".weight", + moe_name, # Pass mapped name (handles both weights and scales) shard_id=shard_id, expert_id=expert_id, ) @@ -1366,7 +1395,25 @@ class Gemma4Model(nn.Module): return loaded_params -class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts): +class Gemma4ForCausalLM( + nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3 +): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + # Gemma4ForConditionalGeneration already loads the text stack + # from `model.language_model.*`. We reuse that same checkpoint + # and adapter naming for the text-only Gemma4ForCausalLM path, + # so LoRA keys from the conditional wrapper map onto `model.*`. + "model.language_model.": "model.", + }, + orig_to_new_substr={ + # Gemma4ForConditionalGeneration names MoE adapter targets under + # `...moe.experts.*`, while the text-only model exposes them + # under `...moe.*`. + ".moe.experts.gate_up_proj": ".moe.gate_up_proj", + ".moe.experts.down_proj": ".moe.down_proj", + }, + ) # Note: qkv_proj packing applies to non-k_eq_v layers (sliding # attention and full attention without k_eq_v). k_eq_v layers use # separate q_proj + k_proj without packing. @@ -1448,7 +1495,7 @@ class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts): intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, **kwargs, - ) -> torch.Tensor | IntermediateTensors: + ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs ) @@ -1499,6 +1546,11 @@ class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts): ".moe.down_proj", ) + # Remap individual 2D expert weights: + # .experts.{id}.{proj} → .moe.experts.{id}.{proj} + # (This handles per-expert 2D quantized weights) + name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name) + # MoE expert weights: checkpoint stores as 3D packed # tensors. Explode into per-expert 2D weights for # FusedMoE weight_loader. diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index fa597fe96a0..e22f23c5c8b 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -64,7 +64,12 @@ from vllm.multimodal.processing.processor import ( from vllm.sequence import IntermediateTensors from vllm.utils.tensor_schema import TensorSchema, TensorShape -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from .interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) from .utils import ( AutoWeightsLoader, WeightsMapper, @@ -845,7 +850,12 @@ class Gemma4MultimodalEmbedder(nn.Module): info=Gemma4ProcessingInfo, dummy_inputs=Gemma4DummyInputsBuilder, ) -class Gemma4ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): +class Gemma4ForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsEagle3, +): packed_modules_mapping = { "qkv_proj": [ "q_proj", diff --git a/vllm/model_executor/models/gpt_oss.py b/vllm/model_executor/models/gpt_oss.py index a9ec8297422..4e4eb581842 100644 --- a/vllm/model_executor/models/gpt_oss.py +++ b/vllm/model_executor/models/gpt_oss.py @@ -560,6 +560,14 @@ class GptOssModel(nn.Module, EagleModelMixin): pcp_rank=get_pcp_group().rank_in_group, ) + def _is_mxfp4(weight_dtype: str | None) -> bool: + """Return True for any MXFP4 weight-dtype variant. + + Covers "gpt_oss_mxfp4" (GptOssMxfp4MoEMethod) and "mxfp4" + (QuarkMoEMethod with fp4 weights) and any future variants. + """ + return weight_dtype is not None and "mxfp4" in weight_dtype + def _get_moe_weight_dtype(layer_id: int = 0) -> str | None: """Helper function to get MoE quantization weight dtype. @@ -578,7 +586,7 @@ class GptOssModel(nn.Module, EagleModelMixin): moe_weight_dtype = _get_moe_weight_dtype(layer_id=0) - if moe_weight_dtype == "mxfp4": + if _is_mxfp4(moe_weight_dtype): # MXFP4 requires OCP_MX_BLOCK_SIZE alignment intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size) @@ -682,7 +690,7 @@ class GptOssModel(nn.Module, EagleModelMixin): continue # Unified handler for mxfp4 weights and scales - elif moe_quant_method == "mxfp4" and any( + elif _is_mxfp4(moe_quant_method) and any( name.endswith(suffix) for suffix in [ ".w13_weight_scale", @@ -1116,8 +1124,22 @@ class GptOssModel(nn.Module, EagleModelMixin): if hasattr(self.config, "quantization_config") else None ) - + # Normalize the checkpoint's quant_method to the internal name. + # Note: there are three places where "mxfp4" -> "gpt_oss_mxfp4" + # normalization occurs, each serving a different data path: + # 1. GptOssMxfp4Config.override_quantization_method() — sets + # ModelConfig.quantization (used to select the QuantizationConfig + # class at model init time), reading from model_arch_config which + # is a snapshot taken before verify_and_update_model_config runs. + # 2. GptOssForCausalLMConfig.verify_and_update_model_config() — + # patches hf_config.quantization_config in-place (a separate copy + # of the dict from model_arch_config) for later hf_config lookups. + # 3. Here — reads directly from self.config (the raw HF config) which + # may still carry the original "mxfp4" string from the checkpoint. if quant_method == "mxfp4": + quant_method = "gpt_oss_mxfp4" + + if quant_method == "gpt_oss_mxfp4": return self._load_weights_mxfp4( ep_rank_end, ep_rank_start, diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 1f1d57493c0..c24798e0840 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1098,6 +1098,12 @@ class SupportsTranscription(Protocol): :meth:`get_language_token_ids`. """ + no_space_languages: ClassVar[set[str]] = {"ja", "zh"} + """ + Languages that don't need a space between words. + For example, Japanese (ja) and Chinese (zh) don't need a space between words. + """ + def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) # language codes in supported_languages @@ -1518,6 +1524,13 @@ class SupportsEncoderCudaGraph(Protocol): def get_encoder_cudagraph_config(self) -> "EncoderCudaGraphConfig": ... + def get_input_modality( + self, + mm_kwargs: dict[str, Any], + ) -> str: + """Return the modality of the inputs.""" + ... + def get_encoder_cudagraph_budget_range( self, vllm_config: "VllmConfig", @@ -1530,7 +1543,7 @@ class SupportsEncoderCudaGraph(Protocol): (e.g. max_num_batched_tokens) Used when ``encoder_cudagraph_token_budgets`` and/or - ``encoder_cudagraph_max_images_per_batch`` are not explicitly + ``encoder_cudagraph_max_vision_items_per_batch`` are not explicitly specified by the user. """ ... @@ -1584,6 +1597,7 @@ class SupportsEncoderCudaGraph(Protocol): self, token_budget: int, max_batch_size: int, + max_frames_per_batch: int, device: torch.device, dtype: torch.dtype, ) -> "EncoderCudaGraphCaptureInputs": @@ -1594,6 +1608,7 @@ class SupportsEncoderCudaGraph(Protocol): self, mm_kwargs: dict[str, Any], max_batch_size: int, + max_frames_per_batch: int, ) -> "EncoderCudaGraphReplayBuffers": """Compute buffer values from actual batch inputs for replay.""" ... diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py new file mode 100644 index 00000000000..980502191dd --- /dev/null +++ b/vllm/model_executor/models/jina.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from https://huggingface.co/jinaai/jina-reranker-v3/blob/main/modeling.py +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.sequence import IntermediateTensors +from vllm.tasks import PoolingTask +from vllm.v1.pool.metadata import PoolingMetadata + +from ..layers.pooler import DispatchPooler +from ..layers.pooler.tokwise import ( + StepPool, + TokenPooler, + TokenPoolingMethodOutputItem, +) +from .interfaces import SupportsLateInteraction +from .qwen3 import Qwen3Model +from .utils import AutoWeightsLoader, maybe_prefix + + +class JinaForRanking(nn.Module, SupportsLateInteraction): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.config = config + self.projector_dim: int = config.embedding_size + + self.vllm_config = vllm_config + self.quant_config = quant_config + self.model = Qwen3Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + + self.projector = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size // 2, bias=False), + nn.ReLU(), + nn.Linear(config.hidden_size // 2, self.projector_dim, bias=False), + ) + + self.pooler = DispatchPooler( + { + "token_embed": TokenPooler( + pooling=JinaForRankingPool(self.projector), + ) + } + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self, skip_prefixes=(["lm_head."])) + return loader.load_weights(weights) + + +class JinaForRankingPool(StepPool): + def __init__(self, projector: nn.Sequential): + super().__init__() + + self.doc_token_id = 151670 + self.query_token_id = 151671 + self.projector = projector + + def get_supported_tasks(self) -> set[PoolingTask]: + return {"token_embed"} + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> list[TokenPoolingMethodOutputItem]: + pooled_data_lst = super().forward(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids() + + embeds_list = list[torch.Tensor | None]() + for data, token_ids in zip(pooled_data_lst, prompt_token_ids): + # for unfinished chunked prefill + if data is None: + embeds_list.append(None) + else: + docs_indexes = torch.where(torch.eq(token_ids, self.doc_token_id))[0] + query_indexes = torch.where(torch.eq(token_ids, self.query_token_id))[0] + + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + indexes = torch.cat([docs_indexes, query_indexes]) + embeds = self.projector(data[indexes]) + embeds_list.append(embeds) + + return embeds_list diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index a9b85f07355..eb7edfca640 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -113,7 +113,29 @@ class KimiK25ProcessingInfo(BaseProcessingInfo): trust_remote_code=self.ctx.model_config.trust_remote_code, ) - self.media_token_id = media_token_id = hf_config.media_placeholder_token_id + # Resolve token ID from the tokenizer because transformers v5 + # may remap token IDs vs config.json. + config_token_id = hf_config.media_placeholder_token_id + resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") + is_valid_resolved = isinstance(resolved_token_id, int) and ( + tokenizer.unk_token_id is None + or resolved_token_id != tokenizer.unk_token_id + ) + if is_valid_resolved and resolved_token_id != config_token_id: + logger.warning_once( + "Kimi-K2.5 config.media_placeholder_token_id (%d) disagrees " + "with tokenizer mapping for <|media_pad|> (%d). " + "Using tokenizer value.", + config_token_id, + resolved_token_id, + ) + media_token_id = resolved_token_id + # Patch config so downstream code also sees the correct ID. + hf_config.media_placeholder_token_id = resolved_token_id + else: + media_token_id = config_token_id + + self.media_token_id = media_token_id self.media_token = tokenizer.decode(media_token_id) self.image_processor = image_processor @@ -232,8 +254,7 @@ class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo]) hf_processor_mm_kwargs: Mapping[str, Any], out_mm_kwargs: MultiModalKwargsItems, ) -> Sequence[PromptUpdate]: - hf_config = self.info.get_hf_config() - media_token_id = hf_config.media_placeholder_token_id + media_token_id = self.info.media_token_id def get_replacement(item_idx: int): media = mm_items.get_items("vision_chunk", (VisionChunkProcessorItems,)) diff --git a/vllm/model_executor/models/longcat_flash.py b/vllm/model_executor/models/longcat_flash.py index a9e2c2268ee..375b0b69b1f 100644 --- a/vllm/model_executor/models/longcat_flash.py +++ b/vllm/model_executor/models/longcat_flash.py @@ -46,7 +46,7 @@ from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import FusedMoE, ZeroExpertFusedMoE +from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -292,12 +292,10 @@ class LongcatMoe(nn.Module): prefix=f"{prefix}.gate", ) - assert config.zero_expert_num is not None assert config.zero_expert_type is not None - self.experts = ZeroExpertFusedMoE( - zero_expert_num=config.zero_expert_num, + self.experts = FusedMoE( zero_expert_type=config.zero_expert_type, - router=self.router, + e_score_correction_bias=self.router.e_score_correction_bias, num_experts=num_experts, top_k=top_k, hidden_size=hidden_size, @@ -332,7 +330,7 @@ class LongcatMoe(nn.Module): hidden_states_padded.to(self.router_params_dtype) ) - # ZeroExpertFusedMoE handles routing memoization and zero expert computation + # FusedMoE handles routing memoization and zero expert computation # internally. Pass full router_logits (including zero experts) so that # zero experts can be properly identified in routing. final_hidden_states = self.experts( diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index 79162eef3f6..cda07ea291e 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -1050,9 +1050,17 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): quant_config=quant_config, prefix=maybe_prefix(prefix, "resampler"), ) + self._resampler_moved = False self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors + def _ensure_resampler_device(self) -> None: + if self._resampler_moved: + return + # Only move device, DO NOT touch dtype (fp8 quant needs its own dtype) + self.resampler.to(current_platform.device_type) + self._resampler_moved = True + def _parse_and_validate_vision_input( self, modality: str, @@ -1171,7 +1179,9 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded def get_mm_mapping(self) -> MultiModelKeys: """ @@ -1276,9 +1286,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel): prefix=prefix, ) - return resampler.to( - device=current_platform.device_type, dtype=torch.get_default_dtype() - ) + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1359,9 +1367,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA): prefix=prefix, ) - return resampler.to( - device=current_platform.device_type, dtype=torch.get_default_dtype() - ) + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1452,11 +1458,8 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): quant_config=quant_config, prefix=prefix, ) - target_device = current_platform.device_type - target_dtype = torch.get_default_dtype() - if any(p.is_meta for p in resampler.parameters()): - return resampler.to_empty(device=target_device).to(dtype=target_dtype) - return resampler.to(device=target_device, dtype=target_dtype) + + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1491,7 +1494,9 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): @@ -1551,10 +1556,7 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): quant_config=quant_config, prefix=prefix, ) - - return resampler.to( - device=current_platform.device_type, dtype=torch.get_default_dtype() - ) + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1589,7 +1591,9 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): @@ -1649,11 +1653,8 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): quant_config=quant_config, prefix=prefix, ) - target_device = current_platform.device_type - target_dtype = torch.get_default_dtype() - if any(p.is_meta for p in resampler.parameters()): - return resampler.to_empty(device=target_device).to(dtype=target_dtype) - return resampler.to(device=target_device, dtype=target_dtype) + + return resampler.to(dtype=torch.get_default_dtype()) def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor: pixel_values = data["pixel_values"] @@ -1692,7 +1693,9 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"]) - return loader.load_weights(weights) + loaded = loader.load_weights(weights) + self._ensure_resampler_device() + return loaded _SUPPORT_VERSION = { diff --git a/vllm/model_executor/models/minimax_m2.py b/vllm/model_executor/models/minimax_m2.py index f10452c5738..2a9e5f4e08a 100644 --- a/vllm/model_executor/models/minimax_m2.py +++ b/vllm/model_executor/models/minimax_m2.py @@ -233,9 +233,7 @@ class MiniMaxM2Attention(nn.Module): ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = MiniMaxText01RMSNormTP.forward_qk( - self.q_norm, self.k_norm, q.contiguous(), k.contiguous() - ) + q, k = MiniMaxText01RMSNormTP.forward_qk(self.q_norm, self.k_norm, q, k) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 0ece3dda2e5..025ce564083 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -382,7 +382,14 @@ class Mistral3ForConditionalGeneration( # Some PEFT LoRAs are trained against the text submodule directly # and produce names like `base_model.model.model.layers.*`. "model.": "language_model.model.", - } + }, + orig_to_new_suffix={ + # FP8 quantized HF checkpoints use "activation_scale" and + # "weight_scale_inv" but vLLM's FP8 linear layers register + # them as "input_scale" and "weight_scale" + ".activation_scale": ".input_scale", + ".weight_scale_inv": ".weight_scale", + }, ) @classmethod @@ -402,13 +409,8 @@ class Mistral3ForConditionalGeneration( self.config = config self.multimodal_config = multimodal_config - # NOTE: These are special cases for Pixtral-12B in the HF-format + # NOTE: This is a special case for Pixtral-12B in the HF-format # https://huggingface.co/mistral-community/pixtral-12b/blob/main/config.json # noqa - if ( - config.text_config.architectures is None - and config.text_config.model_type == "mistral" - ): - config.text_config.architectures = ["MistralForCausalLM"] if ( config.projector_hidden_act is None and config.vision_config.hidden_act == "gelu" diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index 9983015b0ee..b0424675943 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -37,6 +37,7 @@ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM from vllm.model_executor.models.parakeet import ParakeetExtractor, ProjectedParakeet from vllm.model_executor.models.radio import RadioModel, calc_seq_lens from vllm.model_executor.models.utils import ( + WeightsMapper, init_vllm_registered_model, maybe_prefix, ) @@ -597,19 +598,26 @@ class NanoNemotronVLMultiModalProcessor( def _extract_audio_from_videos( self, mm_items: MultiModalDataItems, - ) -> tuple[MultiModalDataItems, list[AudioItem]]: + ) -> tuple[MultiModalDataItems, list[AudioItem], list[bool]]: """Extract audio tracks from video bytes in *mm_items*. + Videos whose bytes are missing or that contain no audio stream are + silently skipped. The returned *has_audio* mask is aligned with + the video list so callers know which ``