forked from Karylab-cklius/vllm
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b2e9eff8ae |
@@ -638,93 +638,9 @@ steps:
|
||||
depends_on:
|
||||
- step: upload-rocm-wheels
|
||||
allow_failure: true
|
||||
- step: input-release-version
|
||||
allow_failure: true
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-rocm-release.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
|
||||
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
|
||||
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
|
||||
- block: "Generate Root Index for ROCm Wheels for Release"
|
||||
key: block-generate-root-index-rocm-wheels
|
||||
depends_on: upload-rocm-wheels
|
||||
|
||||
- label: ":package: Generate Root Index for ROCm Wheels for Release"
|
||||
depends_on: block-generate-root-index-rocm-wheels
|
||||
id: generate-root-index-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm700"
|
||||
|
||||
# ROCm Job 5: Build ROCm Release Docker Image
|
||||
- label: ":rocm: :docker: Build ROCm Release Docker Image"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Download Docker image from S3 (set by build-rocm-base-wheels)
|
||||
DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')"
|
||||
if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then
|
||||
echo "ERROR: rocm-docker-image-s3-path metadata not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Downloading base image from $${DOCKER_IMAGE_S3_PATH}"
|
||||
mkdir -p artifacts/rocm-docker-image
|
||||
aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz
|
||||
|
||||
# Load base Docker image
|
||||
echo "Loading base Docker image..."
|
||||
LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load)
|
||||
BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //')
|
||||
echo "Loaded base image: $${BASE_IMAGE_TAG}"
|
||||
|
||||
# Tag and push the base image to ECR
|
||||
docker tag "$${BASE_IMAGE_TAG}" public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base
|
||||
echo "Pushed base image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base"
|
||||
|
||||
# Get GPU architectures from meta-data
|
||||
PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')"
|
||||
PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}"
|
||||
|
||||
# Build vLLM ROCm release image using cached base
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \
|
||||
--build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.rocm .
|
||||
|
||||
# Push to ECR
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
|
||||
echo "Pushed: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
@@ -32,7 +32,6 @@ To download and upload the image:
|
||||
\`\`\`
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
|
||||
@@ -47,17 +46,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai:rocm
|
||||
docker tag vllm/vllm-openai:rocm vllm/vllm-openai:latest-rocm
|
||||
docker tag vllm/vllm-openai:rocm vllm/vllm-openai:v${RELEASE_VERSION}-rocm
|
||||
docker push vllm/vllm-openai:latest-rocm
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-rocm
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest
|
||||
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
|
||||
|
||||
@@ -3,32 +3,25 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Generate Buildkite annotation for ROCm wheel release
|
||||
|
||||
set -ex
|
||||
|
||||
# Get build configuration from meta-data
|
||||
# Extract ROCm version dynamically from Dockerfile.rocm_base
|
||||
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0"
|
||||
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.1-complete -> extracts "7.1"
|
||||
ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown")
|
||||
PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12")
|
||||
PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151")
|
||||
|
||||
# TODO: Enable the nightly build for ROCm
|
||||
# Get release version, default to 1.0.0.dev for nightly/per-commit builds
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "")
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
RELEASE_VERSION="1.0.0.dev"
|
||||
fi
|
||||
|
||||
# S3 URLs
|
||||
S3_BUCKET="${S3_BUCKET:-vllm-wheels}"
|
||||
S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
|
||||
S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
|
||||
S3_URL="https://${S3_BUCKET}.s3.${S3_REGION}.amazonaws.com"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}"
|
||||
|
||||
# Format ROCm version for path (e.g., "7.1" -> "rocm710")
|
||||
ROCM_VERSION_PATH="rocm$(echo ${ROCM_VERSION} | tr -d '.')"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}"
|
||||
buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF
|
||||
## ROCm Wheel and Docker Image Releases
|
||||
## :rocm: ROCm Wheel Release
|
||||
|
||||
### Build Configuration
|
||||
| Setting | Value |
|
||||
|---------|-------|
|
||||
@@ -41,72 +34,41 @@ buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' <<
|
||||
### :package: Installation
|
||||
|
||||
**Install from this build (by commit):**
|
||||
|
||||
\`\`\`bash
|
||||
pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/{rocm_variant}/
|
||||
|
||||
# Example for ROCm ${ROCM_VERSION}:
|
||||
pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
# Example:
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/rocm700/
|
||||
\`\`\`
|
||||
|
||||
**Install from nightly (if published):**
|
||||
|
||||
\`\`\`bash
|
||||
pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/
|
||||
\`\`\`
|
||||
|
||||
### :floppy_disk: Download Wheels Directly
|
||||
|
||||
\`\`\`bash
|
||||
# List all ROCm wheels
|
||||
aws s3 ls s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/
|
||||
aws s3 ls s3://${S3_BUCKET}/${ROCM_PATH}/
|
||||
|
||||
# Download specific wheels
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/vllm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torch-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-kernels-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/vllm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torch-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/triton_rocm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torchvision-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/amdsmi-*.whl .
|
||||
\`\`\`
|
||||
|
||||
### :gear: Included Packages
|
||||
- **vllm**: vLLM with ROCm support
|
||||
- **torch**: PyTorch built for ROCm ${ROCM_VERSION}
|
||||
- **triton**: Triton
|
||||
- **triton-kernels**: Triton kernels
|
||||
- **triton_rocm**: Triton built for ROCm
|
||||
- **torchvision**: TorchVision for ROCm PyTorch
|
||||
- **torchaudio**: Torchaudio for ROCm PyTorch
|
||||
- **amdsmi**: AMD SMI Python bindings
|
||||
- **aiter**: Aiter for ROCm
|
||||
- **flash-attn**: Flash Attention for ROCm
|
||||
|
||||
### :warning: Notes
|
||||
- These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs
|
||||
- Supported GPU architectures: ${PYTORCH_ROCM_ARCH}
|
||||
- Platform: Linux x86_64 only
|
||||
|
||||
### :package: Docker Image Release
|
||||
|
||||
To download and upload the image:
|
||||
|
||||
\`\`\`
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
\`\`\`
|
||||
|
||||
EOF
|
||||
|
||||
@@ -24,7 +24,6 @@ if [ "$RELEASE_VERSION" != "$GIT_VERSION" ]; then
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
PURE_VERSION=${RELEASE_VERSION#v} # remove leading 'v'
|
||||
|
||||
# check pypi token
|
||||
if [ -z "$PYPI_TOKEN" ]; then
|
||||
@@ -82,16 +81,16 @@ echo "Existing wheels on S3:"
|
||||
aws s3 ls "$S3_COMMIT_PREFIX"
|
||||
echo "Copying wheels to local directory"
|
||||
mkdir -p $DIST_DIR
|
||||
# include only wheels for the release version, ignore all files with "dev" or "rc" in the name (without excluding 'aarch64')
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
|
||||
# include only wheels for the release version, ignore all files with "dev" or "rc" in the name
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm-${RELEASE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
|
||||
echo "Wheels copied to local directory"
|
||||
# generate source tarball
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${RELEASE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
ls -la $DIST_DIR
|
||||
|
||||
|
||||
# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
|
||||
PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${PURE_VERSION}*.whl" -not -name "*+*")
|
||||
PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${RELEASE_VERSION}*.whl" -not -name "*+*")
|
||||
if [ -z "$PYPI_WHEEL_FILES" ]; then
|
||||
echo "No default variant wheels found, quitting..."
|
||||
exit 1
|
||||
|
||||
@@ -29,9 +29,8 @@ 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/cpu.txt --index-strategy unsafe-best-match
|
||||
uv pip install -e . --no-build-isolation
|
||||
uv pip install -e .
|
||||
env:
|
||||
CMAKE_BUILD_PARALLEL_LEVEL: 4
|
||||
|
||||
|
||||
+9
-9
@@ -377,7 +377,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
# preselected input type pairs and schedules.
|
||||
# Generate sources:
|
||||
set(MARLIN_GEN_SCRIPT
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/marlin/generate_kernels.py)
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/gptq_marlin/generate_kernels.py)
|
||||
file(MD5 ${MARLIN_GEN_SCRIPT} MARLIN_GEN_SCRIPT_HASH)
|
||||
list(JOIN CUDA_ARCHS "," CUDA_ARCHS_STR)
|
||||
set(MARLIN_GEN_SCRIPT_HASH_AND_ARCH "${MARLIN_GEN_SCRIPT_HASH}(ARCH:${CUDA_ARCHS_STR})")
|
||||
@@ -412,7 +412,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
if (MARLIN_ARCHS)
|
||||
file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_float16.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_float16.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_ARCHS}")
|
||||
@@ -422,7 +422,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
|
||||
|
||||
file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/marlin/sm80_kernel_*_bfloat16.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_BF16_KERNEL_SRC "csrc/quantization/gptq_marlin/sm80_kernel_*_bfloat16.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_BF16_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_BF16_ARCHS}")
|
||||
@@ -434,7 +434,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
if (MARLIN_SM75_ARCHS)
|
||||
file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/gptq_marlin/sm75_kernel_*.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_SM75_ARCHS}")
|
||||
@@ -446,7 +446,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
if (MARLIN_FP8_ARCHS)
|
||||
file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
|
||||
file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/gptq_marlin/sm89_kernel_*.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
|
||||
CUDA_ARCHS "${MARLIN_FP8_ARCHS}")
|
||||
@@ -459,10 +459,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
|
||||
set(MARLIN_SRCS
|
||||
"csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu"
|
||||
"csrc/quantization/marlin/marlin.cu"
|
||||
"csrc/quantization/marlin/marlin_int4_fp8_preprocess.cu"
|
||||
"csrc/quantization/marlin/gptq_marlin_repack.cu"
|
||||
"csrc/quantization/marlin/awq_marlin_repack.cu")
|
||||
"csrc/quantization/gptq_marlin/gptq_marlin.cu"
|
||||
"csrc/quantization/gptq_marlin/marlin_int4_fp8_preprocess.cu"
|
||||
"csrc/quantization/gptq_marlin/gptq_marlin_repack.cu"
|
||||
"csrc/quantization/gptq_marlin/awq_marlin_repack.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${MARLIN_SRCS}"
|
||||
CUDA_ARCHS "${MARLIN_OTHER_ARCHS}")
|
||||
|
||||
@@ -20,12 +20,8 @@ FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
PROVIDER_CFGS = {
|
||||
"vllm": dict(backend="vllm", is_sf_swizzled_layout=False, enabled=True),
|
||||
"vllm-swizzle": dict(backend="vllm", is_sf_swizzled_layout=True, enabled=True),
|
||||
"flashinfer": dict(backend="flashinfer", is_sf_swizzled_layout=False, enabled=True),
|
||||
"flashinfer-swizzle": dict(
|
||||
backend="flashinfer", is_sf_swizzled_layout=True, enabled=True
|
||||
),
|
||||
"vllm": dict(backend="vllm", enabled=True),
|
||||
"flashinfer": dict(backend="flashinfer", enabled=True),
|
||||
}
|
||||
|
||||
_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
|
||||
@@ -40,7 +36,7 @@ def compute_global_scale(tensor: torch.Tensor) -> torch.Tensor:
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
|
||||
x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096],
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=_enabled,
|
||||
@@ -67,36 +63,19 @@ def benchmark(batch_size, provider, N, K):
|
||||
|
||||
if cfg["backend"] == "vllm":
|
||||
# vLLM's FP4 quantization
|
||||
if cfg["is_sf_swizzled_layout"]:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.scaled_fp4_quant(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
else:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.scaled_fp4_quant(
|
||||
a, a_global_scale, is_sf_swizzled_layout=False
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.scaled_fp4_quant(a, a_global_scale),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
elif cfg["backend"] == "flashinfer":
|
||||
# FlashInfer's FP4 quantization
|
||||
if cfg["is_sf_swizzled_layout"]:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
else:
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=False
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
# Use is_sf_swizzled_layout=True to match vLLM's output format
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
|
||||
# Convert ms to us for better readability at small batch sizes
|
||||
to_us = lambda t_ms: t_ms * 1000
|
||||
@@ -113,9 +92,7 @@ def prepare_shapes(args):
|
||||
return out
|
||||
|
||||
|
||||
def _test_accuracy_once(
|
||||
M: int, K: int, dtype: torch.dtype, device: str, is_sf_swizzled_layout: bool
|
||||
):
|
||||
def _test_accuracy_once(M: int, K: int, dtype: torch.dtype, device: str):
|
||||
"""Test accuracy between vLLM and FlashInfer FP4 quantization."""
|
||||
# Create input tensor
|
||||
a = torch.randn((M, K), device=device, dtype=dtype)
|
||||
@@ -124,13 +101,11 @@ def _test_accuracy_once(
|
||||
a_global_scale = compute_global_scale(a)
|
||||
|
||||
# vLLM quantization
|
||||
vllm_fp4, vllm_scale = ops.scaled_fp4_quant(
|
||||
a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
|
||||
)
|
||||
vllm_fp4, vllm_scale = ops.scaled_fp4_quant(a, a_global_scale)
|
||||
|
||||
# FlashInfer quantization (with swizzled layout to match vLLM's output)
|
||||
flashinfer_fp4, flashinfer_scale = flashinfer_fp4_quantize(
|
||||
a, a_global_scale, is_sf_swizzled_layout=is_sf_swizzled_layout
|
||||
a, a_global_scale, is_sf_swizzled_layout=True
|
||||
)
|
||||
flashinfer_scale = flashinfer_scale.view(torch.float8_e4m3fn)
|
||||
|
||||
@@ -139,14 +114,7 @@ def _test_accuracy_once(
|
||||
vllm_fp4,
|
||||
flashinfer_fp4,
|
||||
)
|
||||
# Compare scales
|
||||
torch.testing.assert_close(
|
||||
vllm_scale,
|
||||
flashinfer_scale,
|
||||
)
|
||||
print(
|
||||
f"M={M}, K={K}, dtype={dtype}, is_sf_swizzled_layout={is_sf_swizzled_layout}: PASSED" # noqa: E501
|
||||
)
|
||||
print(f"M={M}, K={K}, dtype={dtype}: PASSED")
|
||||
|
||||
|
||||
def test_accuracy():
|
||||
@@ -162,10 +130,9 @@ def test_accuracy():
|
||||
Ms = [1, 1024]
|
||||
Ks = [4096]
|
||||
|
||||
for is_sf_swizzled_layout in [True, False]:
|
||||
for M in Ms:
|
||||
for K in Ks:
|
||||
_test_accuracy_once(M, K, dtype, device, is_sf_swizzled_layout)
|
||||
for M in Ms:
|
||||
for K in Ks:
|
||||
_test_accuracy_once(M, K, dtype, device)
|
||||
|
||||
print("\nAll accuracy tests passed!")
|
||||
|
||||
@@ -178,7 +145,7 @@ if __name__ == "__main__":
|
||||
"--models",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=["meta-llama/Llama-3.3-70B-Instruct"],
|
||||
default=["meta-llama/Llama-3.1-8B-Instruct"],
|
||||
choices=list(WEIGHT_SHAPES.keys()),
|
||||
)
|
||||
parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
|
||||
|
||||
@@ -231,7 +231,7 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
|
||||
assert bt.w_tok_s is None
|
||||
assert bt.group_size is not None
|
||||
|
||||
fn = lambda: ops.marlin_gemm(
|
||||
fn = lambda: ops.gptq_marlin_gemm(
|
||||
a=bt.a,
|
||||
c=None,
|
||||
b_q_weight=w_q,
|
||||
|
||||
@@ -239,7 +239,7 @@ def bench_run(
|
||||
"sm_version": sm_version,
|
||||
"CUBLAS_M_THRESHOLD": CUBLAS_M_THRESHOLD,
|
||||
# Kernels
|
||||
"marlin_gemm": ops.marlin_gemm,
|
||||
"gptq_marlin_gemm": ops.gptq_marlin_gemm,
|
||||
"gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
|
||||
"gptq_marlin_repack": ops.gptq_marlin_repack,
|
||||
"allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
|
||||
@@ -263,21 +263,21 @@ def bench_run(
|
||||
|
||||
results.append(
|
||||
benchmark.Timer(
|
||||
stmt="output = marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501
|
||||
stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, False, False)", # noqa: E501
|
||||
globals=globals,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description="marlin_gemm",
|
||||
description="gptq_marlin_gemm",
|
||||
).blocked_autorange(min_run_time=min_run_time)
|
||||
)
|
||||
|
||||
results.append(
|
||||
benchmark.Timer(
|
||||
stmt="output = marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501
|
||||
stmt="output = gptq_marlin_gemm(a, None, marlin_q_w, marlin_s, None, marlin_s2, marlin_zp, marlin_g_idx, marlin_sort_indices, marlin_workspace.scratch, quant_type, size_m, size_n, size_k, is_k_full, False, True, False)", # noqa: E501
|
||||
globals=globals,
|
||||
label=label,
|
||||
sub_label=sub_label,
|
||||
description="marlin_gemm_fp32",
|
||||
description="gptq_marlin_gemm_fp32",
|
||||
).blocked_autorange(min_run_time=min_run_time)
|
||||
)
|
||||
|
||||
|
||||
@@ -13,8 +13,6 @@ endif()
|
||||
#
|
||||
# Define environment variables for special configurations
|
||||
#
|
||||
set(ENABLE_AVX2 $ENV{VLLM_CPU_AVX2})
|
||||
set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
|
||||
set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
|
||||
set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
|
||||
set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
|
||||
@@ -105,16 +103,6 @@ else()
|
||||
find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
|
||||
find_isa(${CPUINFO} "S390" S390_FOUND)
|
||||
find_isa(${CPUINFO} "v" RVV_FOUND) # Check for RISC-V RVV support
|
||||
|
||||
# Support cross-compilation by allowing override via environment variables
|
||||
if (ENABLE_AVX2)
|
||||
set(AVX2_FOUND ON)
|
||||
message(STATUS "AVX2 support enabled via VLLM_CPU_AVX2 environment variable")
|
||||
endif()
|
||||
if (ENABLE_AVX512)
|
||||
set(AVX512_FOUND ON)
|
||||
message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (AVX512_FOUND AND NOT AVX512_DISABLED)
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
|
||||
#endif
|
||||
|
||||
#include "quantization/marlin/marlin.cuh"
|
||||
#include "quantization/marlin/marlin_dtypes.cuh"
|
||||
#include "quantization/gptq_marlin/marlin.cuh"
|
||||
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
|
||||
#include "core/scalar_type.hpp"
|
||||
|
||||
#define MARLIN_KERNEL_PARAMS \
|
||||
|
||||
@@ -23,10 +23,10 @@
|
||||
#define MARLIN_NAMESPACE_NAME marlin_moe_wna16
|
||||
#endif
|
||||
|
||||
#include "quantization/marlin/marlin.cuh"
|
||||
#include "quantization/marlin/marlin_dtypes.cuh"
|
||||
#include "quantization/marlin/dequant.h"
|
||||
#include "quantization/marlin/marlin_mma.h"
|
||||
#include "quantization/gptq_marlin/marlin.cuh"
|
||||
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
|
||||
#include "quantization/gptq_marlin/dequant.h"
|
||||
#include "quantization/gptq_marlin/marlin_mma.h"
|
||||
#include "core/scalar_type.hpp"
|
||||
|
||||
#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \
|
||||
|
||||
+1
-2
@@ -293,8 +293,7 @@ std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_scale,
|
||||
torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
torch::Tensor const& input_scale);
|
||||
|
||||
void scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
|
||||
@@ -27,24 +27,17 @@
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols,
|
||||
int32_t num_padded_cols,
|
||||
Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out,
|
||||
uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = vllm::PackedVec<Type>;
|
||||
__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
float const* SFScale, uint32_t* out,
|
||||
uint32_t* SFout) {
|
||||
using PackedVec = PackedVec<Type>;
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
@@ -56,60 +49,34 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
float const SFScaleVal = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0];
|
||||
|
||||
// Input tensor row/col loops.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
|
||||
if (colIdx < num_padded_cols) {
|
||||
PackedVec in_vec;
|
||||
PackedVec in_vec2;
|
||||
for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD;
|
||||
colIdx += blockDim.x) {
|
||||
int64_t inOffset =
|
||||
rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
int64_t inOffset2 = rowIdx * (numCols * 2 / CVT_FP4_ELTS_PER_THREAD) +
|
||||
numCols / CVT_FP4_ELTS_PER_THREAD + colIdx;
|
||||
PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
PackedVec in_vec2 = reinterpret_cast<PackedVec const*>(in)[inOffset2];
|
||||
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 8],
|
||||
valid);
|
||||
} else {
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec2, &reinterpret_cast<const uint32_t*>(in)[inOffset2 * 4],
|
||||
valid);
|
||||
}
|
||||
// Get the output tensor offset.
|
||||
// Same as inOffset because 8 elements are packed into one uint32_t.
|
||||
int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
auto& out_pos = out[outOffset];
|
||||
|
||||
// Compute silu and mul
|
||||
PackedVec out_silu_mul = compute_silu_mul<Type>(in_vec, in_vec2);
|
||||
PackedVec out_silu_mul = compute_silu_mul(in_vec, in_vec2);
|
||||
|
||||
auto sf_out =
|
||||
cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx, colIdx, numKTiles, SFout);
|
||||
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
out_silu_mul, SFScaleVal, sf_out);
|
||||
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(out_silu_mul, SFScaleVal,
|
||||
sf_out);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -136,23 +103,17 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
|
||||
auto output_ptr = static_cast<int64_t*>(output.data_ptr());
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 1024));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
|
||||
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
|
||||
|
||||
int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
|
||||
int grid_x = std::min(
|
||||
int(m), std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
dim3 grid(std::min(int(m), multiProcessorCount * numBlocksPerSM));
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
|
||||
m, n, sf_n_unpadded, input_ptr, input_sf_ptr,
|
||||
m, n, input_ptr, input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
|
||||
@@ -140,8 +140,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -246,8 +246,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
|
||||
CVT_FP4_NUM_THREADS_PER_SF>(
|
||||
rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
|
||||
|
||||
out_pos = cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
quant_input, SFScaleVal, sf_out);
|
||||
out_pos =
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -21,8 +21,7 @@
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout);
|
||||
torch::Tensor const& input_sf);
|
||||
#endif
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
@@ -52,12 +51,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
#endif
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf) {
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
|
||||
is_sf_swizzled_layout);
|
||||
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf);
|
||||
#endif
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
@@ -27,23 +27,29 @@
|
||||
|
||||
#include "cuda_utils.h"
|
||||
#include "launch_bounds_utils.h"
|
||||
|
||||
// Define before including nvfp4_utils.cuh so the header
|
||||
// can use this macro during compilation.
|
||||
#define NVFP4_ENABLE_ELTS16 1
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
template <typename Int>
|
||||
__host__ __device__ inline Int round_up(Int x, Int y) {
|
||||
static_assert(std::is_integral_v<Int>,
|
||||
"round_up argument must be integral type");
|
||||
return ((x + y - 1) / y) * y;
|
||||
}
|
||||
|
||||
// Compute effective rows for grid configuration with swizzled SF layouts.
|
||||
inline int computeEffectiveRows(int m) {
|
||||
constexpr int ROW_TILE = 128;
|
||||
return round_up(m, ROW_TILE);
|
||||
}
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t num_padded_cols,
|
||||
Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = vllm::PackedVec<Type>;
|
||||
|
||||
cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
|
||||
float const* SFScale, uint32_t* out, uint32_t* SFout) {
|
||||
using PackedVec = PackedVec<Type>;
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
@@ -53,31 +59,33 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
int32_t const numKTiles = (numCols + 63) / 64;
|
||||
|
||||
int sf_m = round_up<int>(numRows, 128);
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
int sf_n_unpadded = numCols / CVT_FP4_SF_VEC_SIZE;
|
||||
int sf_n_int = round_up<int>(sf_n_unpadded, 4) / 4;
|
||||
int num_padded_cols = sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE;
|
||||
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
float const global_scale = SFScale == nullptr ? 1.0f : SFScale[0];
|
||||
|
||||
// Iterate over all rows and cols including padded ones -
|
||||
// ensures we visit every single scale factor address to initialize it.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) {
|
||||
if (colIdx < num_padded_cols) {
|
||||
for (int colIdx = threadIdx.x;
|
||||
colIdx < num_padded_cols / CVT_FP4_ELTS_PER_THREAD;
|
||||
colIdx += blockDim.x) {
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
PackedVec in_vec;
|
||||
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
|
||||
// If we are outside valid rows OR outside valid columns -> Use Zeros
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
if (rowIdx >= numRows || elem_idx >= numCols) {
|
||||
memset(&in_vec, 0, sizeof(PackedVec));
|
||||
|
||||
} else {
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
// Valid Region: Load actual data
|
||||
in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
|
||||
}
|
||||
|
||||
auto sf_out =
|
||||
@@ -86,85 +94,13 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
rowIdx, colIdx, numKTiles, SFout);
|
||||
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
in_vec, global_scale, sf_out);
|
||||
cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, global_scale, sf_out);
|
||||
|
||||
// We do NOT write output for padding because the 'out' tensor is not
|
||||
// padded.
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Use UE4M3 by default.
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols,
|
||||
int32_t sf_n_unpadded, Type const* __restrict__ in,
|
||||
float const* __restrict__ SFScale,
|
||||
uint32_t* __restrict__ out,
|
||||
uint32_t* __restrict__ SFout) {
|
||||
using PackedVec = PackedVec<Type>;
|
||||
|
||||
static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
|
||||
(CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
|
||||
static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
|
||||
"Vec size is not matched.");
|
||||
|
||||
int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x;
|
||||
int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
// Get the global scaling factor, which will be applied to the SF.
|
||||
// Note SFScale is the same as next GEMM's alpha, which is
|
||||
// (448.f / (Alpha_A / 6.f)).
|
||||
float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0];
|
||||
|
||||
// Iterate over all rows and cols including padded ones -
|
||||
// ensures we visit every single scale factor address to initialize it.
|
||||
for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
|
||||
if (colIdx < sf_n_unpadded) {
|
||||
PackedVec in_vec;
|
||||
int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
|
||||
|
||||
// If we are outside valid rows OR outside valid columns -> Use Zeros
|
||||
bool valid = (rowIdx < numRows) && (elem_idx < numCols);
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
ld256_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 8],
|
||||
valid);
|
||||
} else {
|
||||
ld128_or_zero_cg_u32<Type>(
|
||||
in_vec, &reinterpret_cast<const uint32_t*>(in)[inOffset * 4],
|
||||
valid);
|
||||
}
|
||||
|
||||
auto sf_out =
|
||||
sf_out_rowmajor_u8<uint32_t>(rowIdx, colIdx, sf_n_unpadded, SFout);
|
||||
|
||||
auto out_val =
|
||||
cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>(
|
||||
in_vec, global_scale, sf_out);
|
||||
|
||||
// We do NOT write output for padding because the 'out' tensor is not
|
||||
// padded.
|
||||
if (valid) {
|
||||
if constexpr (CVT_FP4_PACK16) {
|
||||
int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2;
|
||||
uint64_t packed64 =
|
||||
(uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo);
|
||||
reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64;
|
||||
} else {
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
if (rowIdx < numRows && elem_idx < numCols) {
|
||||
// Same as inOffset because 8 elements are packed into one uint32_t.
|
||||
out[inOffset] = out_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -175,8 +111,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
torch::Tensor const& input,
|
||||
torch::Tensor const& output_sf,
|
||||
torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
torch::Tensor const& input_sf) {
|
||||
int32_t m = input.size(0);
|
||||
int32_t n = input.size(1);
|
||||
|
||||
@@ -194,48 +129,19 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
|
||||
auto stream = at::cuda::getCurrentCUDAStream(input.get_device());
|
||||
|
||||
int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE);
|
||||
|
||||
// Grid, Block size. Each thread converts 8 values.
|
||||
dim3 block(std::min(int(n / ELTS_PER_THREAD), 512));
|
||||
int const numBlocksPerSM =
|
||||
vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
|
||||
int effectiveRows = vllm::computeEffectiveRows(m);
|
||||
dim3 grid(std::min(effectiveRows, multiProcessorCount * numBlocksPerSM));
|
||||
|
||||
if (is_sf_swizzled_layout) {
|
||||
int sf_n_int = int(vllm::round_up(sf_n_unpadded, 4) / 4);
|
||||
int32_t num_padded_cols =
|
||||
sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
|
||||
|
||||
int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x));
|
||||
int grid_x =
|
||||
std::min(vllm::computeEffectiveRows(m),
|
||||
std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
|
||||
m, n, num_padded_cols, input_ptr, input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
} else {
|
||||
int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast<int>(block.x));
|
||||
int grid_x = std::min(
|
||||
m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y));
|
||||
dim3 grid(grid_x, grid_y);
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>
|
||||
<<<grid, block, 0, stream>>>(m, n, sf_n_unpadded, input_ptr,
|
||||
input_sf_ptr,
|
||||
reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
}
|
||||
}
|
||||
VLLM_DISPATCH_HALF_TYPES(input.scalar_type(), "nvfp4_quant_kernel", [&] {
|
||||
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
|
||||
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
|
||||
// NOTE: We don't support e8m0 scales at this moment.
|
||||
vllm::cvt_fp16_to_fp4<cuda_type, false><<<grid, block, 0, stream>>>(
|
||||
m, n, input_ptr, input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr),
|
||||
reinterpret_cast<uint32_t*>(sf_out));
|
||||
});
|
||||
}
|
||||
@@ -19,17 +19,9 @@
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
|
||||
#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
|
||||
defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
|
||||
#define ELTS_PER_THREAD 16
|
||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 16;
|
||||
constexpr bool CVT_FP4_PACK16 = true;
|
||||
#else
|
||||
#define ELTS_PER_THREAD 8
|
||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
||||
constexpr bool CVT_FP4_PACK16 = false;
|
||||
#endif
|
||||
#define ELTS_PER_THREAD 8
|
||||
|
||||
constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
|
||||
constexpr int CVT_FP4_SF_VEC_SIZE = 16;
|
||||
|
||||
namespace vllm {
|
||||
@@ -76,46 +68,19 @@ struct TypeConverter<__nv_bfloat16> {
|
||||
using Type = __nv_bfloat162;
|
||||
};
|
||||
|
||||
#if (defined(NVFP4_ENABLE_ELTS16) && (CUDART_VERSION >= 12090) && \
|
||||
defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100)
|
||||
// Define a 32 bytes packed data type.
|
||||
template <class Type>
|
||||
struct alignas(32) PackedVec {
|
||||
typename TypeConverter<Type>::Type elts[8];
|
||||
};
|
||||
#else
|
||||
// Define a 16 bytes packed data type.
|
||||
template <class Type>
|
||||
struct alignas(16) PackedVec {
|
||||
struct PackedVec {
|
||||
typename TypeConverter<Type>::Type elts[4];
|
||||
};
|
||||
#endif
|
||||
|
||||
template <>
|
||||
struct PackedVec<__nv_fp8_e4m3> {
|
||||
__nv_fp8x2_e4m3 elts[8];
|
||||
};
|
||||
|
||||
template <typename Int>
|
||||
__host__ __device__ inline Int round_up(Int x, Int y) {
|
||||
static_assert(std::is_integral_v<Int>,
|
||||
"round_up argument must be integral type");
|
||||
return ((x + y - 1) / y) * y;
|
||||
}
|
||||
|
||||
template <typename Int>
|
||||
__host__ __device__ __forceinline__ Int div_round_up(Int x, Int y) {
|
||||
return (x + y - 1) / y;
|
||||
}
|
||||
|
||||
// Compute effective rows for grid configuration with swizzled SF layouts.
|
||||
inline int computeEffectiveRows(int m) {
|
||||
constexpr int ROW_TILE = 128;
|
||||
return round_up(m, ROW_TILE);
|
||||
}
|
||||
|
||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
||||
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
|
||||
inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
|
||||
uint32_t val;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
@@ -136,7 +101,7 @@ inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
|
||||
}
|
||||
|
||||
// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
|
||||
__device__ __forceinline__ uint32_t fp32_vec8_to_e2m1(float2 (&array)[4]) {
|
||||
inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
|
||||
uint32_t val;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
@@ -149,115 +114,20 @@ __device__ __forceinline__ uint32_t fp32_vec8_to_e2m1(float2 (&array)[4]) {
|
||||
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
|
||||
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
|
||||
"}\n"
|
||||
"}"
|
||||
: "=r"(val)
|
||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
|
||||
return val;
|
||||
}
|
||||
|
||||
struct u32x2 {
|
||||
uint32_t lo, hi;
|
||||
};
|
||||
|
||||
using fp4_packed_t = std::conditional_t<CVT_FP4_PACK16, u32x2, uint32_t>;
|
||||
|
||||
__device__ __forceinline__ u32x2 fp32_vec16_to_e2m1(float2 (&array)[8]) {
|
||||
u32x2 out;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b8 b0;\n"
|
||||
".reg .b8 b1;\n"
|
||||
".reg .b8 b2;\n"
|
||||
".reg .b8 b3;\n"
|
||||
".reg .b8 b4;\n"
|
||||
".reg .b8 b5;\n"
|
||||
".reg .b8 b6;\n"
|
||||
".reg .b8 b7;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b0, %3, %2;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b1, %5, %4;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b2, %7, %6;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b3, %9, %8;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b4, %11, %10;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b5, %13, %12;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b6, %15, %14;\n"
|
||||
"cvt.rn.satfinite.e2m1x2.f32 b7, %17, %16;\n"
|
||||
"mov.b32 %0, {b0, b1, b2, b3};\n"
|
||||
"mov.b32 %1, {b4, b5, b6, b7};\n"
|
||||
"}\n"
|
||||
: "=r"(out.lo), "=r"(out.hi)
|
||||
: "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
|
||||
"f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y),
|
||||
"f"(array[4].x), "f"(array[4].y), "f"(array[5].x), "f"(array[5].y),
|
||||
"f"(array[6].x), "f"(array[6].y), "f"(array[7].x), "f"(array[7].y));
|
||||
return out;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t pack_fp4(float2 (&v)[4]) {
|
||||
return fp32_vec8_to_e2m1(v);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ u32x2 pack_fp4(float2 (&v)[8]) {
|
||||
return fp32_vec16_to_e2m1(v);
|
||||
}
|
||||
|
||||
// Fast reciprocal.
|
||||
__device__ __forceinline__ float reciprocal_approximate_ftz(float a) {
|
||||
inline __device__ float reciprocal_approximate_ftz(float a) {
|
||||
float b;
|
||||
asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a));
|
||||
asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
|
||||
return b;
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__device__ __forceinline__ void ld128_or_zero_cg_u32(PackedVec<Type>& out,
|
||||
const void* ptr,
|
||||
bool pred) {
|
||||
uint32_t r0, r1, r2, r3;
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred pr;\n"
|
||||
" setp.ne.u32 pr, %4, 0;\n"
|
||||
" mov.u32 %0, 0;\n"
|
||||
" mov.u32 %1, 0;\n"
|
||||
" mov.u32 %2, 0;\n"
|
||||
" mov.u32 %3, 0;\n"
|
||||
" @pr ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%5];\n"
|
||||
"}\n"
|
||||
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
|
||||
: "r"((int)pred), "l"(ptr));
|
||||
|
||||
*reinterpret_cast<uint4*>(&out) = uint4{r0, r1, r2, r3};
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__device__ __forceinline__ void ld256_or_zero_cg_u32(PackedVec<Type>& out,
|
||||
const void* ptr,
|
||||
bool pred) {
|
||||
uint32_t r0, r1, r2, r3, r4, r5, r6, r7;
|
||||
|
||||
asm volatile(
|
||||
"{\n"
|
||||
" .reg .pred pr;\n"
|
||||
" setp.ne.u32 pr, %8, 0;\n"
|
||||
" mov.u32 %0, 0;\n"
|
||||
" mov.u32 %1, 0;\n"
|
||||
" mov.u32 %2, 0;\n"
|
||||
" mov.u32 %3, 0;\n"
|
||||
" mov.u32 %4, 0;\n"
|
||||
" mov.u32 %5, 0;\n"
|
||||
" mov.u32 %6, 0;\n"
|
||||
" mov.u32 %7, 0;\n"
|
||||
" @pr ld.global.cg.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%9];\n"
|
||||
"}\n"
|
||||
: "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3), "=r"(r4), "=r"(r5), "=r"(r6),
|
||||
"=r"(r7)
|
||||
: "r"((int)pred), "l"(ptr));
|
||||
|
||||
reinterpret_cast<uint4*>(&out)[0] = uint4{r0, r1, r2, r3};
|
||||
reinterpret_cast<uint4*>(&out)[1] = uint4{r4, r5, r6, r7};
|
||||
}
|
||||
|
||||
// Compute SF output offset for swizzled tensor core layout.
|
||||
// SF layout: [numMTiles, numKTiles, 32, 4, 4]
|
||||
// Caller must precompute: numKTiles = (numCols + 63) / 64
|
||||
@@ -296,41 +166,21 @@ __device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(
|
||||
return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
|
||||
}
|
||||
|
||||
template <class SFType>
|
||||
__device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack,
|
||||
int packs_per_row_sf,
|
||||
SFType* SFout) {
|
||||
constexpr int PACK = CVT_FP4_ELTS_PER_THREAD;
|
||||
constexpr int THREADS_PER_SF =
|
||||
CVT_FP4_SF_VEC_SIZE / PACK; // 1 if PACK=16, 2 else PACK=8
|
||||
|
||||
if (threadIdx.x % THREADS_PER_SF != 0) return nullptr;
|
||||
|
||||
int sf_col =
|
||||
pack / THREADS_PER_SF; // PACK=16 => sf_col=pack; PACK=8 => sf_col=pack/2
|
||||
int64_t off = (int64_t)row * packs_per_row_sf + sf_col;
|
||||
|
||||
return (uint8_t*)SFout + off;
|
||||
}
|
||||
|
||||
// Quantizes the provided PackedVec into the uint32_t output
|
||||
template <class Type, int CVT_FP4_NUM_THREADS_PER_SF, bool UE8M0_SF = false>
|
||||
__device__ __forceinline__ fp4_packed_t
|
||||
cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
|
||||
template <class Type, bool UE8M0_SF = false>
|
||||
__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
|
||||
uint8_t* SFout) {
|
||||
// Get absolute maximum values among the local 8 values.
|
||||
auto localMax = __habs2(vec.elts[0]);
|
||||
|
||||
// Local maximum value.
|
||||
// Local maximum value.
|
||||
#pragma unroll
|
||||
for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
|
||||
localMax = __hmax2(localMax, __habs2(vec.elts[i]));
|
||||
}
|
||||
|
||||
// Get the absolute maximum among all 16 values (two threads).
|
||||
|
||||
if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) {
|
||||
localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax);
|
||||
}
|
||||
localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
|
||||
// Get the final absolute maximum values.
|
||||
float vecMax = float(__hmax(localMax.x, localMax.y));
|
||||
|
||||
@@ -355,17 +205,18 @@ cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
|
||||
// Convert back to fp32.
|
||||
SFValue = float(tmp);
|
||||
}
|
||||
|
||||
// Write the SF to global memory (STG.8).
|
||||
if (SFout) *SFout = fp8SFVal;
|
||||
|
||||
// Get the output scale.
|
||||
// Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
|
||||
// reciprocal(SFScaleVal))
|
||||
float outputScale =
|
||||
SFValue != 0.0f ? reciprocal_approximate_ftz(
|
||||
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
||||
: 0.0f;
|
||||
SFValue != 0 ? reciprocal_approximate_ftz(
|
||||
SFValue * reciprocal_approximate_ftz(SFScaleVal))
|
||||
: 0.0f;
|
||||
|
||||
if (SFout) {
|
||||
// Write the SF to global memory (STG.8).
|
||||
*SFout = fp8SFVal;
|
||||
}
|
||||
|
||||
// Convert the input to float.
|
||||
float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
|
||||
@@ -382,7 +233,10 @@ cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal, uint8_t* SFout) {
|
||||
}
|
||||
|
||||
// Convert to e2m1 values.
|
||||
return pack_fp4(fp2Vals);
|
||||
uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
|
||||
|
||||
// Write the e2m1 values to global memory.
|
||||
return e2m1Vec;
|
||||
}
|
||||
|
||||
// silu in float32
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <iostream>
|
||||
#include "../marlin/marlin_dtypes.cuh"
|
||||
#include "../gptq_marlin/marlin_dtypes.cuh"
|
||||
using marlin::MarlinScalarType2;
|
||||
|
||||
namespace allspark {
|
||||
|
||||
@@ -46,7 +46,7 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
|
||||
|
||||
} // namespace marlin
|
||||
|
||||
torch::Tensor marlin_gemm(
|
||||
torch::Tensor gptq_marlin_gemm(
|
||||
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
|
||||
torch::Tensor& b_q_weight,
|
||||
std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
|
||||
@@ -528,7 +528,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
|
||||
|
||||
} // namespace marlin
|
||||
|
||||
torch::Tensor marlin_gemm(
|
||||
torch::Tensor gptq_marlin_gemm(
|
||||
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
|
||||
torch::Tensor& b_q_weight,
|
||||
std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
|
||||
@@ -856,5 +856,5 @@ torch::Tensor marlin_gemm(
|
||||
#endif
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("marlin_gemm", &marlin_gemm);
|
||||
m.impl("gptq_marlin_gemm", &gptq_marlin_gemm);
|
||||
}
|
||||
@@ -303,9 +303,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
ops.impl("permute_cols", torch::kCUDA, &permute_cols);
|
||||
|
||||
// Marlin Optimized Quantized GEMM (supports GPTQ, AWQ, FP8, NVFP4, MXFP4).
|
||||
// gptq_marlin Optimized Quantized GEMM for GPTQ.
|
||||
ops.def(
|
||||
"marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
|
||||
"gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
|
||||
"Tensor? b_bias_or_none,Tensor b_scales, "
|
||||
"Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, "
|
||||
"Tensor? "
|
||||
@@ -546,8 +546,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// Compute NVFP4 block quantized tensor.
|
||||
ops.def(
|
||||
"scaled_fp4_quant(Tensor! output, Tensor input,"
|
||||
" Tensor! output_scale, Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout) -> ()");
|
||||
" Tensor! output_scale, Tensor input_scale) -> ()");
|
||||
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
|
||||
|
||||
// Compute NVFP4 experts quantization.
|
||||
|
||||
+11
-50
@@ -15,11 +15,9 @@
|
||||
# Build arguments:
|
||||
# PYTHON_VERSION=3.13|3.12 (default)|3.11|3.10
|
||||
# VLLM_CPU_DISABLE_AVX512=false (default)|true
|
||||
# VLLM_CPU_AVX2=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512BF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512VNNI=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AMXBF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512BF16=false (default)|true
|
||||
# VLLM_CPU_AVX512VNNI=false (default)|true
|
||||
# VLLM_CPU_AMXBF16=false |true (default)
|
||||
#
|
||||
|
||||
######################### COMMON BASE IMAGE #########################
|
||||
@@ -56,12 +54,9 @@ ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
# Copy requirements files for installation
|
||||
COPY requirements/common.txt requirements/common.txt
|
||||
COPY requirements/cpu.txt requirements/cpu.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/common.txt,target=requirements/common.txt \
|
||||
--mount=type=bind,src=requirements/cpu.txt,target=requirements/cpu.txt \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/cpu.txt
|
||||
|
||||
@@ -93,12 +88,6 @@ ARG GIT_REPO_CHECK=0
|
||||
# Support for building with non-AVX512 vLLM: docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" ...
|
||||
ARG VLLM_CPU_DISABLE_AVX512=0
|
||||
ENV VLLM_CPU_DISABLE_AVX512=${VLLM_CPU_DISABLE_AVX512}
|
||||
# Support for cross-compilation with AVX2 ISA: docker build --build-arg VLLM_CPU_AVX2="1" ...
|
||||
ARG VLLM_CPU_AVX2=0
|
||||
ENV VLLM_CPU_AVX2=${VLLM_CPU_AVX2}
|
||||
# Support for cross-compilation with AVX512 ISA: docker build --build-arg VLLM_CPU_AVX512="1" ...
|
||||
ARG VLLM_CPU_AVX512=0
|
||||
ENV VLLM_CPU_AVX512=${VLLM_CPU_AVX512}
|
||||
# Support for building with AVX512BF16 ISA: docker build --build-arg VLLM_CPU_AVX512BF16="true" ...
|
||||
ARG VLLM_CPU_AVX512BF16=0
|
||||
ENV VLLM_CPU_AVX512BF16=${VLLM_CPU_AVX512BF16}
|
||||
@@ -111,19 +100,18 @@ ENV VLLM_CPU_AMXBF16=${VLLM_CPU_AMXBF16}
|
||||
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
# Copy build requirements
|
||||
COPY requirements/cpu-build.txt requirements/build.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,src=requirements/cpu-build.txt,target=requirements/build.txt \
|
||||
uv pip install -r requirements/build.txt
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
|
||||
RUN --mount=type=bind,source=.git,target=.git \
|
||||
if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=cache,target=/root/.cache/ccache \
|
||||
--mount=type=cache,target=/workspace/vllm/.deps,sharing=locked \
|
||||
--mount=type=bind,source=.git,target=.git \
|
||||
VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
|
||||
|
||||
######################### TEST DEPS #########################
|
||||
@@ -131,11 +119,9 @@ FROM base AS vllm-test-deps
|
||||
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
# Copy test requirements
|
||||
COPY requirements/test.in requirements/cpu-test.in
|
||||
|
||||
# TODO: Update to 2.9.0 when there is a new build for intel_extension_for_pytorch for that version
|
||||
RUN \
|
||||
RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
|
||||
cp requirements/test.in requirements/cpu-test.in && \
|
||||
sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
|
||||
remove_packages_not_supported_on_aarch64() { \
|
||||
case "$(uname -m)" in \
|
||||
@@ -214,29 +200,4 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=vllm-build,src=/workspace/vllm/dist,target=dist \
|
||||
uv pip install dist/*.whl
|
||||
|
||||
# Add labels to document build configuration
|
||||
LABEL org.opencontainers.image.title="vLLM CPU"
|
||||
LABEL org.opencontainers.image.description="vLLM inference engine for CPU platforms"
|
||||
LABEL org.opencontainers.image.vendor="vLLM Project"
|
||||
LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm"
|
||||
|
||||
# Build configuration labels
|
||||
ARG TARGETARCH
|
||||
ARG VLLM_CPU_DISABLE_AVX512
|
||||
ARG VLLM_CPU_AVX2
|
||||
ARG VLLM_CPU_AVX512
|
||||
ARG VLLM_CPU_AVX512BF16
|
||||
ARG VLLM_CPU_AVX512VNNI
|
||||
ARG VLLM_CPU_AMXBF16
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
LABEL ai.vllm.build.target-arch="${TARGETARCH}"
|
||||
LABEL ai.vllm.build.cpu-disable-avx512="${VLLM_CPU_DISABLE_AVX512:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx2="${VLLM_CPU_AVX2:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512="${VLLM_CPU_AVX512:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512bf16="${VLLM_CPU_AVX512BF16:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512vnni="${VLLM_CPU_AVX512VNNI:-false}"
|
||||
LABEL ai.vllm.build.cpu-amxbf16="${VLLM_CPU_AMXBF16:-false}"
|
||||
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
+1
-14
@@ -227,7 +227,7 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
|
||||
# This ensures setuptools_scm sees clean repo state for version detection
|
||||
RUN --mount=type=bind,source=.git,target=vllm/.git \
|
||||
cd vllm \
|
||||
&& pip install setuptools_scm regex \
|
||||
&& pip install setuptools_scm \
|
||||
&& VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
|
||||
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
|
||||
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
|
||||
@@ -342,19 +342,6 @@ RUN mkdir src && mv vllm src/vllm
|
||||
FROM base AS final
|
||||
|
||||
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Clean up sccache from release image (not needed at runtime)
|
||||
# This removes the binary and wrappers that may have been installed during build
|
||||
RUN rm -f /usr/bin/sccache || true \
|
||||
&& rm -rf /opt/sccache-wrappers || true
|
||||
|
||||
# Unset sccache environment variables for the release image
|
||||
# This prevents S3 bucket config from leaking into production images
|
||||
ENV SCCACHE_BUCKET=
|
||||
ENV SCCACHE_REGION=
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=
|
||||
ENV SCCACHE_IDLE_TIMEOUT=
|
||||
|
||||
# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
|
||||
# Manually remove it so that later steps of numpy upgrade can continue
|
||||
RUN case "$(which python3)" in \
|
||||
|
||||
@@ -13,14 +13,14 @@ For x86 CPU environment, please use the image with "-cpu" postfix. For AArch64 C
|
||||
Here is an example for docker run command for CPU. For GPUs skip setting the `ON_CPU` env var.
|
||||
|
||||
```bash
|
||||
export VLLM_COMMIT=7f42dc20bb2800d09faa72b26f25d54e26f1b694 # use full commit hash from the main branch
|
||||
export VLLM_COMMIT=1da94e673c257373280026f75ceb4effac80e892 # use full commit hash from the main branch
|
||||
export HF_TOKEN=<valid Hugging Face token>
|
||||
if [[ "$(uname -m)" == aarch64 || "$(uname -m)" == arm64 ]]; then
|
||||
IMG_SUFFIX="arm64-cpu"
|
||||
else
|
||||
IMG_SUFFIX="cpu"
|
||||
fi
|
||||
docker run -it --entrypoint /bin/bash -v /data/huggingface:/root/.cache/huggingface -e HF_TOKEN=$HF_TOKEN -e ON_CPU=1 --shm-size=16g --name vllm-cpu-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${VLLM_COMMIT}-${IMG_SUFFIX}
|
||||
docker run -it --entrypoint /bin/bash -v /data/huggingface:/root/.cache/huggingface -e HF_TOKEN=$HF_TOKEN -e ON_ARM64_CPU=1 --shm-size=16g --name vllm-cpu-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${VLLM_COMMIT}-${IMG_SUFFIX}
|
||||
```
|
||||
|
||||
Then, run below command inside the docker instance.
|
||||
|
||||
@@ -106,7 +106,6 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
|
||||
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`
|
||||
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
|
||||
Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
|
||||
|
||||
+117
-139
@@ -20,6 +20,67 @@ To input multi-modal data, follow this schema in [vllm.inputs.PromptType][]:
|
||||
- `prompt`: The prompt should follow the format that is documented on HuggingFace.
|
||||
- `multi_modal_data`: This is a dictionary that follows the schema defined in [vllm.multimodal.inputs.MultiModalDataDict][].
|
||||
|
||||
### Stable UUIDs for Caching (multi_modal_uuids)
|
||||
|
||||
When using multi-modal inputs, vLLM normally hashes each media item by content to enable caching across requests. You can optionally pass `multi_modal_uuids` to provide your own stable IDs for each item so caching can reuse work across requests without rehashing the raw content.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_a = Image.open("/path/to/a.jpg")
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [img_a, img_b]},
|
||||
# Provide stable IDs for caching.
|
||||
# Requirements (matched by this example):
|
||||
# - Include every modality present in multi_modal_data.
|
||||
# - For lists, provide the same number of entries.
|
||||
# - Use None to fall back to content hashing for that item.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
Using UUIDs, you can also skip sending media data entirely if you expect cache hits for respective items. Note that the request will fail if the skipped media doesn't have a corresponding UUID, or if the UUID fails to hit the cache.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [None, img_b]},
|
||||
# Since img_a is expected to be cached, we can skip sending the actual
|
||||
# image entirely.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
If both multimodal processor caching and prefix caching are disabled, user-provided `multi_modal_uuids` are ignored.
|
||||
|
||||
### Image Inputs
|
||||
|
||||
You can pass a single image to the `'image'` field of the multi-modal dictionary, as shown in the following examples:
|
||||
@@ -336,8 +397,7 @@ No manual conversion is needed - vLLM handles the channel normalization automati
|
||||
### Embedding Inputs
|
||||
|
||||
To input pre-computed embeddings belonging to a data type (i.e. image, video, or audio) directly to the language model,
|
||||
pass a tensor of shape `(..., hidden_size of LM)` to the corresponding field of the multi-modal dictionary.
|
||||
The exact shape depends on the model being used.
|
||||
pass a tensor of shape `(num_items, feature_size, hidden_size of LM)` to the corresponding field of the multi-modal dictionary.
|
||||
|
||||
You must enable this feature via `enable_mm_embeds=True`.
|
||||
|
||||
@@ -358,7 +418,8 @@ You must enable this feature via `enable_mm_embeds=True`.
|
||||
# Refer to the HuggingFace repo for the correct format to use
|
||||
prompt = "USER: <image>\nWhat is the content of this image?\nASSISTANT:"
|
||||
|
||||
# For most models, `image_embeds` has shape: (num_images, image_feature_size, hidden_size)
|
||||
# Embeddings for single image
|
||||
# torch.Tensor of shape (1, image_feature_size, hidden_size of LM)
|
||||
image_embeds = torch.load(...)
|
||||
|
||||
outputs = llm.generate({
|
||||
@@ -369,8 +430,21 @@ You must enable this feature via `enable_mm_embeds=True`.
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
# Additional examples for models that require extra fields
|
||||
For Qwen2-VL and MiniCPM-V, we accept additional parameters alongside the embeddings:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
# Construct the prompt based on your model
|
||||
prompt = ...
|
||||
|
||||
# Embeddings for multiple images
|
||||
# torch.Tensor of shape (num_images, image_feature_size, hidden_size of LM)
|
||||
image_embeds = torch.load(...)
|
||||
|
||||
# Qwen2-VL
|
||||
llm = LLM(
|
||||
"Qwen/Qwen2-VL-2B-Instruct",
|
||||
limit_mm_per_prompt={"image": 4},
|
||||
@@ -378,15 +452,13 @@ You must enable this feature via `enable_mm_embeds=True`.
|
||||
)
|
||||
mm_data = {
|
||||
"image": {
|
||||
# Shape: (total_feature_size, hidden_size)
|
||||
# total_feature_size = sum(image_feature_size for image in images)
|
||||
"image_embeds": torch.load(...),
|
||||
# Shape: (num_images, 3)
|
||||
"image_embeds": image_embeds,
|
||||
# image_grid_thw is needed to calculate positional encoding.
|
||||
"image_grid_thw": torch.load(...),
|
||||
"image_grid_thw": torch.load(...), # torch.Tensor of shape (1, 3),
|
||||
}
|
||||
}
|
||||
|
||||
# MiniCPM-V
|
||||
llm = LLM(
|
||||
"openbmb/MiniCPM-V-2_6",
|
||||
trust_remote_code=True,
|
||||
@@ -395,14 +467,20 @@ You must enable this feature via `enable_mm_embeds=True`.
|
||||
)
|
||||
mm_data = {
|
||||
"image": {
|
||||
# Shape: (num_images, num_slices, hidden_size)
|
||||
# num_slices can differ for each image
|
||||
"image_embeds": [torch.load(...) for image in images],
|
||||
# Shape: (num_images, 2)
|
||||
"image_embeds": image_embeds,
|
||||
# image_sizes is needed to calculate details of the sliced image.
|
||||
"image_sizes": [image.size for image in images],
|
||||
"image_sizes": [image.size for image in images], # list of image sizes
|
||||
}
|
||||
}
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": mm_data,
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
generated_text = o.outputs[0].text
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
For Qwen3-VL, the `image_embeds` should contain both the base image embedding and deepstack features.
|
||||
@@ -423,8 +501,8 @@ You can pass pre-computed audio embeddings similar to image embeddings:
|
||||
# Refer to the HuggingFace repo for the correct format to use
|
||||
prompt = "USER: <audio>\nWhat is in this audio?\nASSISTANT:"
|
||||
|
||||
# Load pre-computed audio embeddings, usually with shape:
|
||||
# (num_audios, audio_feature_size, hidden_size of LM)
|
||||
# Load pre-computed audio embeddings
|
||||
# torch.Tensor of shape (1, audio_feature_size, hidden_size of LM)
|
||||
audio_embeds = torch.load(...)
|
||||
|
||||
outputs = llm.generate({
|
||||
@@ -437,67 +515,6 @@ You can pass pre-computed audio embeddings similar to image embeddings:
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
### Cached Inputs
|
||||
|
||||
When using multi-modal inputs, vLLM normally hashes each media item by content to enable caching across requests. You can optionally pass `multi_modal_uuids` to provide your own stable IDs for each item so caching can reuse work across requests without rehashing the raw content.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_a = Image.open("/path/to/a.jpg")
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [img_a, img_b]},
|
||||
# Provide stable IDs for caching.
|
||||
# Requirements (matched by this example):
|
||||
# - Include every modality present in multi_modal_data.
|
||||
# - For lists, provide the same number of entries.
|
||||
# - Use None to fall back to content hashing for that item.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
Using UUIDs, you can also skip sending media data entirely if you expect cache hits for respective items. Note that the request will fail if the skipped media doesn't have a corresponding UUID, or if the UUID fails to hit the cache.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from PIL import Image
|
||||
|
||||
# Qwen2.5-VL example with two images
|
||||
llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct")
|
||||
|
||||
prompt = "USER: <image><image>\nDescribe the differences.\nASSISTANT:"
|
||||
img_b = Image.open("/path/to/b.jpg")
|
||||
|
||||
outputs = llm.generate({
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": [None, img_b]},
|
||||
# Since img_a is expected to be cached, we can skip sending the actual
|
||||
# image entirely.
|
||||
"multi_modal_uuids": {"image": ["sku-1234-a", None]},
|
||||
})
|
||||
|
||||
for o in outputs:
|
||||
print(o.outputs[0].text)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
If both multimodal processor caching and prefix caching are disabled, user-provided `multi_modal_uuids` are ignored.
|
||||
|
||||
## Online Serving
|
||||
|
||||
Our OpenAI-compatible server accepts multi-modal data via the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat). Media inputs also support optional UUIDs users can provide to uniquely identify each media, which is used to cache the media results across requests.
|
||||
@@ -862,11 +879,7 @@ Full example: [examples/online_serving/openai_chat_completion_client_for_multimo
|
||||
### Embedding Inputs
|
||||
|
||||
To input pre-computed embeddings belonging to a data type (i.e. image, video, or audio) directly to the language model,
|
||||
pass a tensor of shape `(..., hidden_size of LM)` for each item to the corresponding field of the multi-modal dictionary.
|
||||
|
||||
!!! important
|
||||
Unlike offline inference, the embeddings for each item must be passed separately
|
||||
in order for placeholder tokens to be applied correctly by the chat template.
|
||||
pass a tensor of shape `(num_items, feature_size, hidden_size of LM)` to the corresponding field of the multi-modal dictionary.
|
||||
|
||||
You must enable this feature via the `--enable-mm-embeds` flag in `vllm serve`.
|
||||
|
||||
@@ -884,6 +897,11 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
```python
|
||||
from vllm.utils.serial_utils import tensor2base64
|
||||
|
||||
image_embedding = torch.load(...)
|
||||
grid_thw = torch.load(...) # Required by Qwen/Qwen2-VL-2B-Instruct
|
||||
|
||||
base64_image_embedding = tensor2base64(image_embedding)
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
@@ -894,33 +912,29 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
model = "llava-hf/llava-1.5-7b-hf"
|
||||
embeds = {
|
||||
"type": "image_embeds",
|
||||
"image_embeds": tensor2base64(torch.load(...)), # Shape: (image_feature_size, hidden_size)
|
||||
"image_embeds": f"{base64_image_embedding}",
|
||||
"uuid": image_url, # Optional
|
||||
}
|
||||
|
||||
|
||||
# Additional examples for models that require extra fields
|
||||
# Pass additional parameters (available to Qwen2-VL and MiniCPM-V)
|
||||
model = "Qwen/Qwen2-VL-2B-Instruct"
|
||||
embeds = {
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.load(...)), # Shape: (image_feature_size, hidden_size)
|
||||
"image_grid_thw": tensor2base64(torch.load(...)), # Shape: (3,)
|
||||
"image_embeds": f"{base64_image_embedding}", # Required
|
||||
"image_grid_thw": f"{base64_image_grid_thw}", # Required by Qwen/Qwen2-VL-2B-Instruct
|
||||
},
|
||||
"uuid": image_url, # Optional
|
||||
}
|
||||
|
||||
model = "openbmb/MiniCPM-V-2_6"
|
||||
embeds = {
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.load(...)), # Shape: (num_slices, hidden_size)
|
||||
"image_sizes": tensor2base64(torch.load(...)), # Shape: (2,)
|
||||
"image_embeds": f"{base64_image_embedding}", # Required
|
||||
"image_sizes": f"{base64_image_sizes}", # Required by openbmb/MiniCPM-V-2_6
|
||||
},
|
||||
"uuid": image_url, # Optional
|
||||
}
|
||||
|
||||
# Single image input
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
@@ -940,55 +954,9 @@ The following example demonstrates how to pass image embeddings to the OpenAI se
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Multi image input
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this image?",
|
||||
},
|
||||
embeds,
|
||||
embeds,
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Multi image input (interleaved)
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
embeds,
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this image?",
|
||||
},
|
||||
embeds,
|
||||
],
|
||||
},
|
||||
],
|
||||
model=model,
|
||||
)
|
||||
```
|
||||
|
||||
### Cached Inputs
|
||||
|
||||
Just like with offline inference, you can skip sending media if you expect cache hits with provided UUIDs. You can do so by sending media like this:
|
||||
For Online Serving, you can also skip sending media if you expect cache hits with provided UUIDs. You can do so by sending media like this:
|
||||
|
||||
??? code
|
||||
|
||||
@@ -1022,3 +990,13 @@ Just like with offline inference, you can skip sending media if you expect cache
|
||||
},
|
||||
|
||||
```
|
||||
|
||||
!!! note
|
||||
Multiple messages can now contain `{"type": "image_embeds"}`, enabling you to pass multiple image embeddings in a single request (similar to regular images). The number of embeddings is limited by `--limit-mm-per-prompt`.
|
||||
|
||||
**Important**: The embedding shape format differs based on the number of embeddings:
|
||||
|
||||
- **Single embedding**: 3D tensor of shape `(1, feature_size, hidden_size)`
|
||||
- **Multiple embeddings**: List of 2D tensors, each of shape `(feature_size, hidden_size)`
|
||||
|
||||
If used with a model that requires additional parameters, you must also provide a tensor for each of them, e.g. `image_grid_thw`, `image_sizes`, etc.
|
||||
|
||||
@@ -184,6 +184,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
|
||||
--kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
|
||||
```
|
||||
|
||||
### Cross layers blocks
|
||||
|
||||
By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
|
||||
To enable this feature:
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
|
||||
```
|
||||
|
||||
## Example Scripts/Code
|
||||
|
||||
Refer to these example scripts in the vLLM repository:
|
||||
|
||||
@@ -131,7 +131,7 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
|
||||
|
||||
=== "Apple silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/cpu.apple.inc.md:build-image-from-source"
|
||||
--8<-- "docs/getting_started/installation/cpu.arm.inc.md:build-image-from-source"
|
||||
|
||||
=== "IBM Z (S390X)"
|
||||
--8<-- "docs/getting_started/installation/cpu.s390x.inc.md:build-image-from-source"
|
||||
|
||||
@@ -164,76 +164,21 @@ uv pip install dist/*.whl
|
||||
[https://gallery.ecr.aws/q9t5s3a7/vllm-cpu-release-repo](https://gallery.ecr.aws/q9t5s3a7/vllm-cpu-release-repo)
|
||||
|
||||
!!! warning
|
||||
If deploying the pre-built images on machines without `avx512f`, `avx512_bf16`, or `avx512_vnni` support, an `Illegal instruction` error may be raised. See the build-image-from-source section below for build arguments to match your target CPU capabilities.
|
||||
If deploying the pre-built images on machines without `avx512f`, `avx512_bf16`, or `avx512_vnni` support, an `Illegal instruction` error may be raised. It is recommended to build images for these machines with the appropriate build arguments (e.g., `--build-arg VLLM_CPU_DISABLE_AVX512=true`, `--build-arg VLLM_CPU_AVX512BF16=false`, or `--build-arg VLLM_CPU_AVX512VNNI=false`) to disable unsupported features. Please note that without `avx512f`, AVX2 will be used and this version is not recommended because it only has basic feature support.
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
# --8<-- [start:build-image-from-source]
|
||||
|
||||
## Building for your target CPU
|
||||
|
||||
vLLM supports building Docker images for x86 CPU platforms with automatic instruction set detection.
|
||||
|
||||
### Basic build command
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_DISABLE_AVX512=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX2=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX512=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX512BF16=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=<false (default)|true> \
|
||||
--build-arg VLLM_CPU_AMXBF16=<false|true (default)> \
|
||||
--build-arg VLLM_CPU_AVX512BF16=false (default)|true \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=false (default)|true \
|
||||
--build-arg VLLM_CPU_AMXBF16=false|true (default) \
|
||||
--build-arg VLLM_CPU_DISABLE_AVX512=false (default)|true \
|
||||
--tag vllm-cpu-env \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
!!! note "Instruction set auto-detection"
|
||||
By default, vLLM will auto-detect CPU instruction sets (AVX512, AVX2, etc.) from the build system's CPU flags. Build arguments like `VLLM_CPU_AVX2`, `VLLM_CPU_AVX512`, `VLLM_CPU_AVX512BF16`, `VLLM_CPU_AVX512VNNI`, and `VLLM_CPU_AMXBF16` are primarily used for **cross-compilation** or for building container images on systems that don't have the target platforms ISA:
|
||||
|
||||
- Set `VLLM_CPU_{ISA}=true` to force-enable an instruction set (for cross-compilation to target platforms with that ISA)
|
||||
- Set `VLLM_CPU_{ISA}=false` to rely on auto-detection
|
||||
- When an ISA build arg is set to `true`, vLLM will build with that instruction set regardless of the build system's CPU capabilities
|
||||
|
||||
### Build examples
|
||||
|
||||
**Example 1: Auto-detection (native build)**
|
||||
|
||||
Build on a machine with the same CPU as your target deployment:
|
||||
|
||||
```bash
|
||||
# Auto-detects all CPU features from the build system
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--tag vllm-cpu-env \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Example 2: Cross-compilation for AVX512 deployment**
|
||||
|
||||
Build an AVX512 image on any x86_64 system (even without AVX512):
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_AVX512=true \
|
||||
--build-arg VLLM_CPU_AVX512BF16=true \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=true \
|
||||
--tag vllm-cpu-avx512 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Example 3: Cross-compilation for AVX2 deployment**
|
||||
|
||||
Build an AVX2 image for older CPUs:
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_AVX2=true \
|
||||
--tag vllm-cpu-avx2 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
## Launching the OpenAI server
|
||||
|
||||
```bash
|
||||
# Launching OpenAI server
|
||||
docker run --rm \
|
||||
--security-opt seccomp=unconfined \
|
||||
--cap-add SYS_NICE \
|
||||
|
||||
@@ -118,7 +118,7 @@ There are more environment variables to control the behavior of Python-only buil
|
||||
|
||||
* `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file path of a pre-compiled wheel to use. All other logic to find the wheel will be skipped.
|
||||
* `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit hash to download the pre-compiled wheel. It can be `nightly` to use the last **already built** commit on the main branch.
|
||||
* `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cu130`, `cpu`. If not specified, the variant is auto-detected based on your system's CUDA version (from PyTorch or nvidia-smi). You can also set `VLLM_MAIN_CUDA_VERSION` to override auto-detection.
|
||||
* `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cpu`. If not specified, the CUDA variant with `VLLM_MAIN_CUDA_VERSION` will be tried, then fallback to the default variant on the remote index.
|
||||
|
||||
You can find more information about vLLM's wheels in [Install the latest code](#install-the-latest-code).
|
||||
|
||||
|
||||
@@ -31,7 +31,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/
|
||||
To install a specific version and ROCm variant of vLLM wheel.
|
||||
|
||||
```bash
|
||||
uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.1/rocm700
|
||||
uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
```
|
||||
|
||||
!!! warning "Caveats for using `pip`"
|
||||
@@ -41,7 +41,7 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.14.1/rocm700
|
||||
If you insist on using `pip`, you have to specify the exact vLLM version and full URL of the wheel path `https://wheels.vllm.ai/rocm/<version>/<rocm-variant>` (which can be obtained from the web page).
|
||||
|
||||
```bash
|
||||
pip install vllm==0.14.1+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.1/rocm700
|
||||
pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-wheels]
|
||||
|
||||
+26
-47
@@ -25,55 +25,34 @@ Maintainers form a hierarchy based on sustained, high-quality contributions and
|
||||
|
||||
### Core Maintainers
|
||||
|
||||
Core Maintainers function like a project planning and decision making committee. In other convention, they might be called a Technical Steering Committee (TSC). In vLLM vocabulary, they are often known as "Project Leads". They meet weekly to coordinate roadmap priorities and allocate engineering resources.
|
||||
Core Maintainers function like a project planning and decision making committee. In other convention, they might be called a Technical Steering Committee (TSC). In vLLM vocabulary, they are often known as "Project Leads". They meet weekly to coordinate roadmap priorities and allocate engineering resources. Current active leads: @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, @robertgshaw2-redhat, @tlrmchlsmth, @mgoin, @njhill, @ywang96, @houseroad, @yeqcharlotte, @ApostaC
|
||||
|
||||
**Project Leads:**
|
||||
The responsibilities of the core maintainers are:
|
||||
|
||||
- Woosuk Kwon ([@WoosukKwon](https://github.com/WoosukKwon))
|
||||
- Zhuohan Li ([@zhuohan123](https://github.com/zhuohan123))
|
||||
- Simon Mo ([@simon-mo](https://github.com/simon-mo))
|
||||
- Kaichao You ([@youkaichao](https://github.com/youkaichao))
|
||||
- Robert Shaw ([@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat))
|
||||
- Tyler Michael Smith ([@tlrmchlsmth](https://github.com/tlrmchlsmth))
|
||||
- Michael Goin ([@mgoin](https://github.com/mgoin))
|
||||
- Nick Hill ([@njhill](https://github.com/njhill))
|
||||
- Roger Wang ([@ywang96](https://github.com/ywang96))
|
||||
- Lu Fang ([@houseroad](https://github.com/houseroad))
|
||||
- Ye (Charlotte) Qi ([@yeqcharlotte](https://github.com/yeqcharlotte))
|
||||
- Yihua Cheng ([@ApostaC](https://github.com/ApostaC))
|
||||
|
||||
**Responsibilities:**
|
||||
|
||||
- Author quarterly roadmap and responsible for each development effort.
|
||||
- Making major changes to the technical direction or scope of vLLM and vLLM projects.
|
||||
- Defining the project's release strategy.
|
||||
- Work with model providers, hardware vendors, and key users of vLLM to ensure the project is on the right track.
|
||||
* Author quarterly roadmap and responsible for each development effort.
|
||||
* Making major changes to the technical direction or scope of vLLM and vLLM projects.
|
||||
* Defining the project's release strategy.
|
||||
* Work with model providers, hardware vendors, and key users of vLLM to ensure the project is on the right track.
|
||||
|
||||
### Lead Maintainers
|
||||
|
||||
While Core maintainers assume the day-to-day responsibilities of the project, Lead maintainers are responsible for the overall direction and strategy of the project. The following committee currently shares this role with divided responsibilities:
|
||||
While Core maintainers assume the day-to-day responsibilities of the project, Lead maintainers are responsible for the overall direction and strategy of the project. A committee of @WoosukKwon, @zhuohan123, @simon-mo, @youkaichao, and @robertgshaw2-redhat currently shares this role with divided responsibilities.
|
||||
|
||||
- Woosuk Kwon ([@WoosukKwon](https://github.com/WoosukKwon))
|
||||
- Zhuohan Li ([@zhuohan123](https://github.com/zhuohan123))
|
||||
- Simon Mo ([@simon-mo](https://github.com/simon-mo))
|
||||
- Kaichao You ([@youkaichao](https://github.com/youkaichao))
|
||||
- Robert Shaw ([@robertgshaw2-redhat](https://github.com/robertgshaw2-redhat))
|
||||
The responsibilities of the lead maintainers are:
|
||||
|
||||
**Responsibilities:**
|
||||
|
||||
- Making decisions where consensus among core maintainers cannot be reached.
|
||||
- Adopting changes to the project's technical governance.
|
||||
- Organizing the voting process for new committers.
|
||||
* Making decisions where consensus among core maintainers cannot be reached.
|
||||
* Adopting changes to the project's technical governance.
|
||||
* Organizing the voting process for new committers.
|
||||
|
||||
### Committers and Area Owners
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas and help the community.
|
||||
|
||||
**Responsibilities:**
|
||||
The responsibilities of the committers are:
|
||||
|
||||
- Reviewing PRs and providing feedback.
|
||||
- Addressing issues and questions from the community.
|
||||
- Own specific areas of the codebase and development efforts: reviewing PRs, addressing issues, answering questions, improving documentation.
|
||||
* Reviewing PRs and providing feedback.
|
||||
* Addressing issues and questions from the community.
|
||||
* Own specific areas of the codebase and development efforts: reviewing PRs, addressing issues, answering questions, improving documentation.
|
||||
|
||||
Specially, committers are almost all area owners. They author subsystems, review PRs, refactor code, monitor tests, and ensure compatibility with other areas. All area owners are committers with deep expertise in that area, but not all committers own areas.
|
||||
|
||||
@@ -89,23 +68,23 @@ Any committer can nominate candidates via our private mailing list:
|
||||
|
||||
Committership is highly selective and merit based. The selection criteria requires:
|
||||
|
||||
- **Area expertise**: leading design/implementation of core subsystems, material performance or reliability improvements adopted project‑wide, or accepted RFCs that shape technical direction.
|
||||
- **Sustained contributions**: high‑quality merged contributions and reviews across releases, responsiveness to feedback, and stewardship of code health.
|
||||
- **Community leadership**: mentoring contributors, triaging issues, improving docs, and elevating project standards.
|
||||
* **Area expertise**: leading design/implementation of core subsystems, material performance or reliability improvements adopted project‑wide, or accepted RFCs that shape technical direction.
|
||||
* **Sustained contributions**: high‑quality merged contributions and reviews across releases, responsiveness to feedback, and stewardship of code health.
|
||||
* **Community leadership**: mentoring contributors, triaging issues, improving docs, and elevating project standards.
|
||||
|
||||
To further illustrate, a committer typically satisfies at least two of the following accomplishment patterns:
|
||||
|
||||
- Author of an accepted RFC or design that materially shaped project direction
|
||||
- Measurable, widely adopted performance or reliability improvement in core paths
|
||||
- Long‑term ownership of a subsystem with demonstrable quality and stability gains
|
||||
- Significant cross‑project compatibility or ecosystem enablement work (models, hardware, tooling)
|
||||
* Author of an accepted RFC or design that materially shaped project direction
|
||||
* Measurable, widely adopted performance or reliability improvement in core paths
|
||||
* Long‑term ownership of a subsystem with demonstrable quality and stability gains
|
||||
* Significant cross‑project compatibility or ecosystem enablement work (models, hardware, tooling)
|
||||
|
||||
While there isn't a quantitative bar, past committers have:
|
||||
|
||||
- Submitted approximately 30+ PRs of substantial quality and scope
|
||||
- Provided high-quality reviews of approximately 10+ substantial external contributor PRs
|
||||
- Addressed multiple issues and questions from the community in issues/forums/Slack
|
||||
- Led concentrated efforts on RFCs and their implementation, or significant performance or reliability improvements adopted project‑wide
|
||||
* Submitted approximately 30+ PRs of substantial quality and scope
|
||||
* Provided high-quality reviews of approximately 10+ substantial external contributor PRs
|
||||
* Addressed multiple issues and questions from the community in issues/forums/Slack
|
||||
* Led concentrated efforts on RFCs and their implementation, or significant performance or reliability improvements adopted project‑wide
|
||||
|
||||
### Working Groups
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
| [Intel® Xeon® 6 Processors](https://www.intel.com/content/www/us/en/products/details/processors/xeon.html) |
|
||||
| [Intel® Xeon® 5 Processors](https://www.intel.com/content/www/us/en/products/docs/processors/xeon/5th-gen-xeon-scalable-processors.html) |
|
||||
|
||||
## Recommended Models
|
||||
## Supported Models
|
||||
|
||||
### Text-only Language Models
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
| ----------------------------------------- |
|
||||
| [Intel® Arc™ Pro B-Series Graphics](https://www.intel.com/content/www/us/en/products/docs/discrete-gpus/arc/workstations/b-series/overview.html) |
|
||||
|
||||
## Recommended Models
|
||||
## Supported Models
|
||||
|
||||
### Text-only Language Models
|
||||
|
||||
|
||||
@@ -422,7 +422,7 @@ th {
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | ︎| ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 |`MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 |`MiniMaxAI/MiniMax-M2`, etc. | | ✅︎ |
|
||||
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
|
||||
| `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -686,7 +686,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
|
||||
| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
|
||||
| `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
|
||||
| `KimiK25ForConditionalGeneration` | Kimi-K2.5 | T + I<sup>+</sup> | `moonshotai/Kimi-K2.5` | | ✅︎ |
|
||||
| `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
|
||||
| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -39,7 +39,6 @@ Launch Claude Code with environment variables pointing to your vLLM server:
|
||||
```bash
|
||||
ANTHROPIC_BASE_URL=http://localhost:8000 \
|
||||
ANTHROPIC_API_KEY=dummy \
|
||||
ANTHROPIC_AUTH_TOKEN=dummy \
|
||||
ANTHROPIC_DEFAULT_OPUS_MODEL=my-model \
|
||||
ANTHROPIC_DEFAULT_SONNET_MODEL=my-model \
|
||||
ANTHROPIC_DEFAULT_HAIKU_MODEL=my-model \
|
||||
@@ -52,7 +51,6 @@ The environment variables:
|
||||
| -------------------------------- | --------------------------------------------------------------------- |
|
||||
| `ANTHROPIC_BASE_URL` | Points to your vLLM server (default port is 8000) |
|
||||
| `ANTHROPIC_API_KEY` | Can be any value since vLLM doesn't require authentication by default |
|
||||
| `ANTHROPIC_AUTH_TOKEN` | Is required. Can be any value. |
|
||||
| `ANTHROPIC_DEFAULT_OPUS_MODEL` | Model name for Opus-tier requests |
|
||||
| `ANTHROPIC_DEFAULT_SONNET_MODEL` | Model name for Sonnet-tier requests |
|
||||
| `ANTHROPIC_DEFAULT_HAIKU_MODEL` | Model name for Haiku-tier requests |
|
||||
|
||||
@@ -112,36 +112,10 @@ def get_multi_audios_query() -> QueryResult:
|
||||
)
|
||||
|
||||
|
||||
def get_multi_images_query() -> QueryResult:
|
||||
question = "What are the differences between these two images?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_bos|><|IMAGE|><|vision_eos|>"
|
||||
"<|vision_bos|><|IMAGE|><|vision_eos|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n"
|
||||
)
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"image": [
|
||||
convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
|
||||
convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB"),
|
||||
],
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={
|
||||
"image": 2,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
query_map = {
|
||||
"mixed_modalities": get_mixed_modalities_query,
|
||||
"use_audio_in_video": get_use_audio_in_video_query,
|
||||
"multi_audios": get_multi_audios_query,
|
||||
"multi_images": get_multi_images_query,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -59,10 +59,8 @@ class PrithviMAE:
|
||||
input_data = input_data[0]
|
||||
|
||||
mm_data = {
|
||||
"image": {
|
||||
"pixel_values": input_data,
|
||||
"location_coords": location_coords,
|
||||
}
|
||||
"pixel_values": input_data,
|
||||
"location_coords": location_coords,
|
||||
}
|
||||
|
||||
prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}
|
||||
|
||||
+2
-3
@@ -9,7 +9,7 @@ requires = [
|
||||
"torch == 2.9.1",
|
||||
"wheel",
|
||||
"jinja2",
|
||||
"grpcio-tools",
|
||||
"grpcio-tools>=1.76.0",
|
||||
]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
@@ -132,8 +132,7 @@ python = "./.venv"
|
||||
# these files may be written in non english words
|
||||
extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*",
|
||||
"benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
|
||||
"docs/governance/process.md"]
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*"]
|
||||
ignore-hidden = true
|
||||
ignore-files = true
|
||||
ignore-dot = true
|
||||
|
||||
@@ -9,5 +9,5 @@ wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
build
|
||||
protobuf
|
||||
grpcio-tools
|
||||
protobuf>=6.33.2
|
||||
grpcio-tools>=1.76.0
|
||||
|
||||
@@ -9,7 +9,7 @@ blake3
|
||||
py-cpuinfo
|
||||
transformers >= 4.56.0, < 5
|
||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||
protobuf # Required by LlamaTokenizer, gRPC.
|
||||
protobuf >= 6.30.0 # Required by LlamaTokenizer, gRPC.
|
||||
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
|
||||
aiohttp
|
||||
openai >= 1.99.1 # For Responses API with reasoning content
|
||||
@@ -51,5 +51,5 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
|
||||
anthropic >= 0.71.0
|
||||
model-hosting-container-standards >= 0.1.13, < 1.0.0
|
||||
mcp
|
||||
grpcio
|
||||
grpcio-reflection
|
||||
grpcio>=1.76.0
|
||||
grpcio-reflection>=1.76.0
|
||||
@@ -438,49 +438,6 @@ class precompiled_wheel_utils:
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def detect_system_cuda_variant() -> str:
|
||||
"""Auto-detect CUDA variant from torch, nvidia-smi, or env default."""
|
||||
|
||||
# Map CUDA major version to hosted wheel variants on wheels.vllm.ai
|
||||
supported = {12: "cu129", 13: "cu130"}
|
||||
|
||||
# Respect explicitly set VLLM_MAIN_CUDA_VERSION
|
||||
if envs.is_set("VLLM_MAIN_CUDA_VERSION"):
|
||||
v = envs.VLLM_MAIN_CUDA_VERSION
|
||||
print(f"Using VLLM_MAIN_CUDA_VERSION={v}")
|
||||
return "cu" + v.replace(".", "")[:3]
|
||||
|
||||
# Try torch.version.cuda
|
||||
cuda_version = None
|
||||
try:
|
||||
import torch
|
||||
|
||||
cuda_version = torch.version.cuda
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Try nvidia-smi
|
||||
if not cuda_version:
|
||||
try:
|
||||
out = subprocess.run(
|
||||
["nvidia-smi"], capture_output=True, text=True, timeout=10
|
||||
)
|
||||
if m := re.search(r"CUDA Version:\s*(\d+\.\d+)", out.stdout):
|
||||
cuda_version = m.group(1)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fall back to default
|
||||
if not cuda_version:
|
||||
cuda_version = envs.VLLM_MAIN_CUDA_VERSION
|
||||
|
||||
# Map to supported variant
|
||||
major = int(cuda_version.split(".")[0])
|
||||
variant = supported.get(major, supported[max(supported)])
|
||||
print(f"Detected CUDA {cuda_version}, using variant {variant}")
|
||||
return variant
|
||||
|
||||
@staticmethod
|
||||
def find_local_rocm_wheel() -> str | None:
|
||||
"""Search for a local vllm wheel in common locations."""
|
||||
@@ -556,8 +513,8 @@ class precompiled_wheel_utils:
|
||||
1. user-specified wheel location (can be either local or remote, via
|
||||
VLLM_PRECOMPILED_WHEEL_LOCATION)
|
||||
2. user-specified variant (VLLM_PRECOMPILED_WHEEL_VARIANT) from nightly repo
|
||||
or auto-detected CUDA variant based on system (torch, nvidia-smi)
|
||||
3. the default variant from nightly repo
|
||||
3. the variant corresponding to VLLM_MAIN_CUDA_VERSION from nightly repo
|
||||
4. the default variant from nightly repo
|
||||
|
||||
If downloading from the nightly repo, the commit can be specified via
|
||||
VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch
|
||||
@@ -576,11 +533,9 @@ class precompiled_wheel_utils:
|
||||
import platform
|
||||
|
||||
arch = platform.machine()
|
||||
# try to fetch the wheel metadata from the nightly wheel repo,
|
||||
# detecting CUDA variant from system if not specified
|
||||
variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
|
||||
if variant is None:
|
||||
variant = precompiled_wheel_utils.detect_system_cuda_variant()
|
||||
# try to fetch the wheel metadata from the nightly wheel repo
|
||||
main_variant = "cu" + envs.VLLM_MAIN_CUDA_VERSION.replace(".", "")
|
||||
variant = os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", main_variant)
|
||||
commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
|
||||
if not commit or len(commit) != 40:
|
||||
print(
|
||||
|
||||
@@ -164,40 +164,3 @@ def test_classes_are_types():
|
||||
pass
|
||||
|
||||
assert endswith_fqname(LocalDummy, ".LocalDummy")
|
||||
|
||||
|
||||
def test_envs_compile_factors_stable():
|
||||
"""Test that envs.compile_factors() hash is stable across fresh initializations.
|
||||
|
||||
Uses subprocesses to ensure env vars with dynamic defaults (like UUIDs)
|
||||
are freshly generated each time, verifying they're properly ignored.
|
||||
"""
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
code = """
|
||||
import sys
|
||||
import logging
|
||||
logging.disable(logging.CRITICAL)
|
||||
from vllm import envs
|
||||
from vllm.config.utils import hash_factors
|
||||
print(hash_factors(envs.compile_factors()))
|
||||
"""
|
||||
|
||||
def get_hash_in_subprocess():
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
env={**dict(__import__("os").environ), "VLLM_LOGGING_LEVEL": "ERROR"},
|
||||
)
|
||||
return result.stdout.strip()
|
||||
|
||||
hash1 = get_hash_in_subprocess()
|
||||
hash2 = get_hash_in_subprocess()
|
||||
|
||||
assert hash1 == hash2, (
|
||||
"compile_factors hash differs between fresh initializations - "
|
||||
"dynamic env vars may not be properly ignored"
|
||||
)
|
||||
|
||||
@@ -992,7 +992,7 @@ async def test_mcp_tool_multi_turn(client: OpenAI, model_name: str, server):
|
||||
# First turn - make a calculation
|
||||
response1 = await client.responses.create(
|
||||
model=model_name,
|
||||
input="Calculate 1234 * 4567 using python tool and print the result.",
|
||||
input="Calculate 123 * 456 using python and print the result.",
|
||||
tools=tools,
|
||||
temperature=0.0,
|
||||
instructions=(
|
||||
|
||||
@@ -13,10 +13,31 @@ from vllm.utils.serial_utils import tensor2base64
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
|
||||
def _terratorch_dummy_messages():
|
||||
pixel_values = torch.full((6, 512, 512), 1.0, dtype=torch.float16)
|
||||
location_coords = torch.full((1, 2), 1.0, dtype=torch.float16)
|
||||
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"pixel_values": tensor2base64(pixel_values),
|
||||
"location_coords": tensor2base64(location_coords),
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model_name", ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
|
||||
)
|
||||
def test_single_content(model_name: str):
|
||||
def test_single_request(model_name: str):
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
@@ -38,24 +59,7 @@ def test_single_content(model_name: str):
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"pixel_values": tensor2base64(
|
||||
torch.ones((6, 512, 512), dtype=torch.float16)
|
||||
),
|
||||
"location_coords": tensor2base64(
|
||||
torch.ones((1, 2), dtype=torch.float16)
|
||||
),
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"messages": _terratorch_dummy_messages(),
|
||||
"encoding_format": "base64",
|
||||
},
|
||||
)
|
||||
@@ -65,87 +69,3 @@ def test_single_content(model_name: str):
|
||||
|
||||
np_response = np.frombuffer(base64.b64decode(output), dtype=np.float32)
|
||||
assert len(np_response) == 524288
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["Qwen/Qwen3-VL-2B-Instruct"])
|
||||
def test_multi_content(model_name: str):
|
||||
args = [
|
||||
"--enforce-eager",
|
||||
"--max-num-seqs",
|
||||
"32",
|
||||
"--max-model-len",
|
||||
"8192",
|
||||
"--enable-mm-embeds",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(model_name, args) as server:
|
||||
client = server.get_client()
|
||||
|
||||
# Image only
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
max_tokens=5,
|
||||
)
|
||||
|
||||
assert chat_completion.id is not None
|
||||
assert len(chat_completion.choices) == 1
|
||||
|
||||
# Interleaved text and image
|
||||
chat_completion = client.chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
{"type": "text", "text": "OCR:"},
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(torch.zeros(220, 8192)),
|
||||
"image_grid_thw": tensor2base64(
|
||||
torch.tensor([1, 22, 40])
|
||||
),
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
max_tokens=5,
|
||||
)
|
||||
|
||||
assert chat_completion.id is not None
|
||||
assert len(chat_completion.choices) == 1
|
||||
|
||||
@@ -68,16 +68,6 @@ def phi3v_model_config_image_embeds():
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def qwen25omni_model_config_image_embeds():
|
||||
return ModelConfig(
|
||||
QWEN25OMNI_MODEL_ID,
|
||||
runner="generate",
|
||||
limit_mm_per_prompt={"image": 2},
|
||||
enable_mm_embeds=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def qwen2_audio_model_config():
|
||||
return ModelConfig(
|
||||
@@ -833,8 +823,7 @@ def test_parse_chat_messages_audio_embeds_with_string(
|
||||
import torch
|
||||
|
||||
# Create a sample audio embedding tensor
|
||||
hidden_size = audio_embeds_model_config.get_inputs_embeds_size()
|
||||
audio_embedding = torch.randn(1, 128, hidden_size)
|
||||
audio_embedding = torch.randn(1, 128, 768)
|
||||
|
||||
# Encode it as base64
|
||||
base64_audio_embedding = tensor2base64(audio_embedding)
|
||||
@@ -876,8 +865,7 @@ async def test_parse_chat_messages_audio_embeds_async(
|
||||
import torch
|
||||
|
||||
# Create a sample audio embedding tensor
|
||||
hidden_size = audio_embeds_model_config.get_inputs_embeds_size()
|
||||
audio_embedding = torch.randn(1, 128, hidden_size)
|
||||
audio_embedding = torch.randn(1, 128, 768)
|
||||
|
||||
# Encode it as base64
|
||||
base64_audio_embedding = tensor2base64(audio_embedding)
|
||||
@@ -920,9 +908,8 @@ def test_parse_chat_messages_multiple_image_embeds(
|
||||
can be provided in a single request, similar to regular images.
|
||||
"""
|
||||
# Create two sample image embedding tensors
|
||||
hidden_size = phi3v_model_config_image_embeds.get_inputs_embeds_size()
|
||||
image_embedding_1 = torch.randn(256, hidden_size)
|
||||
image_embedding_2 = torch.randn(128, hidden_size)
|
||||
image_embedding_1 = torch.randn(256, 1024)
|
||||
image_embedding_2 = torch.randn(128, 1024)
|
||||
|
||||
# Encode them as base64 using the convenience function
|
||||
base64_image_embedding_1 = tensor2base64(image_embedding_1)
|
||||
@@ -1035,9 +1022,8 @@ async def test_parse_chat_messages_multiple_image_embeds_async(
|
||||
This validates the AsyncMultiModalItemTracker also supports multiple embeddings.
|
||||
"""
|
||||
# Create two sample image embedding tensors
|
||||
hidden_size = phi3v_model_config_image_embeds.get_inputs_embeds_size()
|
||||
image_embedding_1 = torch.randn(200, hidden_size)
|
||||
image_embedding_2 = torch.randn(150, hidden_size)
|
||||
image_embedding_1 = torch.randn(200, 768)
|
||||
image_embedding_2 = torch.randn(150, 768)
|
||||
|
||||
# Encode them as base64 using the convenience function
|
||||
base64_image_embedding_1 = tensor2base64(image_embedding_1)
|
||||
@@ -1159,14 +1145,13 @@ def test_parse_chat_messages_empty_dict_image_embeds(
|
||||
|
||||
|
||||
def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
qwen25omni_model_config_image_embeds,
|
||||
phi3v_model_config_image_embeds,
|
||||
):
|
||||
"""Test that multiple dictionaries for image_embeds is handled without errors."""
|
||||
# Create two sample image embedding tensors
|
||||
batch_size = 2
|
||||
hidden_size = qwen25omni_model_config_image_embeds.get_inputs_embeds_size()
|
||||
image_embeds = torch.randn(batch_size * 220, hidden_size)
|
||||
image_grid_thw = torch.tensor([[1, 22, 40] for _ in range(batch_size)])
|
||||
image_embedding_1 = torch.randn(batch_size, 256, 1024)
|
||||
image_embedding_2 = torch.randn(batch_size, 3)
|
||||
|
||||
conversation, mm_data, mm_uuids = parse_chat_messages(
|
||||
[
|
||||
@@ -1176,20 +1161,18 @@ def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
{
|
||||
"type": "image_embeds",
|
||||
"image_embeds": {
|
||||
"image_embeds": tensor2base64(embeds),
|
||||
"image_grid_thw": tensor2base64(grid_thw),
|
||||
"image_embedding_1": tensor2base64(p),
|
||||
"image_embedding_2": tensor2base64(i),
|
||||
},
|
||||
}
|
||||
for embeds, grid_thw in zip(
|
||||
image_embeds.chunk(batch_size), image_grid_thw
|
||||
)
|
||||
for p, i in zip(image_embedding_1, image_embedding_2)
|
||||
]
|
||||
+ [
|
||||
{"type": "text", "text": "Describe these two images."},
|
||||
],
|
||||
}
|
||||
],
|
||||
qwen25omni_model_config_image_embeds,
|
||||
phi3v_model_config_image_embeds,
|
||||
content_format="string",
|
||||
)
|
||||
|
||||
@@ -1197,8 +1180,7 @@ def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
assert conversation == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "<|vision_start|><|IMAGE|><|vision_end|>\n"
|
||||
"<|vision_start|><|IMAGE|><|vision_end|>\nDescribe these two images.",
|
||||
"content": "<|image_1|>\n<|image_2|>\nDescribe these two images.",
|
||||
}
|
||||
]
|
||||
|
||||
@@ -1209,10 +1191,10 @@ def test_parse_chat_messages_multiple_dict_image_embeds(
|
||||
assert len(mm_data["image"]) == batch_size
|
||||
|
||||
# Verify each embedding has the correct shape
|
||||
assert isinstance(mm_data["image"]["image_embeds"], torch.Tensor)
|
||||
assert mm_data["image"]["image_embeds"].shape == image_embeds.shape
|
||||
assert isinstance(mm_data["image"]["image_grid_thw"], torch.Tensor)
|
||||
assert mm_data["image"]["image_grid_thw"].shape == image_grid_thw.shape
|
||||
assert isinstance(mm_data["image"]["image_embedding_1"], torch.Tensor)
|
||||
assert mm_data["image"]["image_embedding_1"].shape == image_embedding_1.shape
|
||||
assert isinstance(mm_data["image"]["image_embedding_2"], torch.Tensor)
|
||||
assert mm_data["image"]["image_embedding_2"].shape == image_embedding_2.shape
|
||||
|
||||
# Verify UUIDs (None since we didn't provide any)
|
||||
_assert_mm_uuids(mm_uuids, batch_size, expected_uuids=[None, None])
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
@@ -5,10 +5,12 @@ Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-marlin.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-CT-fi-cutlass-dp-ep.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-marlin.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml
|
||||
Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass-dp-ep.yaml
|
||||
Llama-4-Scout-BF16-fi-cutlass.yaml
|
||||
Llama-4-Scout-BF16-triton.yaml
|
||||
Mixtral-8x7B-BF16-fi-cutlass.yaml
|
||||
|
||||
@@ -107,14 +107,10 @@ def test_flashinfer_nvfp4_gemm(
|
||||
# from checkpoints are in linear scales.
|
||||
# So instead of needing to swizzle for cutlass as in modelopt.py,
|
||||
# we need to unswizzle for trtllm here.
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
|
||||
a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend
|
||||
)
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale, backend)
|
||||
is_sf_128x4_layout = not (backend == "trtllm" and m <= 32)
|
||||
|
||||
b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(
|
||||
b_dtype, b_global_scale, is_sf_swizzled_layout=True
|
||||
)
|
||||
b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
|
||||
|
||||
# get_ref_results unswizzles the scales internally.
|
||||
expected_out = get_ref_results(
|
||||
|
||||
@@ -59,7 +59,7 @@ if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"These tests require gptq_marlin_repack,"
|
||||
"marlin_int4_fp8_preprocess, gptq_marlin_24_gemm,"
|
||||
"or marlin_gemm which are not supported on ROCm.",
|
||||
"or gptq_marlin_gemm which are not supported on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
@@ -417,7 +417,7 @@ def marlin_generate_valid_test_cases():
|
||||
),
|
||||
marlin_generate_valid_test_cases(),
|
||||
)
|
||||
def test_marlin_gemm(
|
||||
def test_gptq_marlin_gemm(
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
@@ -511,7 +511,7 @@ def test_marlin_gemm(
|
||||
|
||||
output = torch.empty((size_m, size_n), dtype=dtype, device=a_input.device)
|
||||
|
||||
output = ops.marlin_gemm(
|
||||
output = ops.gptq_marlin_gemm(
|
||||
a_input,
|
||||
output,
|
||||
marlin_q_w,
|
||||
@@ -646,7 +646,7 @@ def test_marlin_gemm_subset_input():
|
||||
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
|
||||
workspace = marlin_make_workspace_new(a_input.device)
|
||||
|
||||
output = ops.marlin_gemm(
|
||||
output = ops.gptq_marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
@@ -695,7 +695,7 @@ def test_marlin_gemm_with_bias(size_m):
|
||||
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
|
||||
workspace = marlin_make_workspace_new(a_input.device)
|
||||
|
||||
output = ops.marlin_gemm(
|
||||
output = ops.gptq_marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
|
||||
@@ -27,12 +27,6 @@ PAD_SHAPES = [
|
||||
(150, 128),
|
||||
(150, 48),
|
||||
(90, 80),
|
||||
(128, 512),
|
||||
(128, 1024),
|
||||
(128, 2048),
|
||||
(64, 7168),
|
||||
(64, 7152),
|
||||
(32, 14336),
|
||||
]
|
||||
SEEDS = [42]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
@@ -179,25 +173,3 @@ def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_padded_no_sf_swizzled(pad_shape: tuple[int, int]) -> None:
|
||||
dtype = torch.float16
|
||||
set_random_seed(42)
|
||||
torch.set_default_device("cuda:0")
|
||||
|
||||
m, n = pad_shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = ops.scaled_fp4_quant(x, global_scale, is_sf_swizzled_layout=False)
|
||||
scale_ans = out_scale.to(torch.float32)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
@@ -2,11 +2,13 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Integration tests for FlexAttention backend vs default backend"""
|
||||
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from packaging import version
|
||||
|
||||
from tests.utils import set_random_seed
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
@@ -25,6 +27,15 @@ MINIMUM_TORCH_VERSION = version.parse("2.7.0")
|
||||
DIRECT_BUILD_VERSION = version.parse("2.9.dev0")
|
||||
|
||||
|
||||
def set_seed(seed):
|
||||
"""Set seeds for reproducibility"""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
|
||||
reason="CUDA not available or PyTorch version < 2.7",
|
||||
@@ -46,7 +57,7 @@ def test_flex_attention_vs_default_backend(vllm_runner):
|
||||
]
|
||||
|
||||
# Run with flex attention
|
||||
set_random_seed(seed)
|
||||
set_seed(seed)
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
runner="generate",
|
||||
@@ -60,7 +71,7 @@ def test_flex_attention_vs_default_backend(vllm_runner):
|
||||
)
|
||||
|
||||
# Run with default backend
|
||||
set_random_seed(seed)
|
||||
set_seed(seed)
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
runner="generate",
|
||||
|
||||
@@ -7,6 +7,14 @@ import torch
|
||||
from ....conftest import VllmRunner
|
||||
|
||||
|
||||
def generate_test_mm_data():
|
||||
mm_data = {
|
||||
"pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
|
||||
"location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
|
||||
}
|
||||
return mm_data
|
||||
|
||||
|
||||
def _run_test(
|
||||
vllm_runner: type[VllmRunner],
|
||||
model: str,
|
||||
@@ -15,12 +23,7 @@ def _run_test(
|
||||
{
|
||||
# This model deals with no text input
|
||||
"prompt_token_ids": [1],
|
||||
"multi_modal_data": {
|
||||
"image": {
|
||||
"pixel_values": torch.ones((6, 512, 512), dtype=torch.float16),
|
||||
"location_coords": torch.ones((1, 2), dtype=torch.float16),
|
||||
}
|
||||
},
|
||||
"multi_modal_data": generate_test_mm_data(),
|
||||
}
|
||||
for _ in range(10)
|
||||
]
|
||||
|
||||
@@ -48,6 +48,18 @@ VideoInput: TypeAlias = (
|
||||
AudioInput = list[tuple[np.ndarray, int]]
|
||||
|
||||
|
||||
MM_OPTIONS_OVERRIDES = {
|
||||
# Qwen3-VL's default profiling video size (64x64) can cause trouble
|
||||
# after resizing, so we override it here for testing.
|
||||
"qwen3_vl": dict(
|
||||
video=VideoDummyOptions(num_frames=128, width=256, height=256),
|
||||
),
|
||||
"qwen3_vl_moe": dict(
|
||||
video=VideoDummyOptions(num_frames=128, width=256, height=256),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _resize_data(
|
||||
_data: Image.Image | np.ndarray, size_factor: float
|
||||
) -> Image.Image | np.ndarray:
|
||||
@@ -61,12 +73,12 @@ def _resize_data(
|
||||
elif is_list_of(_data, Image.Image):
|
||||
W, H = next(iter(_data)).width, next(iter(_data)).height
|
||||
T = len(_data)
|
||||
T, W, H = map(lambda x: max(int(x * size_factor), 2), (T, W, H))
|
||||
T, W, H = map(lambda x: max(int(x * size_factor), 1), (T, W, H))
|
||||
return [d.resize((W, H)) for d in _data[:T]]
|
||||
# Video input with numpy arrays
|
||||
elif isinstance(_data, np.ndarray) and _data.ndim >= 4:
|
||||
T, H, W, C = _data.shape[-4:]
|
||||
T, H, W = map(lambda x: max(int(x * size_factor), 2), (T, H, W))
|
||||
T, H, W = map(lambda x: max(int(x * size_factor), 1), (T, H, W))
|
||||
return _data[..., :T, :H, :W, :C]
|
||||
# Audio input
|
||||
elif isinstance(_data, np.ndarray) and _data.ndim == 1:
|
||||
@@ -91,6 +103,8 @@ def create_batched_mm_kwargs(
|
||||
processor: BaseMultiModalProcessor,
|
||||
size_factors: tuple[float, ...] = (1.0, 0.5, 0.25),
|
||||
) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
|
||||
model_type = model_config.hf_config.model_type
|
||||
|
||||
processing_info = processor.info
|
||||
dummy_inputs = processor.dummy_inputs
|
||||
supported_mm_limits = processing_info.get_supported_mm_limits()
|
||||
@@ -101,6 +115,7 @@ def create_batched_mm_kwargs(
|
||||
processor_inputs = dummy_inputs.get_dummy_processor_inputs(
|
||||
seq_len=model_config.max_model_len,
|
||||
mm_counts=mm_counts,
|
||||
mm_options=MM_OPTIONS_OVERRIDES.get(model_type),
|
||||
)
|
||||
mm_data = processor_inputs.mm_data
|
||||
resized_mm_data = {
|
||||
|
||||
@@ -771,11 +771,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
)
|
||||
},
|
||||
),
|
||||
"KimiK25ForConditionalGeneration": _HfExamplesInfo(
|
||||
"moonshotai/Kimi-K2.5",
|
||||
trust_remote_code=True,
|
||||
is_available_online=False,
|
||||
),
|
||||
"LightOnOCRForConditionalGeneration": _HfExamplesInfo(
|
||||
"lightonai/LightOnOCR-1B-1025"
|
||||
),
|
||||
|
||||
@@ -349,10 +349,8 @@ class PrithviMultimodalDataProcessor(IOProcessor):
|
||||
{
|
||||
"prompt_token_ids": [1],
|
||||
"multi_modal_data": {
|
||||
"image": {
|
||||
"pixel_values": window.to(torch.float16)[0],
|
||||
"location_coords": location_coords.to(torch.float16),
|
||||
}
|
||||
"pixel_values": window.to(torch.float16)[0],
|
||||
"location_coords": location_coords.to(torch.float16),
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
@@ -20,7 +20,7 @@ MM_BEAM_WIDTHS = [2]
|
||||
MODELS = ["TinyLlama/TinyLlama-1.1B-Chat-v1.0"]
|
||||
|
||||
|
||||
@pytest.mark.skip_v1 # V1 engine does not yet support beam search
|
||||
@pytest.mark.skip_v1 # FIXME: This fails on V1 right now.
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
@@ -62,7 +62,7 @@ def test_beam_search_single_input(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip_v1 # V1 engine does not yet support beam search
|
||||
@pytest.mark.skip_v1 # FIXME: This fails on V1 right now.
|
||||
@pytest.mark.parametrize("model", MODELS)
|
||||
@pytest.mark.parametrize("dtype", ["half"])
|
||||
@pytest.mark.parametrize("max_tokens", MAX_TOKENS)
|
||||
|
||||
+1
-4
@@ -59,10 +59,7 @@ from vllm.tokenizers import get_tokenizer
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
from vllm.utils.mem_constants import GB_bytes
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import (
|
||||
cuda_device_count_stateless,
|
||||
set_random_seed, # noqa: F401 - re-exported for use in test files
|
||||
)
|
||||
from vllm.utils.torch_utils import cuda_device_count_stateless
|
||||
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
|
||||
@@ -13,7 +13,6 @@ from tests.v1.attention.utils import (
|
||||
create_common_attn_metadata,
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
try_backend_includes_kv_cache_update,
|
||||
try_get_attention_backend,
|
||||
)
|
||||
from vllm.config import ModelConfig
|
||||
@@ -296,10 +295,6 @@ def run_attention_backend(
|
||||
# Run forward pass
|
||||
# NOTE: The query, key, and value are already shaped correctly
|
||||
# in the calling test function.
|
||||
if not try_backend_includes_kv_cache_update(actual_backend):
|
||||
impl.do_kv_cache_update(
|
||||
mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
|
||||
)
|
||||
output = impl.forward(
|
||||
mock_layer, query, key, value, kv_cache, attn_metadata, output=output
|
||||
)
|
||||
|
||||
@@ -130,18 +130,6 @@ def try_get_attention_backend(
|
||||
raise AssertionError("unreachable") from None
|
||||
|
||||
|
||||
def try_backend_includes_kv_cache_update(
|
||||
backend: AttentionBackendEnum,
|
||||
) -> bool:
|
||||
"""Try to get the attention backend class, skipping test if not found."""
|
||||
try:
|
||||
backend_class = backend.get_class()
|
||||
return backend_class.forward_includes_kv_cache_update
|
||||
except ImportError as e:
|
||||
pytest.skip(f"{backend.name} not available: {e}")
|
||||
raise AssertionError("unreachable") from None
|
||||
|
||||
|
||||
def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
|
||||
"""Create a FullAttentionSpec from ModelParams only."""
|
||||
return FullAttentionSpec(
|
||||
|
||||
@@ -650,9 +650,9 @@ def test_schedule_order(enable_chunked_prefill: bool):
|
||||
)
|
||||
|
||||
# long requests
|
||||
requests = create_requests(num_requests=2, num_tokens=800, req_ids=["1", "2"])
|
||||
requests = create_requests(num_requests=2, num_tokens=800)
|
||||
# short requests
|
||||
requests += create_requests(num_requests=2, num_tokens=10, req_ids=["3", "4"])
|
||||
requests += create_requests(num_requests=2, num_tokens=10)
|
||||
|
||||
for request in requests:
|
||||
scheduler.add_request(request)
|
||||
@@ -1806,12 +1806,6 @@ def test_priority_scheduling_mixed_priority_and_arrival():
|
||||
assert scheduled_req_ids == ["3", "2", "1", "0"]
|
||||
|
||||
|
||||
# This test had previously been passing due to its use of duplicate
|
||||
# request ids which resulted in incorrect behavior.
|
||||
# Now that the duplicate req ids had been fixed it fails and
|
||||
# investigation is needed into whether the priority scheduling
|
||||
# preemption logic is working as designed or not.
|
||||
@pytest.mark.skip("needs investigation")
|
||||
def test_priority_scheduling_preemption():
|
||||
"""Test that priority scheduling preempts
|
||||
lower priority requests when memory is constrained."""
|
||||
@@ -1828,8 +1822,7 @@ def test_priority_scheduling_preemption():
|
||||
num_requests=2,
|
||||
priorities=[5, 5], # Low priority
|
||||
arrival_times=[1.0, 2.0],
|
||||
num_tokens=30, # Large enough to consume significant memory,
|
||||
req_ids=["lo1", "lo2"],
|
||||
num_tokens=30, # Large enough to consume significant memory
|
||||
)
|
||||
|
||||
# Add and schedule low priority requests
|
||||
@@ -1862,7 +1855,6 @@ def test_priority_scheduling_preemption():
|
||||
priorities=[0], # High priority
|
||||
arrival_times=[3.0],
|
||||
num_tokens=30, # Large enough to require significant memory
|
||||
req_ids=["hi1"],
|
||||
)[0]
|
||||
|
||||
scheduler.add_request(high_priority_request)
|
||||
@@ -1884,13 +1876,13 @@ def test_priority_scheduling_preemption():
|
||||
output2 = scheduler.schedule()
|
||||
assert len(output2.scheduled_new_reqs) == 1
|
||||
# High priority request
|
||||
assert output2.scheduled_new_reqs[0].req_id == "hi1"
|
||||
assert output2.scheduled_new_reqs[0].req_id == "0"
|
||||
else:
|
||||
# No preemption needed - all requests fit
|
||||
# This is also valid behavior if memory allows
|
||||
assert len(output.scheduled_new_reqs) == 1
|
||||
# High priority request
|
||||
assert output.scheduled_new_reqs[0].req_id == "hi1"
|
||||
assert output.scheduled_new_reqs[0].req_id == "0"
|
||||
|
||||
|
||||
def test_priority_scheduling_no_preemption_when_space_available():
|
||||
@@ -1903,11 +1895,7 @@ def test_priority_scheduling_no_preemption_when_space_available():
|
||||
|
||||
# Add two low-priority running requests
|
||||
low_priority_requests = create_requests_with_priority(
|
||||
num_requests=2,
|
||||
priorities=[5, 5],
|
||||
arrival_times=[1.0, 2.0],
|
||||
num_tokens=30,
|
||||
req_ids=["lo1", "lo2"],
|
||||
num_requests=2, priorities=[5, 5], arrival_times=[1.0, 2.0], num_tokens=30
|
||||
)
|
||||
|
||||
for request in low_priority_requests:
|
||||
@@ -1928,11 +1916,7 @@ def test_priority_scheduling_no_preemption_when_space_available():
|
||||
|
||||
# Add high-priority request
|
||||
high_priority_request = create_requests_with_priority(
|
||||
num_requests=1,
|
||||
priorities=[0],
|
||||
arrival_times=[3.0],
|
||||
num_tokens=30,
|
||||
req_ids=["hi1"],
|
||||
num_requests=1, priorities=[0], arrival_times=[3.0], num_tokens=30
|
||||
)[0]
|
||||
|
||||
scheduler.add_request(high_priority_request)
|
||||
|
||||
@@ -1,656 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
End-to-end tests for the streaming input feature in AsyncLLM.
|
||||
|
||||
These tests verify that:
|
||||
1. Streaming inputs work correctly with bunched inputs (queued)
|
||||
2. Streaming inputs work correctly with spaced out inputs
|
||||
3. Outputs are equivalent whether inputs are bunched or spaced
|
||||
4. Cancelling the output stream correctly aborts the session
|
||||
5. Closing the input stream correctly signals completion
|
||||
6. Queued inputs are cancelled when the session is aborted
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from vllm import SamplingParams
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.engine.async_llm import AsyncLLM, StreamingInput
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip(reason="V1 currently only supported on CUDA.", allow_module_level=True)
|
||||
|
||||
# Use a small model that doesn't require authentication for fast tests
|
||||
MODEL = "facebook/opt-125m"
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="module", loop_scope="module")
|
||||
async def engine():
|
||||
"""Create an AsyncLLM engine for the test.
|
||||
|
||||
Note: Using function scope because pytest_asyncio creates a new event loop
|
||||
for each test, and the output_handler task gets cancelled between tests
|
||||
with module scope.
|
||||
"""
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=MODEL, enforce_eager=True, gpu_memory_utilization=0.7
|
||||
)
|
||||
with set_default_torch_num_threads(1):
|
||||
engine = AsyncLLM.from_engine_args(engine_args)
|
||||
try:
|
||||
yield engine
|
||||
finally:
|
||||
engine.shutdown()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
|
||||
def get_sampling_params(max_tokens: int = 20) -> SamplingParams:
|
||||
"""Create sampling params for streaming input tests."""
|
||||
return SamplingParams(
|
||||
max_tokens=max_tokens,
|
||||
ignore_eos=True,
|
||||
output_kind=RequestOutputKind.DELTA,
|
||||
temperature=0.0, # Deterministic for reproducibility
|
||||
)
|
||||
|
||||
|
||||
async def collect_outputs(
|
||||
output_gen: AsyncGenerator[RequestOutput, None],
|
||||
) -> tuple[list[RequestOutput], str]:
|
||||
"""Collect all outputs from a generate call, return outputs and full text."""
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
async for output in output_gen:
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
return outputs, full_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_bunched(engine: AsyncLLM):
|
||||
"""Test streaming input where all inputs are sent at once (bunched/queued).
|
||||
|
||||
This tests the case where multiple inputs arrive before any completes.
|
||||
The inputs should be queued and processed in sequence.
|
||||
"""
|
||||
request_id = "test_bunched"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
# Create an input generator that yields all inputs quickly
|
||||
async def bunched_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# Send multiple inputs rapidly - they should be queued
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
yield StreamingInput(prompt=" Alice and I like")
|
||||
yield StreamingInput(prompt=" to code in Python")
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(
|
||||
bunched_input_generator(),
|
||||
sampling_params,
|
||||
request_id,
|
||||
)
|
||||
)
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify the final output is marked as finished
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify intermediate outputs are not marked as finished
|
||||
for output in outputs[:-1]:
|
||||
assert not output.finished, "Intermediate outputs should not be finished"
|
||||
|
||||
# Verify we generated some text
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
print(f"Bunched test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_spaced(engine: AsyncLLM):
|
||||
"""Test streaming input where inputs are spaced out.
|
||||
|
||||
This tests the case where each input completes processing before the
|
||||
next one is sent. Each chunk should be prefilled, generate tokens,
|
||||
then the next chunk should be processed.
|
||||
"""
|
||||
request_id = "test_spaced"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
# Track when each input is sent
|
||||
input_times: list[float] = []
|
||||
outputs_per_chunk: list[int] = [0, 0, 0]
|
||||
current_chunk = 0
|
||||
|
||||
async def spaced_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal current_chunk
|
||||
import time
|
||||
|
||||
# First input
|
||||
input_times.append(time.time())
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
current_chunk = 0
|
||||
|
||||
# Wait for some outputs to be generated
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Second input
|
||||
input_times.append(time.time())
|
||||
current_chunk = 1
|
||||
yield StreamingInput(prompt=" Alice and I like")
|
||||
|
||||
# Wait for some outputs
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Third input
|
||||
input_times.append(time.time())
|
||||
current_chunk = 2
|
||||
yield StreamingInput(prompt=" to code in Python")
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
|
||||
async for output in engine.generate(
|
||||
spaced_input_generator(),
|
||||
sampling_params,
|
||||
request_id,
|
||||
):
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
outputs_per_chunk[current_chunk] += 1
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify the final output is marked as finished
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify we received outputs from multiple chunks
|
||||
# (with spaced inputs, we should see outputs distributed across chunks)
|
||||
chunks_with_outputs = sum(1 for c in outputs_per_chunk if c > 0)
|
||||
assert chunks_with_outputs >= 1, "Should have outputs from at least one chunk"
|
||||
|
||||
print(f"Spaced test generated: {full_text}")
|
||||
print(f"Outputs per chunk: {outputs_per_chunk}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_output_equivalence(engine: AsyncLLM):
|
||||
"""Test that bunched and spaced inputs produce equivalent outputs.
|
||||
|
||||
When the same prompts are provided either bunched or spaced,
|
||||
the final concatenated output should be the same (with deterministic
|
||||
sampling).
|
||||
"""
|
||||
prompts = ["Hello, my name is", " Bob and I work", " at Anthropic"]
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
# Test bunched inputs
|
||||
async def bunched_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
|
||||
_, bunched_text = await collect_outputs(
|
||||
engine.generate(bunched_gen(), sampling_params, "equiv_bunched")
|
||||
)
|
||||
|
||||
# Test spaced inputs (same prompts, but with delays)
|
||||
async def spaced_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
_, spaced_text = await collect_outputs(
|
||||
engine.generate(spaced_gen(), sampling_params, "equiv_spaced")
|
||||
)
|
||||
|
||||
# Both should produce the same output since we use temperature=0
|
||||
assert bunched_text == spaced_text, (
|
||||
f"Bunched and spaced should produce same output.\n"
|
||||
f"Bunched: {bunched_text!r}\n"
|
||||
f"Spaced: {spaced_text!r}"
|
||||
)
|
||||
|
||||
print(f"Equivalence test passed. Generated: {bunched_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_cancel_output_stream(engine: AsyncLLM):
|
||||
"""Test that cancelling the output stream aborts the entire session.
|
||||
|
||||
When the consumer cancels iteration over the output generator,
|
||||
the session should be aborted including any queued inputs.
|
||||
"""
|
||||
request_id = "test_cancel_output"
|
||||
sampling_params = get_sampling_params(max_tokens=1000)
|
||||
|
||||
input_completed = asyncio.Event()
|
||||
input_task_cancelled = False
|
||||
|
||||
async def slow_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_task_cancelled
|
||||
try:
|
||||
yield StreamingInput(prompt="Tell me a very long story about")
|
||||
yield StreamingInput(prompt=" a dragon and a knight")
|
||||
|
||||
# This should be cancelled before we get here
|
||||
await asyncio.sleep(10)
|
||||
yield StreamingInput(prompt=" who become friends")
|
||||
input_completed.set()
|
||||
except asyncio.CancelledError:
|
||||
input_task_cancelled = True
|
||||
raise
|
||||
|
||||
outputs_received = 0
|
||||
output_gen = engine.generate(slow_input_generator(), sampling_params, request_id)
|
||||
|
||||
# Collect a few outputs then cancel
|
||||
try:
|
||||
async for output in output_gen:
|
||||
outputs_received += 1
|
||||
if outputs_received >= 5:
|
||||
# Cancel by breaking out of the loop (generator will be GC'd)
|
||||
break
|
||||
finally:
|
||||
# Explicitly close the generator to ensure cleanup
|
||||
await output_gen.aclose()
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify we got some outputs before cancelling
|
||||
assert outputs_received >= 5, "Should have received outputs before cancel"
|
||||
|
||||
# Verify the input task was cancelled
|
||||
assert input_task_cancelled, "Input task should have been cancelled"
|
||||
|
||||
# Verify the session is properly cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after cancel"
|
||||
)
|
||||
|
||||
print(f"Cancel test passed. Received {outputs_received} outputs before cancel")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_close_signals_completion(engine: AsyncLLM):
|
||||
"""Test that closing the input stream signals completion.
|
||||
|
||||
When the input generator finishes (naturally or via return),
|
||||
the session should complete with finished=True on the last output.
|
||||
"""
|
||||
request_id = "test_close_completion"
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
input_generator_finished = False
|
||||
|
||||
async def limited_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_generator_finished
|
||||
yield StreamingInput(prompt="What is 2 + 2? The answer is")
|
||||
# Generator finishes naturally here
|
||||
input_generator_finished = True
|
||||
|
||||
outputs, _ = await collect_outputs(
|
||||
engine.generate(limited_input_generator(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
# Verify the input generator completed
|
||||
assert input_generator_finished, "Input generator should have finished"
|
||||
|
||||
# Verify we got a finished output
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
assert outputs[-1].finished, "Last output should be marked as finished"
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests"
|
||||
)
|
||||
|
||||
print("Close completion test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_abort_queued_inputs(engine: AsyncLLM):
|
||||
"""Test that aborting the session cancels queued inputs.
|
||||
|
||||
When multiple inputs are queued and the session is aborted,
|
||||
all pending inputs should be cancelled.
|
||||
"""
|
||||
request_id = "test_abort_queued"
|
||||
# Use large max_tokens to ensure we have time to queue inputs
|
||||
sampling_params = get_sampling_params(max_tokens=2000)
|
||||
|
||||
inputs_sent = 0
|
||||
input_cancelled = False
|
||||
|
||||
async def many_inputs_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal inputs_sent, input_cancelled
|
||||
try:
|
||||
# Send several inputs to fill the queue
|
||||
for i in range(10):
|
||||
yield StreamingInput(prompt=f" Part {i}: Tell me about the number {i}.")
|
||||
inputs_sent += 1
|
||||
# Small delay to interleave with output processing
|
||||
await asyncio.sleep(0.05)
|
||||
except asyncio.CancelledError:
|
||||
input_cancelled = True
|
||||
raise
|
||||
|
||||
outputs_received = 0
|
||||
output_gen = engine.generate(many_inputs_generator(), sampling_params, request_id)
|
||||
|
||||
try:
|
||||
async for output in output_gen:
|
||||
outputs_received += 1
|
||||
# Cancel after receiving some outputs
|
||||
if outputs_received >= 10:
|
||||
break
|
||||
finally:
|
||||
await output_gen.aclose()
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Verify we received some outputs
|
||||
assert outputs_received >= 10, "Should have received outputs before abort"
|
||||
|
||||
# Verify the input generator was cancelled OR finished naturally
|
||||
# (it might finish naturally if all inputs were sent before cancel)
|
||||
assert input_cancelled or inputs_sent == 10, (
|
||||
f"Input generator should have been cancelled or completed. "
|
||||
f"cancelled={input_cancelled}, inputs_sent={inputs_sent}"
|
||||
)
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after abort"
|
||||
)
|
||||
|
||||
print(
|
||||
f"Abort queued test passed. Sent {inputs_sent} inputs, "
|
||||
f"received {outputs_received} outputs"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_error_propagation(engine: AsyncLLM):
|
||||
"""Test that errors in the input generator are propagated to the caller."""
|
||||
request_id = "test_error_propagation"
|
||||
sampling_params = get_sampling_params(max_tokens=20)
|
||||
|
||||
class InputError(Exception):
|
||||
pass
|
||||
|
||||
async def error_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Start with this")
|
||||
await asyncio.sleep(0.1)
|
||||
raise InputError("Simulated input error")
|
||||
|
||||
# Note: The current implementation catches exceptions and puts them
|
||||
# in the queue, so we should get the error when iterating outputs
|
||||
with pytest.raises(InputError, match="Simulated input error"):
|
||||
async for _ in engine.generate(
|
||||
error_input_generator(), sampling_params, request_id
|
||||
):
|
||||
pass
|
||||
|
||||
# Give time for cleanup
|
||||
await asyncio.sleep(0.3)
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests after error"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_multiple_concurrent_sessions(engine: AsyncLLM):
|
||||
"""Test multiple concurrent streaming input sessions.
|
||||
|
||||
Multiple streaming sessions should be able to run concurrently
|
||||
without interfering with each other.
|
||||
"""
|
||||
num_sessions = 3
|
||||
results: list[tuple[str, str]] = []
|
||||
|
||||
async def run_session(session_id: int) -> tuple[str, str]:
|
||||
request_id = f"test_concurrent_{session_id}"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
prompts = [f"Session {session_id}: Hello", f" world from session {session_id}"]
|
||||
|
||||
async def input_gen() -> AsyncGenerator[StreamingInput, None]:
|
||||
for prompt in prompts:
|
||||
yield StreamingInput(prompt=prompt)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
_, text = await collect_outputs(
|
||||
engine.generate(input_gen(), sampling_params, request_id)
|
||||
)
|
||||
return request_id, text
|
||||
|
||||
# Run sessions concurrently
|
||||
tasks = [asyncio.create_task(run_session(i)) for i in range(num_sessions)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all sessions completed
|
||||
assert len(results) == num_sessions
|
||||
|
||||
for request_id, text in results:
|
||||
assert len(text) > 0, f"Session {request_id} should have generated text"
|
||||
print(f"{request_id}: {text}")
|
||||
|
||||
# Verify cleanup
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_per_chunk_sampling_params(engine: AsyncLLM):
|
||||
"""Test that per-chunk sampling params are respected.
|
||||
|
||||
Each StreamingInput can have its own sampling_params.
|
||||
"""
|
||||
request_id = "test_per_chunk_params"
|
||||
base_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
async def variable_params_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# First chunk with base params
|
||||
yield StreamingInput(prompt="Count to five:", sampling_params=base_params)
|
||||
|
||||
# Second chunk with different max_tokens
|
||||
chunk_params = get_sampling_params(max_tokens=5)
|
||||
yield StreamingInput(
|
||||
prompt=" Now count backwards:", sampling_params=chunk_params
|
||||
)
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(variable_params_generator(), base_params, request_id)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
assert outputs[-1].finished, "Last output should be finished"
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
|
||||
print(f"Per-chunk params test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_empty_generator(engine: AsyncLLM):
|
||||
"""Test behavior when the input generator yields nothing.
|
||||
|
||||
An empty generator should still produce a finished output.
|
||||
"""
|
||||
request_id = "test_empty_generator"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
async def empty_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
# Don't yield anything
|
||||
return
|
||||
yield # Make it a generator
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
async for output in engine.generate(empty_generator(), sampling_params, request_id):
|
||||
outputs.append(output)
|
||||
|
||||
# Should still get a finished marker
|
||||
assert len(outputs) >= 1, "Should receive at least one output"
|
||||
assert outputs[-1].finished, "Should have a finished output"
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_single_chunk(engine: AsyncLLM):
|
||||
"""Test streaming input with a single chunk.
|
||||
|
||||
This is effectively the same as a regular non-streaming request,
|
||||
but using the streaming input API.
|
||||
"""
|
||||
request_id = "test_single_chunk"
|
||||
sampling_params = get_sampling_params(max_tokens=15)
|
||||
|
||||
async def single_chunk_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="What color is the sky? The sky is")
|
||||
|
||||
outputs, full_text = await collect_outputs(
|
||||
engine.generate(single_chunk_generator(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
assert len(outputs) > 0
|
||||
assert outputs[-1].finished
|
||||
assert "blue" in full_text.lower() or len(full_text) > 0
|
||||
|
||||
print(f"Single chunk test generated: {full_text}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_reuse_request_id(engine: AsyncLLM):
|
||||
"""Test that request IDs can be reused after a session completes."""
|
||||
request_id = "test_reuse_id"
|
||||
sampling_params = get_sampling_params(max_tokens=5)
|
||||
|
||||
# First session
|
||||
async def gen1() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="First session")
|
||||
|
||||
_, text1 = await collect_outputs(
|
||||
engine.generate(gen1(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
# Second session with same ID
|
||||
async def gen2() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Second session")
|
||||
|
||||
_, text2 = await collect_outputs(
|
||||
engine.generate(gen2(), sampling_params, request_id)
|
||||
)
|
||||
|
||||
assert len(text1) > 0
|
||||
assert len(text2) > 0
|
||||
assert not engine.output_processor.has_unfinished_requests()
|
||||
|
||||
print(f"Reuse ID test: session 1: {text1}, session 2: {text2}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_validation_errors(engine: AsyncLLM):
|
||||
"""Test that invalid configurations raise appropriate errors."""
|
||||
|
||||
async def dummy_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="test")
|
||||
|
||||
# Test n > 1 is rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_n2 = SamplingParams(max_tokens=10, n=2)
|
||||
async for _ in engine.generate(dummy_generator(), params_n2, "test_n2"):
|
||||
pass
|
||||
|
||||
# Test FINAL_ONLY is rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_final = SamplingParams(
|
||||
max_tokens=10, output_kind=RequestOutputKind.FINAL_ONLY
|
||||
)
|
||||
async for _ in engine.generate(dummy_generator(), params_final, "test_final"):
|
||||
pass
|
||||
|
||||
# Test stop strings are rejected
|
||||
with pytest.raises(ValueError, match="Input streaming not currently supported"):
|
||||
params_stop = SamplingParams(max_tokens=10, stop=["stop"])
|
||||
async for _ in engine.generate(dummy_generator(), params_stop, "test_stop"):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio(loop_scope="module")
|
||||
async def test_streaming_input_delayed_generator_exit(engine: AsyncLLM):
|
||||
"""Test that output generator exits when input generator closes after outputs.
|
||||
|
||||
This tests the case where:
|
||||
1. Multiple inputs are sent and fully processed
|
||||
2. The engine has finished
|
||||
3. The input generator doesn't exit until after the engine finishes
|
||||
4. The output generator should exit properly once the input generator exits
|
||||
"""
|
||||
request_id = "test_delayed_exit"
|
||||
sampling_params = get_sampling_params(max_tokens=10)
|
||||
|
||||
engine_finished_event = asyncio.Event()
|
||||
input_generator_exited = False
|
||||
finish_count = 0
|
||||
|
||||
async def delayed_exit_input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
nonlocal input_generator_exited
|
||||
# Send all inputs immediately
|
||||
yield StreamingInput(prompt="Hello, my name is")
|
||||
yield StreamingInput(prompt=" Alice")
|
||||
|
||||
# Wait until the engine has finished generating before exiting
|
||||
await engine_finished_event.wait()
|
||||
|
||||
# Add a small delay to ensure we're testing the "delayed exit" case
|
||||
await asyncio.sleep(0.1)
|
||||
input_generator_exited = True
|
||||
|
||||
outputs: list[RequestOutput] = []
|
||||
full_text = ""
|
||||
|
||||
async for output in engine.generate(
|
||||
delayed_exit_input_generator(), sampling_params, request_id
|
||||
):
|
||||
outputs.append(output)
|
||||
if output.outputs and output.outputs[0].text:
|
||||
full_text += output.outputs[0].text
|
||||
|
||||
# Signal when the engine finishes both input chunks (each gets a finish_reason)
|
||||
# Note: output.finished will be False while input stream is open
|
||||
if output.outputs and output.outputs[0].finish_reason is not None:
|
||||
finish_count += 1
|
||||
if finish_count == 2:
|
||||
engine_finished_event.set()
|
||||
|
||||
# Verify the input generator exited properly
|
||||
assert input_generator_exited, (
|
||||
"Input generator should have exited after engine finished"
|
||||
)
|
||||
|
||||
# Verify we got outputs
|
||||
assert len(outputs) > 0, "Should have received outputs"
|
||||
|
||||
# Verify we generated some text
|
||||
assert len(full_text) > 0, "Should have generated text"
|
||||
|
||||
# Verify the session is cleaned up
|
||||
assert not engine.output_processor.has_unfinished_requests(), (
|
||||
"Should have no unfinished requests"
|
||||
)
|
||||
|
||||
print(f"Delayed exit test passed. Generated: {full_text}")
|
||||
@@ -5,8 +5,14 @@ import pytest
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
||||
from vllm.multimodal import MultiModalUUIDDict
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
DeviceConfig,
|
||||
ModelConfig,
|
||||
MultiModalConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.multimodal import MultiModalRegistry, MultiModalUUIDDict
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine.input_processor import InputProcessor
|
||||
|
||||
@@ -15,26 +21,55 @@ stop_pil_image = ImageAsset("stop_sign").pil_image
|
||||
baby_reading_np_ndarrays = VideoAsset("baby_reading").np_ndarrays
|
||||
|
||||
|
||||
def _build_input_processor(
|
||||
*, mm_cache_gb: float = 4.0, enable_prefix_caching: bool = True
|
||||
def _mock_input_processor(
|
||||
monkeypatch, *, mm_cache_gb: float = 4.0, enable_prefix_caching: bool = True
|
||||
) -> InputProcessor:
|
||||
"""
|
||||
Create a Processor instance with minimal configuration suitable for unit
|
||||
tests without accessing external resources.
|
||||
"""
|
||||
monkeypatch.setattr(
|
||||
ModelConfig, "try_get_generation_config", lambda self: {}, raising=True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ModelConfig, "__post_init__", lambda self, *args: None, raising=True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
ModelConfig,
|
||||
"verify_with_parallel_config",
|
||||
lambda self, parallel_config: None,
|
||||
raising=True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
MultiModalRegistry,
|
||||
"processor_cache_from_config",
|
||||
lambda self, vllm_config: None,
|
||||
raising=True,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(VllmConfig, "__post_init__", lambda self: None, raising=True)
|
||||
|
||||
model_config = ModelConfig(
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
tokenizer="dummy",
|
||||
skip_tokenizer_init=True,
|
||||
max_model_len=128,
|
||||
mm_processor_cache_gb=mm_cache_gb,
|
||||
generation_config="vllm",
|
||||
)
|
||||
model_config.runner_type = "generate"
|
||||
model_config.multimodal_config = MultiModalConfig(mm_processor_cache_gb=mm_cache_gb)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=CacheConfig(enable_prefix_caching=enable_prefix_caching),
|
||||
device_config=DeviceConfig(device="cpu"),
|
||||
)
|
||||
|
||||
return InputProcessor(vllm_config)
|
||||
|
||||
|
||||
def test_multi_modal_uuids_length_mismatch_raises():
|
||||
input_processor = _build_input_processor()
|
||||
def test_multi_modal_uuids_length_mismatch_raises(monkeypatch):
|
||||
input_processor = _mock_input_processor(monkeypatch)
|
||||
|
||||
prompt = {
|
||||
"prompt": "USER: <image>\nDescribe\nASSISTANT:",
|
||||
@@ -43,7 +78,7 @@ def test_multi_modal_uuids_length_mismatch_raises():
|
||||
"multi_modal_uuids": {"image": ["hash_cherry"]},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="must have same length as"):
|
||||
with pytest.raises(ValueError, match="must have same length as data"):
|
||||
input_processor.process_inputs(
|
||||
request_id="req-1",
|
||||
prompt=prompt, # type: ignore[arg-type]
|
||||
@@ -51,21 +86,21 @@ def test_multi_modal_uuids_length_mismatch_raises():
|
||||
)
|
||||
|
||||
|
||||
def test_multi_modal_uuids_missing_modality_raises():
|
||||
input_processor = _build_input_processor()
|
||||
def test_multi_modal_uuids_missing_modality_raises(monkeypatch):
|
||||
input_processor = _mock_input_processor(monkeypatch)
|
||||
|
||||
prompt = {
|
||||
"prompt": "USER: <image><video>\nDescribe\nASSISTANT:",
|
||||
# Two modalities provided in data
|
||||
"multi_modal_data": {
|
||||
"image": [cherry_pil_image],
|
||||
"video": None,
|
||||
"video": [baby_reading_np_ndarrays],
|
||||
},
|
||||
# Only image uuids provided; video missing should raise
|
||||
"multi_modal_uuids": {"image": ["hash_cherry"]},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="is empty but .* is missing"):
|
||||
with pytest.raises(ValueError, match="must be provided if multi_modal_data"):
|
||||
input_processor.process_inputs(
|
||||
request_id="req-2",
|
||||
prompt=prompt, # type: ignore[arg-type]
|
||||
@@ -84,7 +119,8 @@ def test_multi_modal_uuids_missing_modality_raises():
|
||||
def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
|
||||
):
|
||||
input_processor = _build_input_processor(
|
||||
input_processor = _mock_input_processor(
|
||||
monkeypatch,
|
||||
mm_cache_gb=mm_cache_gb,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
@@ -127,8 +163,8 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
# When both processor cache is 0 and prefix caching disabled, the
|
||||
# processor builds overrides from request id instead of using user UUIDs.
|
||||
input_processor = _build_input_processor(
|
||||
mm_cache_gb=0.0, enable_prefix_caching=False
|
||||
input_processor = _mock_input_processor(
|
||||
monkeypatch, mm_cache_gb=0.0, enable_prefix_caching=False
|
||||
)
|
||||
|
||||
captured: dict[str, MultiModalUUIDDict] = {}
|
||||
@@ -144,12 +180,12 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
)
|
||||
|
||||
request_id = "req-42"
|
||||
mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}
|
||||
mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": "hash_video"}
|
||||
prompt = {
|
||||
"prompt": "USER: <image><image><video>\nDescribe\nASSISTANT:",
|
||||
"multi_modal_data": {
|
||||
"image": [cherry_pil_image, stop_pil_image],
|
||||
"video": [baby_reading_np_ndarrays],
|
||||
"video": baby_reading_np_ndarrays,
|
||||
},
|
||||
"multi_modal_uuids": mm_uuids,
|
||||
}
|
||||
@@ -161,15 +197,16 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
|
||||
)
|
||||
|
||||
# Expect request-id-based overrides are passed through
|
||||
mm_uuids = captured["mm_uuids"]
|
||||
assert set(mm_uuids.keys()) == {"image", "video"}
|
||||
assert len(mm_uuids["image"]) == 2
|
||||
assert len(mm_uuids["video"]) == 1
|
||||
assert captured["mm_uuids"]["image"][0].startswith(
|
||||
f"{request_id}-image-"
|
||||
) and captured["mm_uuids"]["image"][0].endswith("-0")
|
||||
assert captured["mm_uuids"]["image"][1].startswith(
|
||||
f"{request_id}-image-"
|
||||
) and captured["mm_uuids"]["image"][1].endswith("-1")
|
||||
assert captured["mm_uuids"]["video"][0].startswith(
|
||||
f"{request_id}-video-"
|
||||
) and captured["mm_uuids"]["video"][0].endswith("-0")
|
||||
assert mm_uuids["image"][0].startswith(f"{request_id}-image-") and mm_uuids[
|
||||
"image"
|
||||
][0].endswith("-0")
|
||||
assert mm_uuids["image"][1].startswith(f"{request_id}-image-") and mm_uuids[
|
||||
"image"
|
||||
][1].endswith("-1")
|
||||
assert mm_uuids["video"][0].startswith(f"{request_id}-video-") and mm_uuids[
|
||||
"video"
|
||||
][0].endswith("-0")
|
||||
|
||||
@@ -70,28 +70,15 @@ async def test_background_cancel(client: openai.AsyncOpenAI):
|
||||
assert response.status == "queued"
|
||||
|
||||
# Cancel the response before it is completed.
|
||||
# Poll until the response is no longer queued (started processing) or timeout
|
||||
loop = asyncio.get_running_loop()
|
||||
start_time = loop.time()
|
||||
max_wait_seconds = 5.0
|
||||
poll_interval = 0.1
|
||||
while loop.time() - start_time < max_wait_seconds:
|
||||
response = await client.responses.retrieve(response.id)
|
||||
if response.status != "queued":
|
||||
# Started processing or completed - try to cancel
|
||||
break
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
# FIXME: This test can be flaky.
|
||||
await asyncio.sleep(0.5)
|
||||
response = await client.responses.cancel(response.id)
|
||||
assert response.status == "cancelled"
|
||||
|
||||
# Make sure the response status remains unchanged after some time.
|
||||
max_retries = 10
|
||||
for _ in range(max_retries):
|
||||
await asyncio.sleep(0.5)
|
||||
response = await client.responses.retrieve(response.id)
|
||||
# Verify status is still cancelled
|
||||
assert response.status == "cancelled"
|
||||
# Make sure the response status remains unchanged.
|
||||
await asyncio.sleep(5)
|
||||
response = await client.responses.retrieve(response.id)
|
||||
assert response.status == "cancelled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -34,11 +34,18 @@ else
|
||||
KV_CONFIG_HETERO_LAYOUT=''
|
||||
fi
|
||||
|
||||
CROSS_LAYERS_BLOCKS=${CROSS_LAYERS_BLOCKS:-"False"} # Default to non cross layers
|
||||
if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
|
||||
KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"cross_layers_blocks": "True"}'
|
||||
else
|
||||
KV_EXTRA_CONFIG=''
|
||||
fi
|
||||
|
||||
# Build the kv-transfer-config once
|
||||
if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
|
||||
else
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
|
||||
fi
|
||||
|
||||
# Models to run
|
||||
|
||||
@@ -86,7 +86,7 @@ class DecodeBenchTestRunner:
|
||||
self._block_hasher = get_request_block_hasher(block_size, sha256)
|
||||
|
||||
self._dummy_ctx: ForwardContext = ForwardContext(
|
||||
no_compile_layers={}, attn_metadata={}, virtual_engine=0, slot_mapping={}
|
||||
no_compile_layers={}, attn_metadata={}, virtual_engine=0
|
||||
)
|
||||
|
||||
def new_request(self, token_ids: list[int]) -> Request:
|
||||
|
||||
@@ -18,8 +18,12 @@ import ray
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import KVTransferConfig
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
|
||||
from vllm.config import KVTransferConfig, set_current_vllm_config
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import (
|
||||
KVOutputAggregator,
|
||||
TpKVTopology,
|
||||
get_current_attn_backend,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
|
||||
@@ -48,8 +52,11 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.output_processor import OutputProcessor
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
|
||||
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
|
||||
from vllm.v1.worker.utils import AttentionGroup
|
||||
|
||||
from .utils import create_request, create_scheduler, create_vllm_config
|
||||
|
||||
@@ -366,6 +373,7 @@ def test_kv_transfer_handshake(dist_init):
|
||||
|
||||
# Decode connector will be able to create handshake with the prefill connector.
|
||||
decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
# Here we are testing the retrieval of NIXLAgentMetadata.
|
||||
# Knowing the implementation detail, we override the add_remote_agent
|
||||
@@ -402,6 +410,23 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
|
||||
self.kv_cache_layout = kv_cache_layout
|
||||
# Mock register_kv_caches attribute needed for tests that do not call it.
|
||||
self.src_xfer_handles_by_block_size = {self.block_size: 1}
|
||||
test_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
self.kv_topo = TpKVTopology(
|
||||
tp_rank=self.tp_rank,
|
||||
engine_id=self.engine_id,
|
||||
remote_tp_size=self._tp_size, # shared state
|
||||
remote_block_size=self._block_size, # shared state
|
||||
is_mla=self.use_mla,
|
||||
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=self.attn_backend,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
self.compat_hash = compute_nixl_compatibility_hash(
|
||||
self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
|
||||
)
|
||||
|
||||
def _nixl_handshake(
|
||||
self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
|
||||
@@ -523,7 +548,6 @@ class TestNixlHandshake:
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
_before_load = time.perf_counter()
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
@@ -594,7 +618,6 @@ class TestNixlHandshake:
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
_before_load = time.perf_counter()
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
@@ -821,7 +844,6 @@ class TestNixlHandshake:
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
_before_load = time.perf_counter()
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
@@ -984,7 +1006,6 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -1370,6 +1391,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
|
||||
),
|
||||
),
|
||||
"TRITON_ATTN",
|
||||
"FLASHINFER",
|
||||
],
|
||||
)
|
||||
def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
@@ -1386,6 +1408,11 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
|
||||
vllm_config = create_vllm_config(attention_backend=attn_backend)
|
||||
|
||||
# Enable cross layers blocks
|
||||
vllm_config.kv_transfer_config.kv_connector_extra_config[
|
||||
"enable_cross_layers_blocks"
|
||||
] = True
|
||||
|
||||
# Import the appropriate backend based on the parameter
|
||||
if attn_backend == "FLASH_ATTN":
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
@@ -1395,49 +1422,11 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
|
||||
|
||||
backend_cls = RocmAttentionBackend
|
||||
else: # TRITON_ATTN
|
||||
else: # TRITON
|
||||
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
|
||||
|
||||
backend_cls = TritonAttentionBackend
|
||||
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
|
||||
# Store tensor info for validation
|
||||
|
||||
test_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
|
||||
|
||||
if is_blocks_first:
|
||||
expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
|
||||
expected_base_addrs = [
|
||||
shared_tensor.data_ptr(),
|
||||
unique_tensor.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 2
|
||||
else:
|
||||
expected_tensor_size = (
|
||||
shared_tensor[0].element_size() * shared_tensor[0].numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor[0].data_ptr(),
|
||||
shared_tensor[1].data_ptr(),
|
||||
unique_tensor[0].data_ptr(),
|
||||
unique_tensor[1].data_ptr(),
|
||||
]
|
||||
expected_num_entries = 4
|
||||
|
||||
nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
|
||||
with (
|
||||
patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
|
||||
@@ -1466,6 +1455,107 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
# Reassure the shutdown() check that the thread is terminated
|
||||
mock_thread.return_value.is_alive.return_value = False
|
||||
|
||||
expected_tensor_size: int
|
||||
expected_base_addrs: list[int]
|
||||
expected_num_entries: int
|
||||
kv_caches: dict[str, torch.Tensor]
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_layers = 32
|
||||
block_size = 16
|
||||
num_blocks = 8
|
||||
kv_cache_spec = AttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(
|
||||
size=kv_cache_spec.page_size_bytes * num_blocks,
|
||||
shared_by=["dummy-layer"],
|
||||
)
|
||||
for i in range(num_layers)
|
||||
],
|
||||
# allocate_uniform_kv_caches does not use this
|
||||
kv_cache_groups=[],
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
_, cross_layers_kv_cache, _ = (
|
||||
KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
|
||||
kv_cache_config=kv_cache_config,
|
||||
attn_groups=[
|
||||
[
|
||||
AttentionGroup(
|
||||
backend=backend_cls,
|
||||
layer_names=[],
|
||||
kv_cache_spec=kv_cache_spec,
|
||||
kv_cache_group_id=0,
|
||||
)
|
||||
]
|
||||
],
|
||||
cache_dtype=torch.bfloat16,
|
||||
device=torch.cuda.current_device(),
|
||||
kernel_block_sizes=[block_size],
|
||||
)
|
||||
)
|
||||
# Store tensor info for validation
|
||||
expected_tensor_size = (
|
||||
cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
cross_layers_kv_cache.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 1
|
||||
|
||||
expected_blocks_count = 8
|
||||
|
||||
kv_caches = {"all-layers": cross_layers_kv_cache}
|
||||
|
||||
else:
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
|
||||
# Store tensor info for validation
|
||||
|
||||
test_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
|
||||
|
||||
if is_blocks_first:
|
||||
expected_tensor_size = (
|
||||
shared_tensor.element_size() * shared_tensor.numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor.data_ptr(),
|
||||
unique_tensor.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 2
|
||||
else:
|
||||
expected_tensor_size = (
|
||||
shared_tensor[0].element_size() * shared_tensor[0].numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor[0].data_ptr(),
|
||||
shared_tensor[1].data_ptr(),
|
||||
unique_tensor[0].data_ptr(),
|
||||
unique_tensor[1].data_ptr(),
|
||||
]
|
||||
expected_num_entries = 4
|
||||
expected_blocks_count = 8
|
||||
|
||||
# Execute register_kv_caches
|
||||
connector.register_kv_caches(kv_caches)
|
||||
|
||||
@@ -1489,16 +1579,19 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
|
||||
|
||||
# Validate blocks_data structure and size
|
||||
expected_blocks_count = 8
|
||||
assert len(blocks_data) == expected_blocks_count, (
|
||||
f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
|
||||
)
|
||||
|
||||
num_blocks = 2
|
||||
if is_blocks_first:
|
||||
expected_block_len = expected_tensor_size // num_blocks // 2
|
||||
else:
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_blocks = 8
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
else:
|
||||
num_blocks = 2
|
||||
if is_blocks_first:
|
||||
expected_block_len = expected_tensor_size // num_blocks // 2
|
||||
else:
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
|
||||
for i, block_entry in enumerate(blocks_data):
|
||||
block_start_addr, block_len, tp_rank = block_entry
|
||||
@@ -1674,7 +1767,6 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -1825,7 +1917,6 @@ def test_transfer_failure_logging(
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
|
||||
# Capture logs from the nixl_connector logger specifically
|
||||
@@ -1926,7 +2017,6 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -1977,7 +2067,6 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
)
|
||||
connector.start_load_kv(dummy_ctx)
|
||||
|
||||
@@ -2049,6 +2138,17 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
|
||||
decode_worker = decode_connector.connector_worker
|
||||
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
remote_config_params: dict[str, Any] = {
|
||||
"model": "facebook/opt-125m",
|
||||
@@ -2071,7 +2171,9 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
)
|
||||
remote_hash = compute_nixl_compatibility_hash(
|
||||
remote_vllm_config, decode_worker.backend_name
|
||||
remote_vllm_config,
|
||||
decode_worker.backend_name,
|
||||
decode_worker.kv_topo.cross_layers_blocks,
|
||||
)
|
||||
|
||||
prefill_block_size = config_overrides.get("block_size", 16)
|
||||
@@ -2150,6 +2252,27 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
|
||||
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
|
||||
decode_worker = decode_connector.connector_worker
|
||||
|
||||
backend = get_current_attn_backend(local_vllm_config)
|
||||
test_shape = backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
decode_worker.kv_topo = TpKVTopology(
|
||||
tp_rank=decode_worker.tp_rank,
|
||||
engine_id=decode_worker.engine_id,
|
||||
remote_tp_size=decode_worker._tp_size, # shared state
|
||||
remote_block_size=decode_worker._block_size, # shared state
|
||||
is_mla=decode_worker.use_mla,
|
||||
total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=backend,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
decode_worker.compat_hash = compute_nixl_compatibility_hash(
|
||||
decode_worker.vllm_config,
|
||||
decode_worker.backend_name,
|
||||
decode_worker.kv_topo.cross_layers_blocks,
|
||||
)
|
||||
|
||||
if error_scenario == "handshake_decode_error":
|
||||
msg_bytes = b"this is not valid msgpack data"
|
||||
elif error_scenario == "handshake_validation_error":
|
||||
|
||||
@@ -209,10 +209,7 @@ class RequestRunner:
|
||||
self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
|
||||
|
||||
self._dummy_ctx: ForwardContext = ForwardContext(
|
||||
no_compile_layers={},
|
||||
attn_metadata={},
|
||||
virtual_engine=0,
|
||||
slot_mapping={},
|
||||
no_compile_layers={}, attn_metadata={}, virtual_engine=0
|
||||
)
|
||||
|
||||
def new_request(self, token_ids: list[int]):
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.utils import create_new_process_for_each_test, set_random_seed
|
||||
from tests.utils import create_new_process_for_each_test
|
||||
from tests.v1.logits_processors.utils import (
|
||||
DUMMY_LOGITPROC_ARG,
|
||||
DUMMY_LOGITPROC_FQCN,
|
||||
@@ -134,7 +135,7 @@ def test_custom_logitsprocs(monkeypatch, logitproc_source: CustomLogitprocSource
|
||||
|
||||
# Test that logitproc info is passed to workers
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
|
||||
set_random_seed(40)
|
||||
random.seed(40)
|
||||
|
||||
# Choose LLM args based on logitproc source
|
||||
if logitproc_source == CustomLogitprocSource.LOGITPROC_SOURCE_NONE:
|
||||
@@ -193,7 +194,7 @@ def test_custom_logitsprocs_req(monkeypatch):
|
||||
|
||||
# Test that logitproc info is passed to workers
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1")
|
||||
set_random_seed(40)
|
||||
random.seed(40)
|
||||
_run_test(
|
||||
{"logits_processors": [WrappedPerReqLogitsProcessor]}, logitproc_loaded=True
|
||||
)
|
||||
@@ -236,7 +237,7 @@ def test_rejects_custom_logitsprocs(
|
||||
logitproc from
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
set_random_seed(40)
|
||||
random.seed(40)
|
||||
|
||||
test_params: dict[str, dict[str, Any]] = {
|
||||
"pooling": {
|
||||
|
||||
@@ -48,11 +48,7 @@ def test_topk_impl_equivalence():
|
||||
assert torch.allclose(result1, result2)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="FlashInfer top-k/top-p renorm comparison fails; "
|
||||
"needs investigation of tolerance threshold or "
|
||||
"interface differences between Python and FlashInfer implementations"
|
||||
)
|
||||
@pytest.mark.skip(reason="FIXME: This test is failing right now.")
|
||||
def test_flashinfer_sampler():
|
||||
"""
|
||||
This test verifies that the FlashInfer top-k and top-p sampling
|
||||
|
||||
@@ -9,7 +9,6 @@ import torch
|
||||
from tests.v1.attention.utils import (
|
||||
create_standard_kv_cache_spec,
|
||||
create_vllm_config,
|
||||
try_backend_includes_kv_cache_update,
|
||||
try_get_attention_backend,
|
||||
)
|
||||
from vllm.config import ParallelConfig, SpeculativeConfig
|
||||
@@ -121,14 +120,6 @@ def forward_attention(
|
||||
key = k.view(-1, num_kv_heads, dim_per_head)
|
||||
value = v.view(-1, num_kv_heads, dim_per_head)
|
||||
output = torch.empty_like(query)
|
||||
if not try_backend_includes_kv_cache_update(backend):
|
||||
instance.do_kv_cache_update(
|
||||
layer=layer,
|
||||
key=key,
|
||||
value=value,
|
||||
kv_cache=kv_cache,
|
||||
slot_mapping=attn_metadata.slot_mapping,
|
||||
)
|
||||
return instance.forward(
|
||||
layer=layer,
|
||||
query=query,
|
||||
|
||||
@@ -1,171 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.outputs import RequestOutput
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.v1.engine.async_llm import AsyncLLM, StreamingInput
|
||||
from vllm.v1.engine.output_processor import RequestOutputCollector
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_llm():
|
||||
"""Create a mock AsyncLLM with mocked dependencies."""
|
||||
# Create a minimal mock without initializing the full engine
|
||||
llm = MagicMock(spec=AsyncLLM)
|
||||
|
||||
# Mock the essential attributes
|
||||
llm.vllm_config = MagicMock()
|
||||
llm.vllm_config.cache_config.kv_sharing_fast_prefill = False
|
||||
llm.model_config = MagicMock()
|
||||
llm.model_config.max_model_len = 2048
|
||||
llm.log_requests = False
|
||||
llm.errored = False
|
||||
llm._pause_cond = asyncio.Condition()
|
||||
llm._paused = False
|
||||
|
||||
# Mock methods
|
||||
llm._run_output_handler = MagicMock()
|
||||
llm.abort = AsyncMock()
|
||||
|
||||
# Use the real generate method from AsyncLLM
|
||||
llm.generate = AsyncLLM.generate.__get__(llm, AsyncLLM)
|
||||
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_normal_flow(mock_async_llm):
|
||||
"""Test normal generation flow with streaming requests."""
|
||||
request_id = "test_request"
|
||||
prompt = "Tell me about Paris"
|
||||
sampling_params = SamplingParams(max_tokens=10)
|
||||
|
||||
# Create a mock queue with outputs
|
||||
queue = RequestOutputCollector(RequestOutputKind.FINAL_ONLY, request_id)
|
||||
output1 = RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="Tell me about Paris",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[],
|
||||
finished=False,
|
||||
)
|
||||
output2 = RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="Tell me about Paris",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[],
|
||||
finished=True,
|
||||
)
|
||||
|
||||
# Feed outputs to queue as they're consumed to avoid aggregation
|
||||
async def feed_outputs():
|
||||
queue.put(output1)
|
||||
await asyncio.sleep(1) # Let first output be consumed
|
||||
queue.put(output2)
|
||||
|
||||
asyncio.create_task(feed_outputs()) # noqa
|
||||
|
||||
# Mock add_request to return the queue
|
||||
async def mock_add_request(*args, **kwargs):
|
||||
return queue
|
||||
|
||||
mock_async_llm.add_request = mock_add_request
|
||||
|
||||
# Collect outputs from generate
|
||||
outputs = []
|
||||
async for output in mock_async_llm.generate(
|
||||
prompt=prompt,
|
||||
sampling_params=sampling_params,
|
||||
request_id=request_id,
|
||||
):
|
||||
outputs.append(output)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0].finished is False
|
||||
assert outputs[1].finished is True
|
||||
|
||||
|
||||
def make_output(request_id: str, finished: bool) -> RequestOutput:
|
||||
"""Helper to create a RequestOutput."""
|
||||
return RequestOutput(
|
||||
request_id=request_id,
|
||||
prompt="test",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
prompt_logprobs=None,
|
||||
outputs=[],
|
||||
finished=finished,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_with_async_generator():
|
||||
"""Test generate with an async input generator.
|
||||
|
||||
With the new streaming input API, completion is signaled by finishing
|
||||
the input generator (not via a resumable flag). Each input chunk
|
||||
produces intermediate outputs, and the final output has finished=True.
|
||||
"""
|
||||
request_id = "test"
|
||||
sampling_params = SamplingParams(max_tokens=10)
|
||||
|
||||
llm = MagicMock(spec=AsyncLLM)
|
||||
llm.vllm_config = MagicMock()
|
||||
llm.vllm_config.cache_config.kv_sharing_fast_prefill = False
|
||||
llm.model_config = MagicMock()
|
||||
llm.model_config.max_model_len = 2048
|
||||
llm.log_requests = False
|
||||
llm.errored = False
|
||||
llm._pause_cond = asyncio.Condition()
|
||||
llm._paused = False
|
||||
llm._run_output_handler = MagicMock()
|
||||
llm.abort = AsyncMock()
|
||||
|
||||
# Bind the real generate method
|
||||
llm.generate = AsyncLLM.generate.__get__(llm, AsyncLLM)
|
||||
|
||||
# Track inputs processed
|
||||
inputs_received = []
|
||||
queue = RequestOutputCollector(RequestOutputKind.DELTA, request_id)
|
||||
|
||||
async def mock_add_request(req_id, prompt, params, *args, **kwargs):
|
||||
# When prompt is an AsyncGenerator, process streaming inputs
|
||||
if isinstance(prompt, AsyncGenerator):
|
||||
# Process inputs in background, produce outputs
|
||||
async def handle_stream():
|
||||
async for input_chunk in prompt:
|
||||
inputs_received.append(input_chunk.prompt)
|
||||
# Each input produces an intermediate output
|
||||
queue.put(make_output(req_id, finished=False))
|
||||
await asyncio.sleep(0.01)
|
||||
# Final output when stream ends
|
||||
queue.put(make_output(req_id, finished=True))
|
||||
|
||||
asyncio.create_task(handle_stream())
|
||||
return queue
|
||||
return queue
|
||||
|
||||
llm.add_request = mock_add_request
|
||||
|
||||
async def input_generator() -> AsyncGenerator[StreamingInput, None]:
|
||||
yield StreamingInput(prompt="Hello", sampling_params=sampling_params)
|
||||
yield StreamingInput(prompt=" world", sampling_params=sampling_params)
|
||||
|
||||
outputs = []
|
||||
async for output in llm.generate(input_generator(), sampling_params, request_id):
|
||||
outputs.append(output)
|
||||
|
||||
# Two intermediate outputs + one final output
|
||||
assert len(outputs) == 3
|
||||
assert outputs[0].finished is False
|
||||
assert outputs[1].finished is False
|
||||
assert outputs[2].finished is True
|
||||
# Both inputs were processed
|
||||
assert inputs_received == ["Hello", " world"]
|
||||
@@ -1,210 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Unit tests for GPUModelRunner._update_streaming_request function."""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalKwargsItem,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
|
||||
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_runner_with_input_batch():
|
||||
"""Create a mock GPUModelRunner with a real InputBatch for e2e testing."""
|
||||
|
||||
runner = Mock(spec=GPUModelRunner)
|
||||
runner.uses_mrope = False
|
||||
runner.requests = {}
|
||||
runner.max_num_reqs = 10
|
||||
runner.max_model_len = 1024
|
||||
|
||||
# Create a real InputBatch for e2e testing
|
||||
runner.input_batch = InputBatch(
|
||||
max_num_reqs=10,
|
||||
max_model_len=1024,
|
||||
max_num_batched_tokens=1024,
|
||||
device="cpu",
|
||||
pin_memory=False,
|
||||
vocab_size=32000,
|
||||
block_sizes=[16],
|
||||
kernel_block_sizes=[16],
|
||||
is_spec_decode=False,
|
||||
logitsprocs=None,
|
||||
is_pooling_model=False,
|
||||
)
|
||||
return runner
|
||||
|
||||
|
||||
def test_e2e_streaming_request_update_basic_flow(mock_model_runner_with_input_batch):
|
||||
"""Test that streaming session are updated correctly.
|
||||
|
||||
This test validates that when a streaming session is updated with new prompt tokens:
|
||||
1. The request is removed from InputBatch before updating (avoids duplication)
|
||||
2. Request state fields are updated correctly
|
||||
3. output_token_ids is cleared (intermediate outputs are now in prompt_token_ids)
|
||||
"""
|
||||
runner = mock_model_runner_with_input_batch
|
||||
req_id = "streaming_req_0"
|
||||
|
||||
# Step 1: Create initial request state with some computed tokens
|
||||
initial_req_state = CachedRequestState(
|
||||
req_id=req_id,
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
mm_features=[],
|
||||
sampling_params=SamplingParams(temperature=0.5),
|
||||
pooling_params=None,
|
||||
generator=None,
|
||||
block_ids=([0],),
|
||||
num_computed_tokens=3,
|
||||
output_token_ids=[10, 11], # Generated 2 tokens
|
||||
)
|
||||
runner.requests[req_id] = initial_req_state
|
||||
|
||||
# Add request to InputBatch
|
||||
runner.input_batch.add_request(initial_req_state)
|
||||
assert req_id in runner.input_batch.req_id_to_index
|
||||
|
||||
# Step 2: Create new request data with extended prompt
|
||||
# The scheduler has already set prompt_token_ids to the full sequence
|
||||
# (original prompt + intermediate outputs + new prompt)
|
||||
new_req_data = Mock()
|
||||
new_req_data.prompt_token_ids = [
|
||||
1,
|
||||
2,
|
||||
3,
|
||||
10,
|
||||
4,
|
||||
5,
|
||||
] # Full sequence with intermediate output (10)
|
||||
new_req_data.mm_features = []
|
||||
new_req_data.prompt_embeds = None
|
||||
new_req_data.sampling_params = SamplingParams(temperature=0.8, max_tokens=50)
|
||||
new_req_data.pooling_params = None
|
||||
new_req_data.block_ids = ([0, 1],)
|
||||
new_req_data.num_computed_tokens = 4 # 3 original prompt + 1 intermediate output
|
||||
|
||||
# Step 3: Update the request
|
||||
updated_req_state = GPUModelRunner._update_streaming_request(
|
||||
runner, req_id, new_req_data
|
||||
)
|
||||
|
||||
# Step 4: Verify the request state was updated correctly
|
||||
assert updated_req_state.prompt_token_ids == [1, 2, 3, 10, 4, 5]
|
||||
assert updated_req_state.num_computed_tokens == 4
|
||||
assert updated_req_state.sampling_params.temperature == 0.8
|
||||
assert updated_req_state.sampling_params.max_tokens == 50
|
||||
assert updated_req_state.block_ids == ([0, 1],)
|
||||
|
||||
# Verify output_token_ids were cleared
|
||||
# (intermediate outputs are now in prompt_token_ids)
|
||||
assert updated_req_state.output_token_ids == []
|
||||
|
||||
# Verify the same object is returned
|
||||
assert runner.requests[req_id] is updated_req_state
|
||||
|
||||
# Verify request was removed from InputBatch during update (avoids duplication)
|
||||
assert req_id not in runner.input_batch.req_id_to_index
|
||||
|
||||
|
||||
def test_e2e_streaming_with_multimodal_features(mock_model_runner_with_input_batch):
|
||||
"""Test that streaming session with multimodal features are updated correctly.
|
||||
|
||||
This test validates that when a streaming session with mm features is updated:
|
||||
1. The request is removed from InputBatch before updating (avoids duplication)
|
||||
2. Multimodal features from both requests are preserved and merged correctly
|
||||
3. New prompt tokens (including intermediate outputs) are appended correctly
|
||||
4. output_token_ids is cleared (intermediate outputs are now in prompt_token_ids)
|
||||
"""
|
||||
runner = mock_model_runner_with_input_batch
|
||||
req_id = "streaming_mm_req_0"
|
||||
|
||||
# Step 1: Create initial request state with one multimodal feature
|
||||
mm_feature_1 = MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem.dummy("audio"),
|
||||
modality="audio",
|
||||
identifier="audio_1",
|
||||
mm_position=PlaceholderRange(offset=2, length=10),
|
||||
)
|
||||
|
||||
initial_req_state = CachedRequestState(
|
||||
req_id=req_id,
|
||||
prompt_token_ids=[1, 2] + [0] * 10 + [3, 4], # 2 + 10 (mm) + 2 = 14 tokens
|
||||
mm_features=[mm_feature_1],
|
||||
sampling_params=SamplingParams(),
|
||||
pooling_params=None,
|
||||
generator=None,
|
||||
block_ids=([0],),
|
||||
num_computed_tokens=14,
|
||||
output_token_ids=[100], # Generated 1 token
|
||||
)
|
||||
runner.requests[req_id] = initial_req_state
|
||||
|
||||
# Add request to InputBatch
|
||||
runner.input_batch.add_request(initial_req_state)
|
||||
assert req_id in runner.input_batch.req_id_to_index
|
||||
|
||||
# Step 2: Create new request data with additional multimodal feature
|
||||
# The scheduler has already set prompt_token_ids to the full sequence
|
||||
# (original prompt + intermediate outputs + new prompt with new multimodal feature)
|
||||
mm_feature_2 = MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem.dummy("audio"),
|
||||
modality="audio",
|
||||
identifier="audio_2",
|
||||
mm_position=PlaceholderRange(offset=15, length=5),
|
||||
)
|
||||
|
||||
new_req_data = Mock()
|
||||
# Full sequence: [1, 2] + [0]*10 + [3, 4] + [100] + [0]*5 + [5] = 21 tokens
|
||||
new_req_data.prompt_token_ids = [1, 2] + [0] * 10 + [3, 4, 100] + [0] * 5 + [5]
|
||||
new_req_data.mm_features = [mm_feature_1, mm_feature_2]
|
||||
new_req_data.prompt_embeds = None
|
||||
new_req_data.sampling_params = SamplingParams(temperature=0.7, max_tokens=30)
|
||||
new_req_data.pooling_params = None
|
||||
new_req_data.block_ids = ([0, 1],)
|
||||
new_req_data.num_computed_tokens = 14 # 14 tokens from initial request
|
||||
|
||||
# Step 3: Update the request
|
||||
updated_req_state = GPUModelRunner._update_streaming_request(
|
||||
runner, req_id, new_req_data
|
||||
)
|
||||
|
||||
# Step 4: Verify the request state was updated correctly
|
||||
# Verify multimodal features are preserved
|
||||
assert len(updated_req_state.mm_features) == 2
|
||||
assert updated_req_state.mm_features[0] == mm_feature_1
|
||||
assert updated_req_state.mm_features[1] == mm_feature_2
|
||||
|
||||
# Verify prompt tokens include intermediate output (100) and new tokens
|
||||
# Initial: 2 + 10 (mm1) + 2 = 14 tokens
|
||||
# New: 2 + 10 (mm1) + 2 + 1 (output 100) + 5 (mm2) + 1 = 21 tokens
|
||||
assert len(updated_req_state.prompt_token_ids) == 21
|
||||
assert updated_req_state.prompt_token_ids == [1, 2] + [0] * 10 + [3, 4, 100] + [
|
||||
0
|
||||
] * 5 + [5]
|
||||
|
||||
# Verify output_token_ids were cleared
|
||||
# (intermediate outputs are now in prompt_token_ids)
|
||||
assert updated_req_state.output_token_ids == []
|
||||
|
||||
# Verify other parameters were updated
|
||||
assert updated_req_state.num_computed_tokens == 14
|
||||
assert updated_req_state.sampling_params.temperature == 0.7
|
||||
assert updated_req_state.sampling_params.max_tokens == 30
|
||||
assert updated_req_state.block_ids == ([0, 1],)
|
||||
|
||||
# Verify the same object is returned
|
||||
assert runner.requests[req_id] is updated_req_state
|
||||
|
||||
# Verify request was removed from InputBatch during update (avoids duplication)
|
||||
assert req_id not in runner.input_batch.req_id_to_index
|
||||
@@ -1,575 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import DeviceConfig, VllmConfig
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFeatureSpec,
|
||||
MultiModalKwargsItem,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.scheduler import Scheduler
|
||||
from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
)
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm.v1.structured_output import StructuredOutputManager
|
||||
|
||||
STOP_TOKEN = 128001
|
||||
|
||||
|
||||
class DummyRequest(Request):
|
||||
def __init__(
|
||||
self,
|
||||
request_id,
|
||||
resumable=True,
|
||||
prompt_token_ids=None,
|
||||
mm_features: list[MultiModalFeatureSpec] | None = None,
|
||||
max_tokens: int | None = 16,
|
||||
):
|
||||
super().__init__(
|
||||
request_id=request_id,
|
||||
prompt_token_ids=prompt_token_ids if prompt_token_ids is not None else [],
|
||||
sampling_params=SamplingParams(
|
||||
stop_token_ids=[STOP_TOKEN], max_tokens=max_tokens
|
||||
),
|
||||
pooling_params=None,
|
||||
eos_token_id=None,
|
||||
mm_features=mm_features,
|
||||
resumable=resumable,
|
||||
)
|
||||
|
||||
|
||||
def create_scheduler() -> Scheduler:
|
||||
vllm_config = VllmConfig(device_config=DeviceConfig("cpu"))
|
||||
vllm_config.model_config = MagicMock()
|
||||
vllm_config.model_config.skip_tokenizer_init = True
|
||||
vllm_config.model_config.is_multimodal_model = False
|
||||
vllm_config.model_config.max_model_len = 1024
|
||||
vllm_config.model_config.enable_return_routed_experts = False
|
||||
vllm_config.cache_config = MagicMock()
|
||||
vllm_config.cache_config.num_gpu_blocks = 1000
|
||||
vllm_config.cache_config.enable_prefix_caching = False
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=1000,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["layer"],
|
||||
FullAttentionSpec(
|
||||
block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
|
||||
),
|
||||
)
|
||||
],
|
||||
)
|
||||
return Scheduler(
|
||||
vllm_config=vllm_config,
|
||||
kv_cache_config=kv_cache_config,
|
||||
log_stats=True,
|
||||
structured_output_manager=StructuredOutputManager(vllm_config),
|
||||
block_size=16,
|
||||
)
|
||||
|
||||
|
||||
class TestStreamingScheduler(unittest.TestCase):
|
||||
def test_add_request(self):
|
||||
scheduler = create_scheduler()
|
||||
|
||||
request = DummyRequest(
|
||||
request_id="test_request",
|
||||
resumable=True,
|
||||
)
|
||||
|
||||
scheduler.add_request(request)
|
||||
|
||||
assert "test_request" in scheduler.requests
|
||||
assert request.status == RequestStatus.WAITING
|
||||
assert len(scheduler.waiting) == 1
|
||||
|
||||
next_request = DummyRequest(
|
||||
request_id="test_request",
|
||||
resumable=True,
|
||||
)
|
||||
scheduler.add_request(next_request)
|
||||
|
||||
assert next_request.status == RequestStatus.WAITING
|
||||
assert len(scheduler.requests["test_request"].streaming_queue) == 1
|
||||
|
||||
def test_update_request_as_session_max_token(self):
|
||||
scheduler = create_scheduler()
|
||||
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
)
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
session.max_tokens = 10 # Initial max_tokens
|
||||
session._output_token_ids = [1] * 10 # reach max_tokens
|
||||
|
||||
new_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[4, 5, 6],
|
||||
)
|
||||
new_request.sampling_params = SamplingParams(max_tokens=10)
|
||||
new_request.max_tokens = 10 # Additional max_tokens from new request
|
||||
|
||||
update = StreamingUpdate.from_request(new_request)
|
||||
scheduler._update_request_as_session(session, update)
|
||||
|
||||
assert session.sampling_params.max_tokens == 10
|
||||
# _update_request_as_session clears output tokens first, so
|
||||
# max_tokens = num_output_tokens (0) + update.max_tokens (10) = 10
|
||||
assert session.max_tokens == 10
|
||||
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
|
||||
# Simulate generating 5 more output tokens
|
||||
session._output_token_ids = [1] * 5
|
||||
new_request2 = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[7, 8, 9],
|
||||
)
|
||||
new_request2.sampling_params = SamplingParams(max_tokens=10)
|
||||
new_request2.max_tokens = 10
|
||||
update2 = StreamingUpdate.from_request(new_request2)
|
||||
scheduler._update_request_as_session(session, update2)
|
||||
|
||||
assert session.sampling_params.max_tokens == 10
|
||||
# Again, output tokens are cleared first, so max_tokens = 0 + 10 = 10
|
||||
assert session.max_tokens == 10
|
||||
|
||||
def test_update_request_as_session(self):
|
||||
scheduler = create_scheduler()
|
||||
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
)
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
|
||||
new_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[4, 5, 6],
|
||||
)
|
||||
new_request.sampling_params = SamplingParams(max_tokens=10)
|
||||
|
||||
update = StreamingUpdate.from_request(new_request)
|
||||
scheduler._update_request_as_session(session, update)
|
||||
|
||||
assert session.prompt_token_ids == [1, 2, 3, 4, 5, 6]
|
||||
assert session._all_token_ids == [1, 2, 3, 4, 5, 6]
|
||||
assert session.sampling_params.max_tokens == 10
|
||||
assert session.status == RequestStatus.WAITING
|
||||
|
||||
def test_update_request_as_session_with_multimodal(self):
|
||||
scheduler = create_scheduler()
|
||||
|
||||
mm_feature = MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem.dummy("audio"),
|
||||
modality="audio",
|
||||
identifier="",
|
||||
mm_position=PlaceholderRange(offset=1, length=1),
|
||||
)
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
mm_features=[mm_feature],
|
||||
)
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
|
||||
mm_feature = MultiModalFeatureSpec(
|
||||
data=MultiModalKwargsItem.dummy("audio"),
|
||||
modality="audio",
|
||||
identifier="",
|
||||
mm_position=PlaceholderRange(offset=2, length=1),
|
||||
)
|
||||
new_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[4, 5, 6, 7],
|
||||
mm_features=[mm_feature],
|
||||
)
|
||||
update = StreamingUpdate.from_request(new_request)
|
||||
scheduler._update_request_as_session(session, update)
|
||||
|
||||
assert len(session.mm_features) == 2
|
||||
assert session.mm_features[0].mm_position.offset == 1
|
||||
# 2 + len([1, 2, 3])
|
||||
assert session.mm_features[1].mm_position.offset == 5
|
||||
|
||||
def test_process_streaming_requests_with_finish_session(self):
|
||||
"""Test that a non-resumable request signals stream completion.
|
||||
|
||||
With the new streaming API, completion is signaled by closing/finishing
|
||||
the input generator. When a non-resumable request is added to a session
|
||||
in WAITING_FOR_STREAMING_REQ state, the session is finished immediately
|
||||
with FINISHED_ABORTED status.
|
||||
"""
|
||||
scheduler = create_scheduler()
|
||||
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
resumable=True,
|
||||
)
|
||||
scheduler.add_request(session)
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
|
||||
# A non-resumable request signals stream completion
|
||||
close_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[0],
|
||||
resumable=False,
|
||||
max_tokens=1,
|
||||
)
|
||||
scheduler.add_request(close_request)
|
||||
|
||||
# The session should be immediately finished (stream completed)
|
||||
assert session.status == RequestStatus.FINISHED_ABORTED
|
||||
# The session should be removed from the scheduler
|
||||
assert session.request_id not in scheduler.requests
|
||||
|
||||
def test_streaming_request_session_update(self):
|
||||
"""Test that a resumable request updates a waiting session directly.
|
||||
|
||||
When a session is in WAITING_FOR_STREAMING_REQ state and a new resumable
|
||||
request arrives, the update is applied directly via _update_request_as_session,
|
||||
not queued.
|
||||
"""
|
||||
scheduler = create_scheduler()
|
||||
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
resumable=True,
|
||||
)
|
||||
scheduler.add_request(session)
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
|
||||
next_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[4, 5],
|
||||
resumable=True,
|
||||
)
|
||||
|
||||
scheduler.add_request(next_request)
|
||||
|
||||
# With the new behavior, when session is in WAITING_FOR_STREAMING_REQ,
|
||||
# the update is applied directly (not queued), and session status
|
||||
# becomes WAITING
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert session.prompt_token_ids == [1, 2, 3, 4, 5]
|
||||
|
||||
_ = scheduler.schedule()
|
||||
|
||||
assert session.status == RequestStatus.RUNNING
|
||||
|
||||
def test_update_request_as_session_with_output_tokens(self):
|
||||
scheduler = create_scheduler()
|
||||
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3], # 3 prompt tokens
|
||||
)
|
||||
session.append_output_token_ids([10, 11])
|
||||
"""
|
||||
The last output token (11) hasn't been "scheduled" yet, so `num_computed_tokens`
|
||||
only includes: 3 prompt + 1 output (the 10) = 4
|
||||
"""
|
||||
session.num_computed_tokens = 4
|
||||
|
||||
new_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[4, 5],
|
||||
)
|
||||
|
||||
update = StreamingUpdate.from_request(new_request)
|
||||
scheduler._update_request_as_session(session, update)
|
||||
|
||||
# _update_request_as_session keeps computed output tokens (they become
|
||||
# part of the prompt) and only discards the final uncomputed sampled
|
||||
# token. Computed output token 10 is kept, uncomputed token 11 is
|
||||
# discarded.
|
||||
assert session._all_token_ids == [1, 2, 3, 10, 4, 5]
|
||||
assert session.prompt_token_ids == [1, 2, 3, 10, 4, 5]
|
||||
# Output tokens list is cleared
|
||||
assert session._output_token_ids == []
|
||||
# num_computed_tokens is unchanged (KV cache still valid for computed
|
||||
# tokens)
|
||||
assert session.num_computed_tokens == 4
|
||||
# Verify that the next schedule will only process the new prompt tokens
|
||||
# num_new_tokens = num_tokens - num_computed_tokens = 6 - 4 = 2
|
||||
num_new_tokens = session.num_tokens - session.num_computed_tokens
|
||||
assert num_new_tokens == 2
|
||||
|
||||
def test_streaming_e2e_lifecycle(self):
|
||||
"""
|
||||
Comprehensive integration test covering complete streaming request lifecycle
|
||||
including scheduler state management and aliasing bug prevention.
|
||||
|
||||
FULL LIFECYCLE:
|
||||
================
|
||||
CYCLE 1 (Initial Decode):
|
||||
1. Add streaming request (seq_id=0) with prompt tokens [1,2,3]
|
||||
2. Schedule() creates NewRequestData with prompt_token_ids
|
||||
3. Model runner caches this prompt_token_ids reference (simulated)
|
||||
4. Model executes and generates output token 10
|
||||
5. update_from_output() appends token 10 to request._all_token_ids
|
||||
6. Request transitions to RUNNING state
|
||||
|
||||
CYCLE 2 (Continue Decode):
|
||||
7. Schedule() again - request is now in scheduled_cached_reqs (not new)
|
||||
8. Model runner uses CACHED state to calculate num_tokens
|
||||
9. Model generates output token (STOP_TOKEN)
|
||||
10. update_from_output() appends STOP_TOKEN to request._all_token_ids
|
||||
11. Request transitions to WAITING_FOR_STREAMING_REQ
|
||||
|
||||
CYCLE 3 (New Streaming Request):
|
||||
12. Add new streaming request (seq_id=1) with prompt tokens [4,5]
|
||||
13. Scheduler merges into session, creates NewRequestData again
|
||||
14. Model runner caches new prompt_token_ids reference
|
||||
15. Verify cached state from Cycle 1 wasn't corrupted by mutations
|
||||
|
||||
CRITICAL BUG PREVENTION:
|
||||
========================
|
||||
Without .copy() in _create_new_request_data():
|
||||
- Cycle 1 Step 3: cached_state["prompt_token_ids"] aliases
|
||||
request._all_token_ids
|
||||
- Cycle 1 Step 5: When appending token 10, cached state mutates:
|
||||
[1,2,3] -> [1,2,3,10]
|
||||
- Cycle 2 Step 8: num_tokens = len([1,2,3,10]) + len([10])
|
||||
= 5 (WRONG! Should be 4)
|
||||
- Cycle 2: Discard logic would see seq_lens=4 < num_tokens=5
|
||||
-> INCORRECTLY DISCARDS
|
||||
|
||||
With .copy() in _create_new_request_data():
|
||||
- Cycle 1 Step 3: cached_state["prompt_token_ids"] is independent copy
|
||||
- Cycle 1 Step 5: Only request._all_token_ids mutates, cached stays [1,2,3]
|
||||
- Cycle 2 Step 8: num_tokens = len([1,2,3]) + len([10]) = 4 (CORRECT)
|
||||
- Cycle 2: Discard logic works correctly
|
||||
"""
|
||||
scheduler = create_scheduler()
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
# CYCLE 1: Initial Request Scheduling and First Decode
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
session = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
)
|
||||
scheduler.add_request(session)
|
||||
|
||||
# Step 2: Schedule creates NewRequestData
|
||||
scheduler_output_cycle1 = scheduler.schedule()
|
||||
|
||||
# Verify request is in scheduled_new_reqs (first time scheduling)
|
||||
assert len(scheduler_output_cycle1.scheduled_new_reqs) == 1
|
||||
new_req_data_cycle1 = scheduler_output_cycle1.scheduled_new_reqs[0]
|
||||
assert new_req_data_cycle1.prompt_token_ids == [1, 2, 3]
|
||||
assert (
|
||||
scheduler_output_cycle1.num_scheduled_tokens[session.request_id] == 3
|
||||
) # [1, 2, 3]
|
||||
assert (
|
||||
session.request_id
|
||||
not in scheduler_output_cycle1.scheduled_cached_reqs.req_ids
|
||||
)
|
||||
|
||||
# Step 3: Simulate model runner caching the prompt_token_ids
|
||||
# This simulates gpu_model_runner.py:706-720 CachedRequestState creation
|
||||
# The model runner makes a copy of prompt_token_ids when creating
|
||||
# CachedRequestState
|
||||
cached_state_cycle1 = {
|
||||
"req_id": session.request_id,
|
||||
"prompt_token_ids": list(
|
||||
new_req_data_cycle1.prompt_token_ids
|
||||
), # Explicit copy
|
||||
"output_token_ids": [],
|
||||
"num_computed_tokens": 0,
|
||||
}
|
||||
|
||||
# Store original for verification
|
||||
original_cached_prompt_cycle1 = cached_state_cycle1["prompt_token_ids"].copy()
|
||||
|
||||
# Step 4-5: Model execution generates token, scheduler updates request
|
||||
output_token_1 = 10
|
||||
cached_state_cycle1["output_token_ids"].append(output_token_1)
|
||||
|
||||
mro_cycle1 = ModelRunnerOutput(
|
||||
req_ids=[session.request_id],
|
||||
req_id_to_index={session.request_id: 0},
|
||||
sampled_token_ids=[[output_token_1]],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={session.request_id: None},
|
||||
pooler_output=[],
|
||||
)
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
eco_dict_cycle1 = scheduler.update_from_output(
|
||||
scheduler_output_cycle1, mro_cycle1
|
||||
)
|
||||
|
||||
# Step 6: Verify request state after Cycle 1
|
||||
eco_cycle1 = eco_dict_cycle1[session.client_index].outputs[0]
|
||||
assert eco_cycle1.finish_reason is None # Not stopped yet
|
||||
assert session.status == RequestStatus.RUNNING
|
||||
assert session in scheduler.running
|
||||
assert session._all_token_ids == [1, 2, 3, 10] # Mutation happened here
|
||||
|
||||
# CRITICAL ASSERTION: Cached prompt_token_ids must NOT have changed
|
||||
assert (
|
||||
cached_state_cycle1["prompt_token_ids"] == original_cached_prompt_cycle1
|
||||
), (
|
||||
f"ALIASING BUG DETECTED in Cycle 1! "
|
||||
f"cached_state['prompt_token_ids'] was mutated from "
|
||||
f"{original_cached_prompt_cycle1} to "
|
||||
f"{cached_state_cycle1['prompt_token_ids']}. "
|
||||
f"This means _create_new_request_data() didn't call .copy()!"
|
||||
)
|
||||
assert cached_state_cycle1["prompt_token_ids"] is not session._all_token_ids, (
|
||||
"ALIASING BUG! cached_state['prompt_token_ids'] is the same object as "
|
||||
"session._all_token_ids. They must be independent copies."
|
||||
)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
# CYCLE 2: Continue Decoding (Using Cached State)
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
# Step 7: Schedule again - now request uses cached state
|
||||
scheduler_output_cycle2 = scheduler.schedule()
|
||||
|
||||
# Verify request is NOT in scheduled_new_reqs (already cached)
|
||||
assert not scheduler_output_cycle2.scheduled_new_reqs
|
||||
assert (
|
||||
session.request_id in scheduler_output_cycle2.scheduled_cached_reqs.req_ids
|
||||
)
|
||||
assert (
|
||||
scheduler_output_cycle2.num_scheduled_tokens[session.request_id] == 1
|
||||
) # Only the output token [10]
|
||||
|
||||
# Step 8: Calculate num_tokens like gpu_model_runner.py:1284 does
|
||||
# This is where the bug would manifest!
|
||||
num_tokens_cycle2 = len(cached_state_cycle1["prompt_token_ids"]) + len(
|
||||
cached_state_cycle1["output_token_ids"]
|
||||
)
|
||||
|
||||
# CRITICAL ASSERTION: num_tokens must be correct (3 prompt + 1 output = 4)
|
||||
# Without .copy(), cached_state["prompt_token_ids"] would be [1,2,3,10]
|
||||
# and num_tokens would incorrectly be 5, causing the discard bug
|
||||
expected_num_tokens_cycle2 = 4
|
||||
assert num_tokens_cycle2 == expected_num_tokens_cycle2, (
|
||||
f"DISCARD BUG WOULD TRIGGER! num_tokens calculation is wrong. "
|
||||
f"Expected {expected_num_tokens_cycle2}, got {num_tokens_cycle2}. "
|
||||
f"cached_state['prompt_token_ids'] = "
|
||||
f"{cached_state_cycle1['prompt_token_ids']} (should be [1,2,3], not [1,2,3,"
|
||||
f"10]). Without .copy(), this would be 5 = len([1,2,3,10]) + len([10]). "
|
||||
f"Discard logic would see: seq_lens={session.num_computed_tokens} "
|
||||
f"< num_tokens={num_tokens_cycle2}, triggering incorrect discard!"
|
||||
)
|
||||
|
||||
# Step 9-10: Model generates STOP_TOKEN, scheduler updates
|
||||
output_token_2 = STOP_TOKEN
|
||||
cached_state_cycle1["output_token_ids"].append(output_token_2)
|
||||
|
||||
mro_cycle2 = ModelRunnerOutput(
|
||||
req_ids=[session.request_id],
|
||||
req_id_to_index={session.request_id: 0},
|
||||
sampled_token_ids=[[output_token_2]],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={session.request_id: None},
|
||||
pooler_output=[],
|
||||
)
|
||||
eco_dict_cycle2 = scheduler.update_from_output(
|
||||
scheduler_output_cycle2, mro_cycle2
|
||||
)
|
||||
|
||||
# Step 11: Verify request transitioned to WAITING_FOR_STREAMING_REQ
|
||||
eco_cycle2 = eco_dict_cycle2[session.client_index].outputs[0]
|
||||
assert eco_cycle2.finish_reason == FinishReason.STOP
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
assert session in scheduler.waiting
|
||||
assert session._all_token_ids == [1, 2, 3, 10, STOP_TOKEN]
|
||||
|
||||
# CRITICAL ASSERTION: Cached prompt_token_ids STILL must not have changed
|
||||
assert cached_state_cycle1["prompt_token_ids"] == [1, 2, 3], (
|
||||
f"ALIASING BUG DETECTED in Cycle 2! "
|
||||
f"cached_state['prompt_token_ids'] = "
|
||||
f"{cached_state_cycle1['prompt_token_ids']} (should still be [1,2,3]). "
|
||||
f"Mutations from update_from_output() leaked through!"
|
||||
)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
# CYCLE 3: New Streaming Request (Session Continuation)
|
||||
# ═══════════════════════════════════════════════════════════════════
|
||||
|
||||
# Step 12: Add new streaming request with seq_id=1
|
||||
new_request = DummyRequest(
|
||||
request_id="session",
|
||||
prompt_token_ids=[4, 5],
|
||||
)
|
||||
scheduler.add_request(new_request)
|
||||
|
||||
# With the new streaming API, when session is in WAITING_FOR_STREAMING_REQ,
|
||||
# the update is applied directly via _update_request_as_session (not queued).
|
||||
# The session status becomes WAITING after the update is applied.
|
||||
assert session.status == RequestStatus.WAITING
|
||||
|
||||
# Step 13: Scheduler schedules the updated session
|
||||
scheduler_output_cycle3 = scheduler.schedule()
|
||||
|
||||
# Verify scheduler created NewRequestData with merged prompt_token_ids
|
||||
assert len(scheduler_output_cycle3.scheduled_new_reqs) == 1
|
||||
assert (
|
||||
scheduler_output_cycle3.scheduled_new_reqs[0].prompt_token_ids
|
||||
== session.prompt_token_ids
|
||||
)
|
||||
assert (
|
||||
scheduler_output_cycle3.num_scheduled_tokens[session.request_id] == 2
|
||||
) # Only new tokens [4, 5]
|
||||
# Computed output tokens are kept (become part of prompt), only the
|
||||
# final uncomputed sampled token (STOP_TOKEN) is discarded
|
||||
assert session._all_token_ids == [1, 2, 3, 10, 4, 5]
|
||||
assert session.prompt_token_ids == [1, 2, 3, 10, 4, 5] # Includes kept output
|
||||
assert session._output_token_ids == [] # Output tokens are cleared
|
||||
|
||||
# Step 14: Model runner caches NEW prompt_token_ids reference
|
||||
# The model runner makes a copy of prompt_token_ids when creating
|
||||
# CachedRequestState
|
||||
new_req_data_cycle3 = scheduler_output_cycle3.scheduled_new_reqs[0]
|
||||
cached_state_cycle3 = {
|
||||
"req_id": session.request_id,
|
||||
"prompt_token_ids": list(
|
||||
new_req_data_cycle3.prompt_token_ids
|
||||
), # Explicit copy
|
||||
"output_token_ids": [],
|
||||
"num_computed_tokens": session.num_computed_tokens,
|
||||
}
|
||||
|
||||
# Step 15: FINAL CRITICAL VERIFICATION
|
||||
# The old cached state from Cycle 1 must still be unchanged
|
||||
assert cached_state_cycle1["prompt_token_ids"] == [1, 2, 3], (
|
||||
f"PERSISTENT ALIASING BUG! Even after new scheduling cycle, "
|
||||
f"old cached_state was mutated to "
|
||||
f"{cached_state_cycle1['prompt_token_ids']}. This proves the aliasing bug "
|
||||
f"exists!"
|
||||
)
|
||||
|
||||
# The new cached state must be independent
|
||||
assert cached_state_cycle3["prompt_token_ids"] is not session._all_token_ids, (
|
||||
"ALIASING BUG in Cycle 3! Cached state is aliased to _all_token_ids."
|
||||
)
|
||||
|
||||
# Both cached states must be independent of each other
|
||||
assert (
|
||||
cached_state_cycle1["prompt_token_ids"]
|
||||
is not cached_state_cycle3["prompt_token_ids"]
|
||||
), "Cached states from different cycles should be independent objects."
|
||||
@@ -8,7 +8,6 @@ def test_request_status_fmt_str():
|
||||
assert f"{RequestStatus.WAITING}" == "WAITING"
|
||||
assert f"{RequestStatus.WAITING_FOR_FSM}" == "WAITING_FOR_FSM"
|
||||
assert f"{RequestStatus.WAITING_FOR_REMOTE_KVS}" == "WAITING_FOR_REMOTE_KVS"
|
||||
assert f"{RequestStatus.WAITING_FOR_STREAMING_REQ}" == "WAITING_FOR_STREAMING_REQ"
|
||||
assert f"{RequestStatus.RUNNING}" == "RUNNING"
|
||||
assert f"{RequestStatus.PREEMPTED}" == "PREEMPTED"
|
||||
assert f"{RequestStatus.FINISHED_STOPPED}" == "FINISHED_STOPPED"
|
||||
|
||||
@@ -1,233 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Generate S3 PyPI Root Index for Latest Version
|
||||
#
|
||||
# Creates a PEP 503 compatible index.html at rocm/ pointing to the latest
|
||||
# semantic version's packages. This enables users to install with:
|
||||
# uv pip install vllm --extra-index-url s3://vllm-wheels/rocm
|
||||
#
|
||||
# Usage:
|
||||
# generate-root-index.sh [options]
|
||||
#
|
||||
# Options:
|
||||
# --dry-run Preview changes without uploading
|
||||
# --version VER Use specific version instead of auto-detecting latest
|
||||
#
|
||||
# Environment variables:
|
||||
# S3_BUCKET - Bucket name (default: vllm-wheels)
|
||||
# VARIANT - ROCm variant (default: rocm700)
|
||||
# DRY_RUN - Set to 1 for preview mode (same as --dry-run)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# ======== Configuration ========
|
||||
BUCKET="${S3_BUCKET:-vllm-wheels}"
|
||||
VARIANT="${VARIANT:-rocm700}"
|
||||
DRY_RUN="${DRY_RUN:-0}"
|
||||
FORCE_VERSION=""
|
||||
|
||||
# Parse command line arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
--version)
|
||||
FORCE_VERSION="$2"
|
||||
shift 2
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Working directory for generated files
|
||||
WORK_DIR=$(mktemp -d)
|
||||
trap 'rm -rf "$WORK_DIR"' EXIT
|
||||
|
||||
echo "========================================"
|
||||
echo "Generate Root Index for Latest Version"
|
||||
echo "========================================"
|
||||
echo "S3 Bucket: $BUCKET"
|
||||
echo "ROCm Variant: $VARIANT"
|
||||
echo "Dry Run: $DRY_RUN"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
|
||||
# ======== Step 1: Find latest semantic version ========
|
||||
|
||||
echo "Step 1: Finding latest semantic version..."
|
||||
|
||||
# List all directories under rocm/
|
||||
aws s3api list-objects-v2 \
|
||||
--bucket "$BUCKET" \
|
||||
--prefix "rocm/" \
|
||||
--delimiter "/" \
|
||||
--query 'CommonPrefixes[].Prefix' \
|
||||
--output text | tr '\t' '\n' > "$WORK_DIR/all_prefixes.txt"
|
||||
|
||||
# Filter for semantic versions (x.y.z pattern)
|
||||
grep -oE 'rocm/[0-9]+\.[0-9]+\.[0-9]+/' "$WORK_DIR/all_prefixes.txt" | \
|
||||
sed 's|rocm/||; s|/||' | \
|
||||
sort -V > "$WORK_DIR/versions.txt" || true
|
||||
|
||||
if [[ ! -s "$WORK_DIR/versions.txt" ]]; then
|
||||
echo "ERROR: No semantic versions found under s3://$BUCKET/rocm/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Found versions:"
|
||||
cat "$WORK_DIR/versions.txt"
|
||||
echo ""
|
||||
|
||||
if [[ -n "$FORCE_VERSION" ]]; then
|
||||
LATEST_VERSION="$FORCE_VERSION"
|
||||
echo "Using forced version: $LATEST_VERSION"
|
||||
else
|
||||
LATEST_VERSION=$(tail -1 "$WORK_DIR/versions.txt")
|
||||
echo "Latest version (auto-detected): $LATEST_VERSION"
|
||||
fi
|
||||
|
||||
# Verify the version exists
|
||||
if ! grep -qx "$LATEST_VERSION" "$WORK_DIR/versions.txt"; then
|
||||
echo "ERROR: Version $LATEST_VERSION not found in bucket"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# ======== Step 2: List packages from latest version ========
|
||||
|
||||
echo ""
|
||||
echo "Step 2: Listing packages from rocm/$LATEST_VERSION/$VARIANT/..."
|
||||
|
||||
VERSION_PREFIX="rocm/$LATEST_VERSION/$VARIANT/"
|
||||
|
||||
# List package directories
|
||||
aws s3api list-objects-v2 \
|
||||
--bucket "$BUCKET" \
|
||||
--prefix "$VERSION_PREFIX" \
|
||||
--delimiter "/" \
|
||||
--query 'CommonPrefixes[].Prefix' \
|
||||
--output text | tr '\t' '\n' > "$WORK_DIR/package_prefixes.txt" || true
|
||||
|
||||
if [[ ! -s "$WORK_DIR/package_prefixes.txt" ]]; then
|
||||
echo "ERROR: No packages found under s3://$BUCKET/$VERSION_PREFIX"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Extract package names
|
||||
sed "s|${VERSION_PREFIX}||; s|/||g" "$WORK_DIR/package_prefixes.txt" | \
|
||||
grep -v '^$' > "$WORK_DIR/packages.txt"
|
||||
|
||||
echo "Found packages:"
|
||||
cat "$WORK_DIR/packages.txt"
|
||||
echo ""
|
||||
|
||||
# ======== Step 3: Generate root index.html ========
|
||||
|
||||
echo "Step 3: Generating root index.html..."
|
||||
|
||||
mkdir -p "$WORK_DIR/output"
|
||||
|
||||
{
|
||||
cat <<'EOF'
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<meta name="pypi:repository-version" content="1.0">
|
||||
</head>
|
||||
<body>
|
||||
EOF
|
||||
|
||||
while read -r pkg; do
|
||||
echo " <a href=\"$pkg/\">$pkg</a><br>"
|
||||
done < "$WORK_DIR/packages.txt"
|
||||
|
||||
cat <<'EOF'
|
||||
</body>
|
||||
</html>
|
||||
EOF
|
||||
} > "$WORK_DIR/output/index.html"
|
||||
|
||||
echo "Generated root index.html:"
|
||||
cat "$WORK_DIR/output/index.html"
|
||||
echo ""
|
||||
|
||||
# ======== Step 4: Copy and adjust package index files ========
|
||||
|
||||
echo "Step 4: Copying and adjusting package index files..."
|
||||
|
||||
while read -r pkg; do
|
||||
echo "Processing package: $pkg"
|
||||
|
||||
# Download existing index.html from versioned path
|
||||
SOURCE_INDEX="s3://$BUCKET/$VERSION_PREFIX$pkg/index.html"
|
||||
|
||||
mkdir -p "$WORK_DIR/output/$pkg"
|
||||
|
||||
if aws s3 cp "$SOURCE_INDEX" "$WORK_DIR/output/$pkg/index.html" 2>/dev/null; then
|
||||
# Adjust relative paths:
|
||||
# Original: href="../../../{commit}/wheel.whl" (from rocm/0.13.0/rocm710/vllm/)
|
||||
# New: href="../{commit}/wheel.whl" (from rocm/vllm/)
|
||||
sed -i 's|href="\.\./\.\./\.\./|href="../|g' "$WORK_DIR/output/$pkg/index.html"
|
||||
echo " - Downloaded and adjusted: $pkg/index.html"
|
||||
else
|
||||
echo " - WARNING: Could not download index for $pkg"
|
||||
fi
|
||||
done < "$WORK_DIR/packages.txt"
|
||||
|
||||
echo ""
|
||||
|
||||
# ======== Step 5: Upload to S3 ========
|
||||
|
||||
echo "Step 5: Uploading to s3://$BUCKET/rocm/..."
|
||||
echo ""
|
||||
|
||||
# List what would be uploaded
|
||||
echo "Files to upload:"
|
||||
find "$WORK_DIR/output" -name "*.html" -type f | while read -r file; do
|
||||
rel_path="${file#$WORK_DIR/output/}"
|
||||
echo " rocm/$rel_path"
|
||||
done
|
||||
echo ""
|
||||
|
||||
if [[ "$DRY_RUN" == "1" ]]; then
|
||||
echo "DRY RUN - Skipping upload"
|
||||
echo ""
|
||||
echo "Preview of generated files:"
|
||||
echo "----------------------------------------"
|
||||
echo "rocm/index.html:"
|
||||
cat "$WORK_DIR/output/index.html"
|
||||
echo ""
|
||||
echo "----------------------------------------"
|
||||
echo "Sample package index (first package):"
|
||||
FIRST_PKG=$(head -1 "$WORK_DIR/packages.txt")
|
||||
if [[ -f "$WORK_DIR/output/$FIRST_PKG/index.html" ]]; then
|
||||
echo "rocm/$FIRST_PKG/index.html:"
|
||||
cat "$WORK_DIR/output/$FIRST_PKG/index.html"
|
||||
fi
|
||||
else
|
||||
# Upload all generated files
|
||||
aws s3 cp --recursive "$WORK_DIR/output/" "s3://$BUCKET/rocm/" \
|
||||
--content-type "text/html"
|
||||
|
||||
echo "Upload complete!"
|
||||
fi
|
||||
|
||||
# ======== Summary ========
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo "Root Index Generation Complete!"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
echo "Latest version: $LATEST_VERSION"
|
||||
echo "Packages indexed: $(wc -l < "$WORK_DIR/packages.txt")"
|
||||
echo ""
|
||||
echo "Install command:"
|
||||
echo " uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/"
|
||||
echo "========================================"
|
||||
+17
-22
@@ -591,8 +591,8 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
|
||||
) -> torch.Tensor:
|
||||
return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)
|
||||
|
||||
@register_fake("_C::marlin_gemm")
|
||||
def _marlin_gemm_fake(
|
||||
@register_fake("_C::gptq_marlin_gemm")
|
||||
def _gptq_marlin_gemm_fake(
|
||||
a: torch.Tensor,
|
||||
c: torch.Tensor | None,
|
||||
b_q_weight: torch.Tensor,
|
||||
@@ -1312,7 +1312,7 @@ def marlin_int4_fp8_preprocess(
|
||||
return torch.ops._C.marlin_int4_fp8_preprocess(qweight, qzeros_or_none, inplace)
|
||||
|
||||
|
||||
def marlin_gemm(
|
||||
def gptq_marlin_gemm(
|
||||
a: torch.Tensor,
|
||||
c: torch.Tensor | None,
|
||||
b_q_weight: torch.Tensor,
|
||||
@@ -1333,7 +1333,7 @@ def marlin_gemm(
|
||||
use_fp32_reduce: bool = False,
|
||||
is_zp_float: bool = False,
|
||||
) -> torch.Tensor:
|
||||
return torch.ops._C.marlin_gemm(
|
||||
return torch.ops._C.gptq_marlin_gemm(
|
||||
a,
|
||||
c,
|
||||
b_q_weight,
|
||||
@@ -1534,7 +1534,6 @@ def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
|
||||
def scaled_fp4_quant(
|
||||
input: torch.Tensor,
|
||||
input_global_scale: torch.Tensor,
|
||||
is_sf_swizzled_layout: bool = True,
|
||||
backend: str = "none",
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
@@ -1578,26 +1577,22 @@ def scaled_fp4_quant(
|
||||
else:
|
||||
# Two fp4 values will be packed into an uint8.
|
||||
output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
|
||||
if is_sf_swizzled_layout:
|
||||
# We use the rounded values to store the swizzled values. Due to the
|
||||
# requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
|
||||
# So, we first pad the scales to multiples of 128 and 4. Then, the scales
|
||||
# (in float8_e4m3fn) are packed into an int32 for every 4 values. More:
|
||||
# https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x
|
||||
round_up = lambda x, y: (x + y - 1) // y * y
|
||||
rounded_m = round_up(m, 128)
|
||||
scale_n = n // block_size
|
||||
rounded_n = round_up(scale_n, 4)
|
||||
output_scale = torch.empty(
|
||||
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
output_scale = torch.empty((m, n // 16), device=device, dtype=torch.uint8)
|
||||
|
||||
torch.ops._C.scaled_fp4_quant(
|
||||
output, input, output_scale, input_global_scale, is_sf_swizzled_layout
|
||||
# We use the rounded values to store the swizzled values. Due to the
|
||||
# requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
|
||||
# So, we first pad the scales to multiples of 128 and 4. Then, the scales
|
||||
# (in float8_e4m3fn) are packed into an int32 for every 4 values. More:
|
||||
# https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x
|
||||
round_up = lambda x, y: (x + y - 1) // y * y
|
||||
rounded_m = round_up(m, 128)
|
||||
scale_n = n // block_size
|
||||
rounded_n = round_up(scale_n, 4)
|
||||
output_scale = torch.empty(
|
||||
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
|
||||
)
|
||||
|
||||
torch.ops._C.scaled_fp4_quant(output, input, output_scale, input_global_scale)
|
||||
|
||||
output_scale = output_scale.view(torch.float8_e4m3fn)
|
||||
return output, output_scale
|
||||
|
||||
|
||||
+16
-85
@@ -390,43 +390,29 @@ class Attention(nn.Module, AttentionLayerBase):
|
||||
if value is not None:
|
||||
value = value.view(-1, self.num_kv_heads, self.head_size_v)
|
||||
if self.use_direct_call:
|
||||
kv_cache_dummy_dep = None
|
||||
if not self.attn_backend.forward_includes_kv_cache_update:
|
||||
kv_cache_dummy_dep = unified_kv_cache_update(
|
||||
key, value, self.layer_name
|
||||
)
|
||||
unified_attention_with_output(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
output,
|
||||
self.layer_name,
|
||||
kv_cache_dummy_dep=kv_cache_dummy_dep,
|
||||
forward_context: ForwardContext = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
if isinstance(attn_metadata, dict):
|
||||
attn_metadata = attn_metadata[self.layer_name]
|
||||
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
|
||||
self.impl.forward(
|
||||
self, query, key, value, self_kv_cache, attn_metadata, output=output
|
||||
)
|
||||
else:
|
||||
kv_cache_dummy_dep = None
|
||||
if not self.attn_backend.forward_includes_kv_cache_update and (
|
||||
# torch can only dispatch custom op if a tensor is passed
|
||||
key is not None or value is not None
|
||||
):
|
||||
kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
|
||||
key, value, self.layer_name
|
||||
)
|
||||
torch.ops.vllm.unified_attention_with_output(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
output,
|
||||
self.layer_name,
|
||||
kv_cache_dummy_dep=kv_cache_dummy_dep,
|
||||
query, key, value, output, self.layer_name
|
||||
)
|
||||
return output.view(-1, hidden_size)
|
||||
else:
|
||||
assert self.attn_backend.forward_includes_kv_cache_update, (
|
||||
"Split KV cache update not supported when output tensor not provided."
|
||||
)
|
||||
if self.use_direct_call:
|
||||
return unified_attention(query, key, value, self.layer_name)
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
if isinstance(attn_metadata, dict):
|
||||
attn_metadata = attn_metadata[self.layer_name]
|
||||
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
|
||||
return self.impl.forward(
|
||||
self, query, key, value, self_kv_cache, attn_metadata
|
||||
)
|
||||
else:
|
||||
return torch.ops.vllm.unified_attention(
|
||||
query, key, value, self.layer_name
|
||||
@@ -816,55 +802,6 @@ direct_register_custom_op(
|
||||
)
|
||||
|
||||
|
||||
def unified_kv_cache_update(
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> 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.
|
||||
"""
|
||||
forward_context = get_forward_context()
|
||||
attn_layer = forward_context.no_compile_layers[layer_name]
|
||||
kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
|
||||
|
||||
slot_mapping = forward_context.slot_mapping
|
||||
assert isinstance(slot_mapping, dict), (
|
||||
f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
|
||||
)
|
||||
layer_slot_mapping = slot_mapping.get(layer_name)
|
||||
if layer_slot_mapping is not None:
|
||||
assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
|
||||
f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
|
||||
)
|
||||
attn_layer.impl.do_kv_cache_update(
|
||||
attn_layer,
|
||||
key,
|
||||
value,
|
||||
kv_cache,
|
||||
layer_slot_mapping,
|
||||
)
|
||||
|
||||
return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
|
||||
|
||||
|
||||
def unified_kv_cache_update_fake(
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
layer_name: str,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty(0, device=key.device, dtype=key.dtype)
|
||||
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="unified_kv_cache_update",
|
||||
op_func=unified_kv_cache_update,
|
||||
fake_impl=unified_kv_cache_update_fake,
|
||||
mutates_args=[],
|
||||
)
|
||||
|
||||
|
||||
@maybe_transfer_kv_layer
|
||||
def unified_attention_with_output(
|
||||
query: torch.Tensor,
|
||||
@@ -874,12 +811,7 @@ def unified_attention_with_output(
|
||||
layer_name: str,
|
||||
output_scale: torch.Tensor | None = None,
|
||||
output_block_scale: torch.Tensor | None = None,
|
||||
kv_cache_dummy_dep: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
# kv_cache_dummy_dep is not used but accepting it creates a data dependency
|
||||
# that ensures torch.compile preserves ordering between KV cache update and
|
||||
# attention forward.
|
||||
del kv_cache_dummy_dep
|
||||
attn_metadata, self, kv_cache = get_attention_context(layer_name)
|
||||
|
||||
self.impl.forward(
|
||||
@@ -903,7 +835,6 @@ def unified_attention_with_output_fake(
|
||||
layer_name: str,
|
||||
output_scale: torch.Tensor | None = None,
|
||||
output_block_scale: torch.Tensor | None = None,
|
||||
kv_cache_dummy_dep: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
return
|
||||
|
||||
|
||||
+41
-141
@@ -22,10 +22,6 @@ from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm.benchmarks.datasets import (
|
||||
MultiModalConversationDataset,
|
||||
VisionArenaDataset,
|
||||
)
|
||||
from vllm.benchmarks.throughput import get_requests
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.multimodal.processing.context import (
|
||||
@@ -42,7 +38,6 @@ except ImportError:
|
||||
|
||||
def collect_mm_processor_stats(
|
||||
llm_engine: Any,
|
||||
num_warmup_reqs: int = 0,
|
||||
) -> dict[str, list[float]]:
|
||||
"""
|
||||
Collect multimodal processor timing stats.
|
||||
@@ -50,24 +45,26 @@ def collect_mm_processor_stats(
|
||||
"""
|
||||
all_stats = get_timing_stats_from_engine_client(llm_engine)
|
||||
|
||||
stat_keys = [
|
||||
"hf_processor_time",
|
||||
"hashing_time",
|
||||
"cache_lookup_time",
|
||||
"prompt_update_time",
|
||||
"preprocessor_total_time",
|
||||
"encoder_forward_time",
|
||||
"num_encoder_calls",
|
||||
]
|
||||
stats_by_stage = {key: [] for key in stat_keys}
|
||||
stats_by_stage = {
|
||||
"hf_processor_time": [],
|
||||
"hashing_time": [],
|
||||
"cache_lookup_time": [],
|
||||
"prompt_update_time": [],
|
||||
"total_time": [],
|
||||
}
|
||||
|
||||
# Skip warmup requests
|
||||
stats_list = list(all_stats.values())[num_warmup_reqs:]
|
||||
|
||||
for stats_dict in stats_list:
|
||||
for key in stat_keys:
|
||||
if key in stats_dict:
|
||||
stats_by_stage[key].append(stats_dict[key])
|
||||
for stats_dict in all_stats.values():
|
||||
stats_by_stage["hf_processor_time"].append(
|
||||
stats_dict.get("hf_processor_time", 0.0)
|
||||
)
|
||||
stats_by_stage["hashing_time"].append(stats_dict.get("hashing_time", 0.0))
|
||||
stats_by_stage["cache_lookup_time"].append(
|
||||
stats_dict.get("cache_lookup_time", 0.0)
|
||||
)
|
||||
stats_by_stage["prompt_update_time"].append(
|
||||
stats_dict.get("prompt_update_time", 0.0)
|
||||
)
|
||||
stats_by_stage["total_time"].append(stats_dict.get("total_time", 0.0))
|
||||
|
||||
return stats_by_stage
|
||||
|
||||
@@ -91,14 +88,14 @@ def calculate_mm_processor_metrics(
|
||||
}
|
||||
continue
|
||||
|
||||
is_count_metric = stage_name == "num_encoder_calls"
|
||||
values = times if is_count_metric else [t * 1000 for t in times]
|
||||
|
||||
times_ms = [t * 1000 for t in times]
|
||||
metrics[stage_name] = {
|
||||
"mean": float(np.mean(values)),
|
||||
"median": float(np.median(values)),
|
||||
"std": float(np.std(values)),
|
||||
**{f"p{p}": float(np.percentile(values, p)) for p in selected_percentiles},
|
||||
"mean": float(np.mean(times_ms)),
|
||||
"median": float(np.median(times_ms)),
|
||||
"std": float(np.std(times_ms)),
|
||||
**{
|
||||
f"p{p}": float(np.percentile(times_ms, p)) for p in selected_percentiles
|
||||
},
|
||||
}
|
||||
|
||||
return metrics
|
||||
@@ -117,23 +114,6 @@ def validate_args(args):
|
||||
if not hasattr(args, "max_loras"):
|
||||
args.max_loras = None
|
||||
|
||||
if args.dataset_name == "hf" and not args.dataset_path:
|
||||
raise ValueError(
|
||||
"--dataset-path is required when using --dataset-name hf. "
|
||||
"For multimodal benchmarking, specify a dataset like "
|
||||
"'lmarena-ai/VisionArena-Chat'."
|
||||
)
|
||||
if args.dataset_name == "hf":
|
||||
supported_mm_datasets = (
|
||||
VisionArenaDataset.SUPPORTED_DATASET_PATHS.keys()
|
||||
| MultiModalConversationDataset.SUPPORTED_DATASET_PATHS
|
||||
)
|
||||
if args.dataset_path not in supported_mm_datasets:
|
||||
raise ValueError(
|
||||
f"{args.dataset_path} is not a supported multimodal dataset. "
|
||||
f"Supported multimodal datasets are: {sorted(supported_mm_datasets)}"
|
||||
)
|
||||
|
||||
|
||||
def benchmark_multimodal_processor(
|
||||
args: argparse.Namespace,
|
||||
@@ -182,25 +162,6 @@ def benchmark_multimodal_processor(
|
||||
|
||||
freeze_gc_heap()
|
||||
|
||||
num_warmups = getattr(args, "num_warmups", 0)
|
||||
if num_warmups > 0:
|
||||
print(f"Processing {num_warmups} warmup requests...")
|
||||
# Create a temporary args object for warmup requests
|
||||
warmup_args = argparse.Namespace(**vars(args))
|
||||
warmup_args.num_prompts = num_warmups
|
||||
warmup_args.seed += 1
|
||||
warmup_requests = get_requests(warmup_args, tokenizer)
|
||||
warmup_prompts = [req.prompt for req in warmup_requests]
|
||||
warmup_output_lens = [req.expected_output_len for req in warmup_requests]
|
||||
warmup_sampling_params = [
|
||||
SamplingParams(max_tokens=output_len) for output_len in warmup_output_lens
|
||||
]
|
||||
llm.chat(
|
||||
warmup_prompts,
|
||||
warmup_sampling_params,
|
||||
use_tqdm=not getattr(args, "disable_tqdm", False),
|
||||
)
|
||||
|
||||
print(f"Processing {len(prompts)} requests...")
|
||||
start_time = time.perf_counter()
|
||||
|
||||
@@ -211,7 +172,9 @@ def benchmark_multimodal_processor(
|
||||
end_time = time.perf_counter()
|
||||
total_time = end_time - start_time
|
||||
|
||||
mm_stats_by_stage = collect_mm_processor_stats(llm.llm_engine, num_warmups)
|
||||
mm_stats_by_stage = collect_mm_processor_stats(
|
||||
llm.llm_engine,
|
||||
)
|
||||
|
||||
if not any(mm_stats_by_stage.values()):
|
||||
print(
|
||||
@@ -233,23 +196,17 @@ def benchmark_multimodal_processor(
|
||||
if not output.finished or output.metrics is None:
|
||||
continue
|
||||
metrics = output.metrics
|
||||
# Calculate E2E latency as: TTFT + (last_token_ts - first_token_ts)
|
||||
if (
|
||||
getattr(metrics, "first_token_latency", None) is not None
|
||||
and getattr(metrics, "last_token_ts", None) is not None
|
||||
and getattr(metrics, "first_token_ts", None) is not None
|
||||
):
|
||||
ttft = metrics.first_token_latency
|
||||
# Decode time is the duration between the first and last token generation
|
||||
decode_time = max(0.0, metrics.last_token_ts - metrics.first_token_ts)
|
||||
e2el_times.append((ttft + decode_time) * 1000)
|
||||
for attr in ("finished_time", "last_token_time"):
|
||||
if (
|
||||
getattr(metrics, attr, None) is not None
|
||||
and getattr(metrics, "arrival_time", None) is not None
|
||||
):
|
||||
e2el_times.append(
|
||||
(getattr(metrics, attr) - metrics.arrival_time) * 1000
|
||||
)
|
||||
break
|
||||
|
||||
if not e2el_times and completed > 0:
|
||||
print(
|
||||
"\n⚠️ Warning: Detailed end-to-end latency metrics not available.\n"
|
||||
" Falling back to average request latency "
|
||||
"(total_time / num_completed_requests).\n"
|
||||
)
|
||||
avg_time_per_request = total_time / completed
|
||||
e2el_times = [avg_time_per_request * 1000] * completed
|
||||
|
||||
@@ -266,17 +223,6 @@ def benchmark_multimodal_processor(
|
||||
std_e2el_ms = 0.0
|
||||
percentiles_e2el_ms = [(p, 0.0) for p in selected_percentiles]
|
||||
|
||||
encoder_summary = {}
|
||||
if (
|
||||
"num_encoder_calls" in mm_stats_by_stage
|
||||
and mm_stats_by_stage["num_encoder_calls"]
|
||||
):
|
||||
encoder_calls = mm_stats_by_stage["num_encoder_calls"]
|
||||
encoder_summary = {
|
||||
"total_encoder_calls": int(sum(encoder_calls)),
|
||||
"num_requests_with_encoder_calls": len(encoder_calls),
|
||||
}
|
||||
|
||||
benchmark_result = {
|
||||
"completed": completed,
|
||||
"failed": failed,
|
||||
@@ -285,7 +231,6 @@ def benchmark_multimodal_processor(
|
||||
"std_e2el_ms": std_e2el_ms,
|
||||
"percentiles_e2el_ms": percentiles_e2el_ms,
|
||||
"mm_processor_stats": mm_processor_metrics,
|
||||
"encoder_summary": encoder_summary,
|
||||
}
|
||||
|
||||
return benchmark_result
|
||||
@@ -303,7 +248,7 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
|
||||
"--dataset-name",
|
||||
type=str,
|
||||
default="random-mm",
|
||||
choices=["random-mm", "hf"],
|
||||
choices=["random-mm", "random-rerank"],
|
||||
help="Name of the dataset to benchmark on. Defaults to 'random-mm'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -312,12 +257,6 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
|
||||
default=10,
|
||||
help="Number of prompts to process.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-warmups",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of warmup prompts to process.",
|
||||
)
|
||||
|
||||
from vllm.benchmarks.datasets import (
|
||||
add_random_dataset_base_args,
|
||||
@@ -327,34 +266,6 @@ def add_cli_args(parser: argparse.ArgumentParser) -> None:
|
||||
add_random_dataset_base_args(parser)
|
||||
add_random_multimodal_dataset_args(parser)
|
||||
|
||||
# HuggingFace dataset arguments
|
||||
parser.add_argument(
|
||||
"--dataset-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to the dataset file or HuggingFace dataset name "
|
||||
"(e.g., 'yale-nlp/MMVU', 'lmarena-ai/VisionArena-Chat').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-subset",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Subset of the HuggingFace dataset (optional).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-split",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Split of the HuggingFace dataset (e.g., 'train', 'test', 'validation').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Output length for each request. "
|
||||
"Overrides the default output lengths from the dataset.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--output-json",
|
||||
type=str,
|
||||
@@ -385,17 +296,14 @@ def main(args: argparse.Namespace) -> None:
|
||||
print("=" * 80)
|
||||
|
||||
if "mm_processor_stats" in result:
|
||||
print("\nMM Processor Metrics:")
|
||||
print("\nMM Processor Timing (ms):")
|
||||
selected_percentiles = [
|
||||
float(p) for p in getattr(args, "metric_percentiles", "99").split(",")
|
||||
]
|
||||
mm_data = []
|
||||
for stage, metrics in result["mm_processor_stats"].items():
|
||||
is_count = stage == "num_encoder_calls"
|
||||
unit = "" if is_count else " (ms)"
|
||||
|
||||
row = {
|
||||
"Stage": stage + unit,
|
||||
"Stage": stage,
|
||||
"Mean": f"{metrics['mean']:.2f}",
|
||||
"Median": f"{metrics['median']:.2f}",
|
||||
"Std": f"{metrics['std']:.2f}",
|
||||
@@ -407,14 +315,6 @@ def main(args: argparse.Namespace) -> None:
|
||||
mm_df = pd.DataFrame(mm_data)
|
||||
print(mm_df.to_string(index=False))
|
||||
|
||||
if "encoder_summary" in result and result["encoder_summary"]:
|
||||
total_calls = result["encoder_summary"]["total_encoder_calls"]
|
||||
num_requests = result["encoder_summary"]["num_requests_with_encoder_calls"]
|
||||
print(
|
||||
f"\nSummary: {total_calls} total encoder calls "
|
||||
f"across {num_requests} requests."
|
||||
)
|
||||
|
||||
if "mean_e2el_ms" in result:
|
||||
print("\nEnd-to-End Latency (ms):")
|
||||
selected_percentiles = [
|
||||
|
||||
@@ -152,7 +152,6 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
|
||||
input=result_silu_mul,
|
||||
output_scale=output_scale,
|
||||
input_scale=scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
return at[1], at[2]
|
||||
|
||||
|
||||
@@ -946,7 +946,6 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
|
||||
input=rms,
|
||||
output_scale=output_scale,
|
||||
input_scale=input_global_scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
|
||||
# quant_out, allreduce_output, output_scale
|
||||
@@ -1044,7 +1043,6 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
|
||||
input=rms,
|
||||
output_scale=output_scale,
|
||||
input_scale=input_global_scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
|
||||
# quant_out, allreduce_output, output_scale
|
||||
|
||||
@@ -248,7 +248,6 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
|
||||
input=attn_out_view,
|
||||
output_scale=output_scale,
|
||||
input_scale=input_scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
)
|
||||
output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
|
||||
return at2[1], output_scale_view
|
||||
|
||||
+9
-11
@@ -1339,9 +1339,10 @@ class ModelConfig:
|
||||
Returns:
|
||||
A dictionary containing the non-default sampling parameters.
|
||||
"""
|
||||
src = self.generation_config
|
||||
|
||||
config = {} if src == "vllm" else self.try_get_generation_config()
|
||||
if self.generation_config == "vllm":
|
||||
config = {}
|
||||
else:
|
||||
config = self.try_get_generation_config()
|
||||
|
||||
# Overriding with given generation config
|
||||
config.update(self.override_generation_config)
|
||||
@@ -1367,16 +1368,13 @@ class ModelConfig:
|
||||
else:
|
||||
diff_sampling_param = {}
|
||||
|
||||
if diff_sampling_param and src != "vllm":
|
||||
if diff_sampling_param:
|
||||
logger.warning_once(
|
||||
"Default vLLM sampling parameters have been overridden by %s: `%s`. "
|
||||
"If this is not intended, please relaunch vLLM instance "
|
||||
"with `--generation-config vllm`.",
|
||||
"the model's `generation_config.json`" if src == "auto" else src,
|
||||
str(diff_sampling_param),
|
||||
scope="local",
|
||||
"Default sampling parameters have been overridden by the "
|
||||
"model's Hugging Face generation config recommended from the "
|
||||
"model creator. If this is not intended, please relaunch "
|
||||
"vLLM instance with `--generation-config vllm`."
|
||||
)
|
||||
|
||||
return diff_sampling_param
|
||||
|
||||
@property
|
||||
|
||||
@@ -213,8 +213,7 @@ class MultiModalConfig:
|
||||
factors: list[Any] = [
|
||||
self.mm_encoder_attn_backend.name
|
||||
if self.mm_encoder_attn_backend is not None
|
||||
else None,
|
||||
self.mm_encoder_tp_mode,
|
||||
else None
|
||||
]
|
||||
hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
|
||||
return hash_str
|
||||
|
||||
@@ -263,12 +263,6 @@ class VllmConfig:
|
||||
vllm_factors.append(__version__)
|
||||
if self.model_config:
|
||||
vllm_factors.append(self.model_config.compute_hash())
|
||||
if (
|
||||
self.compilation_config
|
||||
and getattr(self.compilation_config, "compile_mm_encoder", False)
|
||||
and self.model_config.multimodal_config
|
||||
):
|
||||
vllm_factors.append(self.model_config.multimodal_config.compute_hash())
|
||||
else:
|
||||
vllm_factors.append("None")
|
||||
if self.cache_config:
|
||||
|
||||
@@ -72,8 +72,7 @@ class ncclDataTypeEnum:
|
||||
ncclFloat64 = 8
|
||||
ncclDouble = 8
|
||||
ncclBfloat16 = 9
|
||||
ncclFloat8e4m3 = 10
|
||||
ncclNumTypes = 11
|
||||
ncclNumTypes = 10
|
||||
|
||||
@classmethod
|
||||
def from_torch(cls, dtype: torch.dtype) -> int:
|
||||
@@ -93,12 +92,9 @@ class ncclDataTypeEnum:
|
||||
return cls.ncclFloat64
|
||||
if dtype == torch.bfloat16:
|
||||
return cls.ncclBfloat16
|
||||
if dtype == torch.float8_e4m3fn:
|
||||
return cls.ncclFloat8e4m3
|
||||
raise ValueError(
|
||||
f"Unsupported dtype {dtype}: should be one of "
|
||||
f"int8, uint8, int32, int64, float16, float32, float64, bfloat16,"
|
||||
" float8e4m3."
|
||||
f"int8, uint8, int32, int64, float16, float32, float64, bfloat16."
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -316,12 +316,13 @@ class TpKVTopology:
|
||||
attn_backend: type[AttentionBackend]
|
||||
engine_id: EngineId
|
||||
remote_block_size: dict[EngineId, int]
|
||||
tensor_shape: torch.Size | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
# Figure out whether the first dimension of the cache is K/V
|
||||
# or num_blocks. This is used to register the memory regions correctly.
|
||||
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
num_blocks=1, block_size=16, num_kv_heads=4, head_size=1
|
||||
)
|
||||
# Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
|
||||
# we just mock num_blocks to 1 for the dimension check below.
|
||||
@@ -329,6 +330,32 @@ class TpKVTopology:
|
||||
len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
|
||||
)
|
||||
|
||||
self._kv_heads_position: int | None = None
|
||||
self._cross_layers_blocks = False
|
||||
if self.tensor_shape is not None:
|
||||
self._cross_layers_blocks = (
|
||||
len(self.tensor_shape) == len(kv_cache_shape) + 1
|
||||
)
|
||||
|
||||
if self._cross_layers_blocks:
|
||||
# prepend layers dimension
|
||||
kv_cache_shape = (80,) + kv_cache_shape
|
||||
try:
|
||||
kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(
|
||||
include_num_layers_dimension=self._cross_layers_blocks
|
||||
)
|
||||
except (AttributeError, NotImplementedError):
|
||||
kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
|
||||
|
||||
# permute kv_cache_shape according to stride_order
|
||||
kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
|
||||
|
||||
physical_block_size_position = kv_cache_shape.index(16)
|
||||
assert physical_block_size_position is not None
|
||||
self._physical_block_size_position = -(
|
||||
len(kv_cache_shape) - physical_block_size_position
|
||||
)
|
||||
|
||||
@property
|
||||
def is_kv_layout_blocks_first(self) -> bool:
|
||||
return self._is_kv_layout_blocks_first
|
||||
@@ -336,7 +363,9 @@ class TpKVTopology:
|
||||
@property
|
||||
def split_k_and_v(self) -> bool:
|
||||
# Whether to register regions for K and V separately (when present).
|
||||
return not (self.is_mla or self.is_kv_layout_blocks_first)
|
||||
return not (
|
||||
self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
|
||||
)
|
||||
|
||||
@property
|
||||
def tp_size(self) -> int:
|
||||
@@ -346,6 +375,14 @@ class TpKVTopology:
|
||||
def block_size(self) -> int:
|
||||
return self.remote_block_size[self.engine_id]
|
||||
|
||||
@property
|
||||
def cross_layers_blocks(self) -> bool:
|
||||
return self._cross_layers_blocks
|
||||
|
||||
@property
|
||||
def block_size_position(self) -> int:
|
||||
return self._physical_block_size_position
|
||||
|
||||
def tp_ratio(
|
||||
self,
|
||||
remote_tp_size: int,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user