Compare commits

..
Author SHA1 Message Date
Kevin H. LuuandGitHub b2e9eff8ae Update upload-release-wheels.sh
Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
2026-01-23 13:36:25 -08:00
216 changed files with 2375 additions and 7767 deletions
-84
View File
@@ -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"
+5 -12
View File
@@ -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
+18 -56
View File
@@ -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
+4 -5
View File
@@ -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
+1 -2
View File
@@ -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
View File
@@ -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}")
+22 -55
View File
@@ -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])
+1 -1
View File
@@ -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,
+5 -5
View File
@@ -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)
)
-12
View File
@@ -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)
+2 -2
View File
@@ -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 \
+4 -4
View File
@@ -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
View File
@@ -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));
});
+4 -4
View File
@@ -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);
}
}
+3 -6
View File
@@ -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");
}
+46 -140
View File
@@ -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));
});
}
+25 -171
View File
@@ -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);
}
+3 -4
View File
@@ -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
View File
@@ -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
View File
@@ -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 \
+2 -2
View File
@@ -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.
-1
View File
@@ -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
View File
@@ -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.
+9
View File
@@ -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:
+1 -1
View File
@@ -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
View File
@@ -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 projectwide, or accepted RFCs that shape technical direction.
- **Sustained contributions**: highquality 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 projectwide, or accepted RFCs that shape technical direction.
* **Sustained contributions**: highquality 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
- Longterm ownership of a subsystem with demonstrable quality and stability gains
- Significant crossproject 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
* Longterm ownership of a subsystem with demonstrable quality and stability gains
* Significant crossproject 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 projectwide
* 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 projectwide
### Working Groups
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -2
View File
@@ -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. | ✅︎ | ✅︎ |
-2
View File
@@ -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
View File
@@ -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
+2 -2
View File
@@ -9,5 +9,5 @@ wheel
jinja2>=3.1.6
regex
build
protobuf
grpcio-tools
protobuf>=6.33.2
grpcio-tools>=1.76.0
+3 -3
View File
@@ -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
+5 -50
View File
@@ -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(
-37
View File
@@ -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=(
+23 -103
View File
@@ -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
+18 -36
View File
@@ -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"
@@ -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)
+14 -3
View File
@@ -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 = {
-5
View File
@@ -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),
},
}
)
+2 -2
View File
@@ -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
View File
@@ -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
)
-12
View File
@@ -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(
+7 -23
View File
@@ -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)
-656
View File
@@ -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:
+178 -55
View File
@@ -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": {
+1 -5
View File
@@ -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."
-1
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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]
-2
View File
@@ -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
-1
View File
@@ -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
View File
@@ -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
+1 -2
View File
@@ -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
-6
View File
@@ -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