Compare commits

..
Author SHA1 Message Date
NickLucche e0dd86f6c3 deprecate timeout
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-04-28 09:03:27 +02:00
NickLucche 01f19dee7f deprecate timeout
Signed-off-by: NickLucche <nlucches@redhat.com>
2026-04-28 09:01:01 +02:00
Moritz SanftandGitHub 2c06cf3486 [Bugfix] use served_model_name for multimodal error message (#41003)
Signed-off-by: Moritz Sanft <58110325+msanft@users.noreply.github.com>
2026-04-27 08:22:35 -07:00
Harry MellorGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
e6f710a87f Deprecate support for Transformers v4 (#40389)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-04-27 08:19:57 -07:00
c245d35ff4 [Model] Add MiMo-V2.5 support (#40967)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: zjy0516 <zhujiangyun@inferact.ai>
Co-authored-by: yasong <yasong.wang@inferact.ai>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Copilot <copilot@github.com>
2026-04-27 13:26:51 +00:00
Xiaoshuang WangGitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
f8ac0c7cf0 [Bugfix] Fix k_norm weight sharding in MiniMaxM2Attention when total_num_kv_heads < tp_size (#38191)
Signed-off-by: wxsIcey <1790571317@qq.com>
Signed-off-by: Icey <1790571317@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-04-27 05:57:13 -07:00
ebf862c351 Add system_fingerprint field to OpenAI-compatible API responses (#40537)
Co-authored-by: Claude <noreply@anthropic.com>
2026-04-27 16:17:52 +08:00
wang.yuqiandGitHub 8d8062d0a7 [Examples] Resettle generate examples. (#36464)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
2026-04-27 07:48:37 +00:00
985961345a [Bugfix] Install libcublas-dev in Dockerfile for FlashInfer CuTe DSL JIT (#39855)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
2026-04-27 15:47:39 +08:00
Yongye ZhuandGitHub 706a04d34b [DSV4] Add silu clamp limit to shared expert (#40950)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
2026-04-27 00:37:43 -07:00
Isotr0pyandGitHub 22631f80a0 [Bugfix] Remove invalid deepstack boundary check for Qwen3-VL (#40932)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
2026-04-27 07:27:06 +00:00
BhoomitandGitHub 2cc008e7b4 [Attention][TurboQuant] Share dequant buffers, eliminate float16_copy (#40941)
Signed-off-by: Bhoomit Vasani <bhoomit.2010@gmail.com>
Signed-off-by: Vasani Bhoomit <bhoomit.2010@gmail.com>
2026-04-27 13:48:36 +08:00
5d5c776444 [Perf] FP8 FlashInfer Attn for ViT (#38065)
Signed-off-by: Zhanda Zhu <zhandazhu@gmail.com>
Co-authored-by: Yubo Gao <ybgao-nvidia@users.noreply.github.com>
2026-04-27 13:44:15 +08:00
ojhaanshikaandGitHub 592ae6805c Cutlass W4A16 (Machete) Tests (#35450)
Signed-off-by: Anshika Ojha <anshikao@nvidia.com>
2026-04-27 05:15:29 +00:00
7b1bc0a3eb [Bugfix] Cap SWA/chunked-local runtime admission to startup pool-sizing bound (#40946)
Signed-off-by: Dao Le <Dao007forever@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
2026-04-27 04:33:13 +00:00
Silu PandaandGitHub c0879d9483 [Tests] Gate Isaac under Transformers v5 (#40907)
Signed-off-by: Silu Panda <31051721+SiluPanda@users.noreply.github.com>
2026-04-26 19:26:51 -07:00
Giancarlo DelfinandGitHub f5f9878514 [Model Runner V2] Fix rejection sampling acceptance rate gap vs MRV1 (#40651)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
2026-04-26 19:12:08 -07:00
youkaichaoGitHubClaudegemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2ce95a761b Auto-disable expandable_segments around cumem memory pool (#40812)
Signed-off-by: youkaichao <youkaichao@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-04-27 09:37:22 +08:00
169 changed files with 8047 additions and 2286 deletions
+8 -156
View File
@@ -121,25 +121,7 @@ steps:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
@@ -152,25 +134,7 @@ steps:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
- label: "Build release image - x86_64 - CUDA 12.9"
@@ -180,24 +144,7 @@ steps:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
@@ -210,24 +157,7 @@ steps:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
@@ -237,27 +167,7 @@ steps:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
@@ -269,27 +179,7 @@ steps:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
@@ -299,26 +189,7 @@ steps:
queue: cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
@@ -330,26 +201,7 @@ steps:
queue: arm64_cpu_queue_release
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
- |
DOCKER_BUILDKIT=1 docker buildx build \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 \
--target vllm-openai \
--progress plain \
--load \
-f docker/Dockerfile .
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- block: "Build release image for x86_64 CPU"
+14 -14
View File
@@ -388,10 +388,10 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# Multi-modal models
- python3 offline_inference/audio_language.py --seed 0
- python3 offline_inference/vision_language.py --seed 0
- python3 offline_inference/vision_language_multi_image.py --seed 0
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
# Pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# Features demo
@@ -1647,10 +1647,10 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# Multi-modal models
- python3 offline_inference/audio_language.py --seed 0
- python3 offline_inference/vision_language.py --seed 0
- python3 offline_inference/vision_language_multi_image.py --seed 0
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
# Pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# Features demo
@@ -1951,8 +1951,8 @@ steps:
- pytest -v -s tests/models/multimodal/processing/
- pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/basic/offline_inference/chat.py
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
@@ -2930,10 +2930,10 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# Multi-modal models
- python3 offline_inference/audio_language.py --seed 0
- python3 offline_inference/vision_language.py --seed 0
- python3 offline_inference/vision_language_multi_image.py --seed 0
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
# Pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# Features demo
+2
View File
@@ -95,11 +95,13 @@ steps:
- tests/kernels/moe/test_deepgemm.py
- tests/kernels/moe/test_batched_deepgemm.py
- tests/kernels/attention/test_deepgemm_attention.py
- tests/quantization/test_cutlass_w4a16.py
commands:
- pytest -v -s kernels/quantization/test_block_fp8.py
- pytest -v -s kernels/moe/test_deepgemm.py
- pytest -v -s kernels/moe/test_batched_deepgemm.py
- pytest -v -s kernels/attention/test_deepgemm_attention.py
- pytest -v -s quantization/test_cutlass_w4a16.py
- label: Kernels (B200)
timeout_in_minutes: 30
+4 -4
View File
@@ -113,10 +113,10 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# for multi-modal models
- python3 offline_inference/audio_language.py --seed 0
- python3 offline_inference/vision_language.py --seed 0
- python3 offline_inference/vision_language_multi_image.py --seed 0
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
# for pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# for features demo
+4 -4
View File
@@ -44,10 +44,10 @@ steps:
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
#- python3 basic/offline_inference/embed.py # TODO
# for multi-modal models
- python3 offline_inference/audio_language.py --seed 0
- python3 offline_inference/vision_language.py --seed 0
- python3 offline_inference/vision_language_multi_image.py --seed 0
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
- python3 generate/multimodal/audio_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_offline.py --seed 0
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
# for pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# for features demo
+5 -5
View File
@@ -69,9 +69,9 @@ steps:
- pytest -v -s tests/models/multimodal/processing/
- pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/basic/offline_inference/chat.py
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
- label: Transformers Backward Compatibility Models Test
working_dir: "/vllm-workspace/"
@@ -83,7 +83,7 @@ steps:
- pytest -v -s tests/models/test_transformers.py
- pytest -v -s tests/models/multimodal/processing/
- pytest -v -s tests/models/multimodal/test_mapping.py
- python3 examples/offline_inference/basic/chat.py
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
- python3 examples/basic/offline_inference/chat.py
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
+1 -5
View File
@@ -389,11 +389,7 @@ pull_request_rules:
- files~=^tests/entrypoints/anthropic/.*tool.*
- files~=^vllm/tool_parsers/
- files=docs/features/tool_calling.md
- files~=^examples/tool_chat_*
- files=examples/offline_inference/chat_with_tools.py
- files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
- files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
- files=examples/online_serving/openai_chat_completion_client_with_tools.py
- files~=^examples/tool_calling/
actions:
label:
add:
@@ -0,0 +1,324 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
#
# == Usage Examples ==
#
# Benchmark mode (default, FlashInfer CUDAGraph Bench)
# python3 benchmark_vit_fp8_attn.py
#
# Profile mode (PyTorch profiler, saves TensorBoard traces):
# python3 benchmark_vit_fp8_attn.py --profile
# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
#
# Custom seq_lens:
# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
from functools import partial
import numpy as np
import torch
from torch.profiler import ProfilerActivity, profile, record_function
from vllm.utils.argparse_utils import FlexibleArgumentParser
# Qwen3-VL defaults
NUM_HEADS = 16
HEAD_DIM = 72
DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
"""Create FP8 and BF16 attention modules + workspace."""
from types import SimpleNamespace
from unittest.mock import patch
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.config.multimodal import MultiModalConfig
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
_get_flashinfer_workspace_buffer,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
old_dtype = torch.get_default_dtype()
torch.set_default_dtype(torch.bfloat16)
backend_patch = patch(
"vllm.model_executor.layers.attention.mm_encoder_attention"
".get_vit_attn_backend",
return_value=AttentionBackendEnum.FLASHINFER,
)
# FP8 attention
mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
vllm_config_fp8 = VllmConfig()
vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
with set_current_vllm_config(vllm_config_fp8), backend_patch:
attn_fp8 = MMEncoderAttention(
num_heads=num_heads,
head_size=head_dim,
prefix="visual.blocks.0.attn",
).to("cuda")
# BF16 attention (no FP8)
with set_current_vllm_config(VllmConfig()), backend_patch:
attn_bf16 = MMEncoderAttention(
num_heads=num_heads,
head_size=head_dim,
prefix="visual.blocks.0.attn",
).to("cuda")
torch.set_default_dtype(old_dtype)
workspace = _get_flashinfer_workspace_buffer()
return attn_fp8, attn_bf16, workspace
def _build_meta(
seq_len: int,
num_heads: int,
head_dim: int,
fp8: bool,
):
"""Build cu_seqlens, max_seqlen, sequence_lengths."""
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
from vllm.utils.math_utils import round_up
from vllm.v1.attention.backends.registry import AttentionBackendEnum
cu_np = np.array([0, seq_len], dtype=np.int32)
fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
)
max_seqlen = torch.tensor(
MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
dtype=torch.int32,
)
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
AttentionBackendEnum.FLASHINFER,
cu_np,
num_heads * head_dim,
1,
torch.device("cuda"),
fp8_padded_hidden_size=fp8_padded,
)
return cu_seqlens, max_seqlen, seq_lengths
def run_benchmark(
seq_lens: list[int],
num_heads: int,
head_dim: int,
method: str,
):
"""Benchmark FP8 vs BF16 attention across seq_lens.
Uses FlashInfer GPU-level timing to measure pure kernel time,
excluding CPU launch overhead.
"""
if method == "cupti":
from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
elif method == "cudagraph":
from flashinfer.testing import (
bench_gpu_time_with_cudagraph as bench_fn,
)
bench_fn = partial(bench_fn, cold_l2_cache=False)
else:
raise ValueError(f"Invalid method: {method}")
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
print(f"Timing method: {method}")
print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
print("-" * 46)
for seq_len in seq_lens:
torch.manual_seed(42)
q = torch.randn(
seq_len,
num_heads,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
k = torch.randn_like(q)
v = torch.randn_like(q)
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
cu_bf16 = cu_fp8.clone()
def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
# bench_fn returns List[float] of per-iteration times in ms
bf16_times = bench_fn(bf16_fn)
fp8_times = bench_fn(fp8_fn)
bf16_us = np.median(bf16_times) * 1e3 # ms -> us
fp8_us = np.median(fp8_times) * 1e3
speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
def _make_trace_handler(output_dir: str, worker_name: str, label: str):
"""Create a trace handler that saves to TensorBoard and prints summary."""
def handler(prof):
torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
print(f"\n{'=' * 80}")
print(label)
print(f"{'=' * 80}")
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
return handler
def run_profile(
seq_len: int,
num_heads: int,
head_dim: int,
warmup: int,
output_dir: str,
):
"""Profile FP8 vs BF16 attention with PyTorch profiler."""
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
torch.manual_seed(42)
q = torch.randn(
seq_len,
num_heads,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
k = torch.randn_like(q)
v = torch.randn_like(q)
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
cu_bf16 = cu_fp8.clone()
sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
# Profile BF16 (warmup handled by profiler schedule)
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=sched,
on_trace_ready=_make_trace_handler(
output_dir,
f"bf16_h{head_dim}_s{seq_len}",
f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
f"head_dim={head_dim})",
),
) as prof_bf16:
for _ in range(warmup + 1):
with record_function("bf16_attention"):
attn_bf16._forward_flashinfer(
q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
)
torch.accelerator.synchronize()
prof_bf16.step()
# Profile FP8 (warmup handled by profiler schedule)
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=sched,
on_trace_ready=_make_trace_handler(
output_dir,
f"fp8_h{head_dim}_s{seq_len}",
f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
f"head_dim={head_dim})",
),
) as prof_fp8:
for _ in range(warmup + 1):
with record_function("fp8_attention"):
attn_fp8._forward_flashinfer(
q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
)
torch.accelerator.synchronize()
prof_fp8.step()
print(f"\nTensorBoard traces saved to: {output_dir}")
print(f"View with: tensorboard --logdir={output_dir}")
if __name__ == "__main__":
parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
parser.add_argument(
"--seq-lens",
type=int,
nargs="+",
default=DEFAULT_SEQ_LENS,
help="Sequence lengths to benchmark",
)
parser.add_argument(
"--num-heads",
type=int,
default=NUM_HEADS,
)
parser.add_argument(
"--head-dim",
type=int,
default=HEAD_DIM,
)
parser.add_argument(
"--method",
choices=["cupti", "cudagraph"],
default="cudagraph",
help="GPU timing method: cupti (CUPTI kernel timing) or "
"cudagraph (CUDA graph capture/replay). Default: cudagraph",
)
parser.add_argument(
"--warmup",
type=int,
default=10,
help="Warmup iterations (profile mode only)",
)
parser.add_argument(
"--profile",
action="store_true",
help="Run PyTorch profiler instead of benchmark",
)
parser.add_argument(
"--profile-seq-len",
type=int,
default=8192,
help="Sequence length for profiling (default: 8192)",
)
parser.add_argument(
"--profile-output-dir",
type=str,
default="./profile_traces",
help="Output directory for TensorBoard traces (default: ./profile_traces)",
)
args = parser.parse_args()
if args.profile:
run_profile(
args.profile_seq_len,
args.num_heads,
args.head_dim,
args.warmup,
args.profile_output_dir,
)
else:
run_benchmark(
args.seq_lens,
args.num_heads,
args.head_dim,
args.method,
)
+4 -59
View File
@@ -82,73 +82,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
// Query occupancy for the instantiation that will actually launch;
// overestimating it deadlocks the cooperative barrier.
int occupancy = 1;
cudaError_t occ_err = cudaSuccess;
if (vec_size == 4) {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
smem_size);
} else if (vec_size == 2) {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
smem_size);
} else {
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
smem_size);
}
TORCH_CHECK(occ_err == cudaSuccess,
"persistent_topk occupancy query failed: ",
cudaGetErrorString(occ_err));
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
smem_size);
if (occupancy < 1) occupancy = 1;
// The cooperative spin-wait barrier only runs when at least one row hits
// the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
// early-exit, so oversubscription can't deadlock and headroom is wasted.
const bool needs_cooperative =
static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
const uint32_t hw_resident_cap =
static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
uint32_t max_resident_ctas = hw_resident_cap;
if (needs_cooperative) {
// Reserve one CTA per SM when occupancy allows; fall back to a single
// CTA when occupancy == 1 (the most deadlock-prone case — any straggler
// kernel that takes the only slot on one SM hangs the barrier). Never
// drop below one full group's worth.
uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
if (max_resident_ctas >= headroom + ctas_per_group) {
max_resident_ctas -= headroom;
}
}
uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
static_cast<uint32_t>(num_rows));
if (num_groups == 0) num_groups = 1;
uint32_t total_ctas = num_groups * ctas_per_group;
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
// instead of deadlocking. Only relevant when needs_cooperative.
if (needs_cooperative && total_ctas > hw_resident_cap) {
TORCH_CHECK(max_smem_per_block >= 128 * 1024,
"persistent_topk would oversubscribe and the FilteredTopK "
"fallback requires >=128KB smem per block (have ",
max_smem_per_block, "). total_ctas=", total_ctas,
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
", smem=", smem_size, ").");
cudaError_t status =
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
stream);
TORCH_CHECK(status == cudaSuccess,
"FilteredTopK fallback failed: ", cudaGetErrorString(status));
return;
}
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
"workspace too small, need ", state_bytes, " bytes");
+27 -34
View File
@@ -490,46 +490,41 @@ WORKDIR /vllm-workspace
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
# Install Python from source and system dependencies
# Install Python and system dependencies
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
software-properties-common \
curl \
sudo \
ffmpeg \
libsm6 \
libxext6 \
libgl1 \
&& if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
mkdir -p -m 0755 /etc/apt/keyrings ; \
curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
fi ; \
else \
for i in 1 2 3; do \
add-apt-repository -y ppa:deadsnakes/ppa && break || \
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
done ; \
fi \
&& apt-get update -y \
&& apt-get install -y --no-install-recommends \
python${PYTHON_VERSION} \
python${PYTHON_VERSION}-dev \
python${PYTHON_VERSION}-venv \
libibverbs-dev \
build-essential \
libssl-dev \
libffi-dev \
zlib1g-dev \
libbz2-dev \
libreadline-dev \
libsqlite3-dev \
libncurses-dev \
liblzma-dev \
libgdbm-dev \
uuid-dev \
tk-dev \
&& PYTHON_MAJOR_MINOR=${PYTHON_VERSION} \
&& PYTHON_FULL_VERSION=$(curl -s https://www.python.org/ftp/python/ \
| grep -oE "${PYTHON_MAJOR_MINOR}\.[0-9]+" \
| sort -t. -k3 -n | uniq \
| tail -1) \
&& echo "Building Python ${PYTHON_FULL_VERSION} from source..." \
&& curl -fSL https://www.python.org/ftp/python/${PYTHON_FULL_VERSION}/Python-${PYTHON_FULL_VERSION}.tgz -o /tmp/python.tgz \
&& tar -xzf /tmp/python.tgz -C /tmp \
&& cd /tmp/Python-${PYTHON_FULL_VERSION} \
&& ./configure --enable-optimizations --with-ensurepip=install --prefix=/usr/local \
&& make -j$(nproc) \
&& make install \
&& cd / \
&& rm -rf /tmp/python.tgz /tmp/Python-${PYTHON_FULL_VERSION} \
&& ln -sf /usr/local/bin/python${PYTHON_MAJOR_MINOR} /usr/bin/python3 \
&& ln -sf /usr/local/bin/python${PYTHON_MAJOR_MINOR}-config /usr/bin/python3-config \
&& ln -sf /usr/local/bin/pip${PYTHON_MAJOR_MINOR} /usr/bin/pip \
&& rm -rf /var/lib/apt/lists/* \
&& update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
&& update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
&& ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
&& rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
&& curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
&& python3 --version && python3 -m pip --version
# Install CUDA development tools for runtime JIT compilation
@@ -543,11 +538,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
cuda-nvrtc-${CUDA_VERSION_DASH} \
cuda-cuobjdump-${CUDA_VERSION_DASH} \
libcurand-dev-${CUDA_VERSION_DASH} \
libcublas-${CUDA_VERSION_DASH} \
libcublas-dev-${CUDA_VERSION_DASH} \
# Required by fastsafetensors (fixes #20384)
libnuma-dev \
# numactl CLI for NUMA binding at runtime
numactl && \
libnuma-dev && \
# Fixes nccl_allocator requiring nccl.h at runtime
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
+1 -1
View File
@@ -36,7 +36,7 @@ th {
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
!!! info "Table key"
1. All types: mxfp4, nvfp4, int4, int8, fp8
+7 -7
View File
@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
print(generated_text)
```
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
print(generated_text)
```
Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
!!! note
'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
### Audio Inputs
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
#### Chunking Long Audio for Transcription
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
print("Chat completion output:", chat_response.choices[0].message.content)
```
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
!!! tip
Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
print("Chat completion output from image url:", result)
```
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
!!! note
By default, the timeout for fetching videos through HTTP URL is `30` seconds.
@@ -958,7 +958,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
print("Chat completion output from audio url:", result)
```
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
!!! note
By default, the timeout for fetching audios through HTTP URL is `10` seconds.
+1
View File
@@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM:
- [AMD Quark](quark.md)
- [Quantized KV Cache](quantized_kvcache.md)
- [TorchAO](torchao.md)
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
## Supported Hardware
+109
View File
@@ -0,0 +1,109 @@
# FP8 ViT Encoder Attention
For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
short text prompts/generation, the ViT encoder attention can become a significant
bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
supports optional FP8 quantization for the ViT encoder attention via the
FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
cuDNN attention call.
!!! note
- Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
`qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
- Dynamic scaling is not compatible with ViT full CUDA graphs.
- Performance gains are mostly visible at QHD/4K resolutions or multi-image
requests. Smaller images may see no speedup due to quantization overhead
(3 quantization kernel launches + un-padding).
- FP8 tensor-core speedup is more pronounced on GB300 than GB200.
## Requirements
- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
## Usage
Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
with `--mm-encoder-attn-backend FLASHINFER`:
```bash
vllm serve $MODEL \
--mm-encoder-attn-backend FLASHINFER \
--mm-encoder-attn-dtype fp8
```
By default (no scale file), **dynamic scaling** is used: a 16-entry circular
buffer of observed Q/K/V amax values drives per-forward scale updates. This
matches BF16 accuracy without any calibration but adds a small per-forward
overhead.
## Calibrate-Once, Reuse Workflow (Recommended)
For production, calibrate static scales on a representative dataset once and
reuse them to avoid the dynamic overhead:
```bash
# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
# then dumps the learned scales to JSON).
vllm bench mm-processor \
--model $MODEL --mm-encoder-attn-backend FLASHINFER \
--mm-encoder-attn-dtype fp8 \
--mm-encoder-fp8-scale-save-path /path/to/scales.json \
--dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
--num-prompts 100
# Step 2: serve with static scales (no dynamic overhead).
vllm serve $MODEL \
--mm-encoder-attn-backend FLASHINFER \
--mm-encoder-attn-dtype fp8 \
--mm-encoder-fp8-scale-path /path/to/scales.json
```
Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
`1.5`) to leave headroom against activation outliers not present in the
calibration set. The default has been validated to generalize across datasets
(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
## Scale File Format
```json
{
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
}
```
Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
## Performance
**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
| Hardware | BF16 | FP8 | Speedup |
| -------- | ---- | ---- | ------- |
| GB200 | 350 us | 312 us | **1.12x** |
| GB300 | 300 us | 211 us | **1.42x** |
**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
| Resolution | BF16 median | FP8 median | Speedup |
| ---------- | ----------- | ---------- | ------- |
| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
## Accuracy
ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
on VisionArena-Chat (with default 1.5x margin):
| Metric | BF16 | FP8 dynamic | FP8 static |
| ------ | ---- | ----------- | ---------- |
| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
| exact_match | 0.584 | 0.582 | 0.578 |
All three configurations match within statistical noise, confirming that
static scales calibrated on one dataset generalize to another.
+1 -1
View File
@@ -202,7 +202,7 @@ The reasoning content is also available when both tool calling and the reasoning
print(f"Arguments: {tool_call.arguments}")
```
For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
## Server-Level Default Chat Template Kwargs
+3 -3
View File
@@ -292,10 +292,10 @@ Pooling models now support token-wise task.
### Score task
`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
equal to 1 can it be used as a scoring model and have its scoring API enabled.
### Pooling multitask support
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
+2
View File
@@ -439,6 +439,7 @@ th {
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
| `MiMoV2ProForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
@@ -590,6 +591,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
+3 -3
View File
@@ -251,7 +251,7 @@ The following extra parameters are supported:
Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
#### Extra parameters
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
!!! note
To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
!!! note
To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
#### Extra Parameters
@@ -6,15 +6,15 @@ This folder provides several example scripts on how to inference Qwen2.5-Omni of
```bash
# Audio + image + video
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
-q mixed_modalities
# Read vision and audio inputs from a single video file
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
-q use_audio_in_video
# Multiple audios
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
-q multi_audios
```
@@ -24,16 +24,16 @@ You can also test Qwen2.5-Omni on a single modality:
```bash
# Process audio inputs
python examples/offline_inference/audio_language.py \
python examples/generate/multimodal/audio_language_offline.py \
--model-type qwen2_5_omni
# Process image inputs
python examples/offline_inference/vision_language.py \
python examples/generate/multimodal/vision_language_offline.py \
--modality image \
--model-type qwen2_5_omni
# Process video inputs
python examples/offline_inference/vision_language.py \
python examples/generate/multimodal/vision_language_offline.py \
--modality video \
--model-type qwen2_5_omni
```
@@ -1402,7 +1402,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
# MiniCPM-V
def run_minicpmv_base(questions: list[str], modality: str, model_name):
assert modality in ["image", "video", "image+video"]
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language_offline.py` # noqa
# 2.0
# The official repo doesn't work yet, so we need to use a fork for now
@@ -4,68 +4,74 @@
import torch
from vllm import LLM
from vllm.config import PoolerConfig
from vllm.inputs import TextPrompt
from vllm.multimodal.utils import fetch_image
# Initialize model
model = LLM(
model="jinaai/jina-embeddings-v4-vllm-text-matching",
runner="pooling",
max_model_len=1024,
gpu_memory_utilization=0.8,
)
# Create text prompts
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
def main():
# Initialize model
model = LLM(
model="jinaai/jina-embeddings-v4-vllm-text-matching",
pooler_config=PoolerConfig(task="token_embed"),
runner="pooling",
max_model_len=1024,
gpu_memory_utilization=0.8,
)
text2 = "浜辺に沈む美しい夕日"
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
# Create text prompts
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
# Create image prompt
image = fetch_image(
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
)
image_prompt = TextPrompt(
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
multi_modal_data={"image": image},
)
text2 = "浜辺に沈む美しい夕日"
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
# Encode all prompts
prompts = [text1_prompt, text2_prompt, image_prompt]
outputs = model.encode(prompts, pooling_task="token_embed")
# Create image prompt
image = fetch_image(
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
)
image_prompt = TextPrompt(
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
multi_modal_data={"image": image},
)
# Encode all prompts
prompts = [text1_prompt, text2_prompt, image_prompt]
outputs = model.encode(prompts, pooling_task="token_embed")
def get_embeddings(outputs):
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
embeddings = []
for output in outputs:
if VISION_START_TOKEN_ID in output.prompt_token_ids:
# Gather only vision tokens
img_start_pos = torch.where(
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
)[0][0]
img_end_pos = torch.where(
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
)[0][0]
embeddings_tensor = output.outputs.data.detach().clone()[
img_start_pos : img_end_pos + 1
]
else:
# Use all tokens for text-only prompts
embeddings_tensor = output.outputs.data.detach().clone()
# Pool and normalize embeddings
pooled_output = (
embeddings_tensor.sum(dim=0, dtype=torch.float32)
/ embeddings_tensor.shape[0]
)
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
return embeddings
embeddings = get_embeddings(outputs)
for embedding in embeddings:
print(embedding.shape)
def get_embeddings(outputs):
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
embeddings = []
for output in outputs:
if VISION_START_TOKEN_ID in output.prompt_token_ids:
# Gather only vision tokens
img_start_pos = torch.where(
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
)[0][0]
img_end_pos = torch.where(
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
)[0][0]
embeddings_tensor = output.outputs.data.detach().clone()[
img_start_pos : img_end_pos + 1
]
else:
# Use all tokens for text-only prompts
embeddings_tensor = output.outputs.data.detach().clone()
# Pool and normalize embeddings
pooled_output = (
embeddings_tensor.sum(dim=0, dtype=torch.float32)
/ embeddings_tensor.shape[0]
)
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
return embeddings
embeddings = get_embeddings(outputs)
for embedding in embeddings:
print(embedding.shape)
if __name__ == "__main__":
main()
@@ -4,6 +4,7 @@
from argparse import Namespace
from vllm import LLM, EngineArgs
from vllm.config import PoolerConfig
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -13,6 +14,7 @@ def parse_args():
# Set example specific arguments
parser.set_defaults(
model="BAAI/bge-m3",
pooler_config=PoolerConfig(task="token_embed"),
runner="pooling",
enforce_eager=True,
)
@@ -32,15 +34,6 @@ def main(args: Namespace):
# You should pass runner="pooling" for embedding models
llm = LLM(**vars(args))
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
outputs = llm.embed(prompts)
# Print the outputs.
print("\nGenerated Outputs:\n" + "-" * 60)
for prompt, output in zip(prompts, outputs):
embeds = output.outputs.embedding
print(len(embeds))
# Generate embedding for each token. The output is a list of PoolingRequestOutput.
outputs = llm.encode(prompts, pooling_task="token_embed")
@@ -50,6 +43,20 @@ def main(args: Namespace):
multi_vector = output.outputs.data
print(multi_vector.shape)
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
# Generate scores.
outputs = llm.score(query, documents)
# Print the outputs.
print("\nGenerated Outputs:\n" + "-" * 60)
for document, output in zip(documents, outputs):
score = output.outputs.score
print(f"Pair: {[query, document]!r} \nScore: {score}")
print("-" * 60)
if __name__ == "__main__":
args = parse_args()
@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
Run `vllm serve <model> --runner pooling`
to start up the server in vLLM. e.g.
vllm serve BAAI/bge-m3
vllm serve BAAI/bge-m3 --pooler-config.task token_embed
"""
import argparse
import pprint
import requests
import torch
@@ -32,7 +33,8 @@ def parse_args():
def main(args):
api_url = f"http://{args.host}:{args.port}/pooling"
pooling_url = f"http://{args.host}:{args.port}/pooling"
score_url = f"http://{args.host}:{args.port}/score"
model_name = args.model
prompts = [
@@ -43,11 +45,23 @@ def main(args):
]
prompt = {"model": model_name, "input": prompts}
pooling_response = post_http_request(prompt=prompt, api_url=api_url)
pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
for output in pooling_response.json()["data"]:
multi_vector = torch.tensor(output["data"])
print(multi_vector.shape)
queries = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
prompt = {"model": model_name, "queries": queries, "documents": documents}
score_response = post_http_request(prompt=prompt, api_url=score_url)
print("\nPrompt when queries is string and documents is a list:")
pprint.pprint(prompt)
print("\nScore Response:")
pprint.pprint(score_response.json())
if __name__ == "__main__":
args = parse_args()
+2 -5
View File
@@ -34,10 +34,7 @@ def _run_vllm(vllm_runner):
mode=CompilationMode.VLLM_COMPILE,
cudagraph_mode=CUDAGraphMode.NONE,
),
# Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
# at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
# 32 for headroom.
num_gpu_blocks_override=32,
num_gpu_blocks_override=8,
):
pass
@@ -193,7 +190,7 @@ def _run_model(vllm_runner, spec: ModelStartupSpec):
cudagraph_mode=CUDAGraphMode.NONE,
pass_config=PassConfig(fuse_allreduce_rms=False),
),
num_gpu_blocks_override=16,
num_gpu_blocks_override=8,
):
pass
-3
View File
@@ -405,9 +405,6 @@ def test_should_split():
(None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0),
# truncated to nearest multiple of 8 or 16
(None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
# max_num_batched_tokens <= max_cudagraph_capture_size should always be
# captured even if not landing on a 16-stride step
(None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257),
# max from list
([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
# SP forces full-graph compilation, sizes are filtered by TP
+18
View File
@@ -41,3 +41,21 @@ def test_language_model_only_affects_model_hash():
base_hash = ModelConfig(model).compute_hash()
lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
assert base_hash != lm_only_hash
def test_mm_encoder_fp8_scale_path_requires_fp8():
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
scale_file = tmp_path / "scales.json"
scale_file.write_text("{}")
base_hash = MultiModalConfig().compute_hash()
fp8_hash = MultiModalConfig(mm_encoder_attn_dtype="fp8").compute_hash()
fp8_static_hash = MultiModalConfig(
mm_encoder_attn_dtype="fp8",
mm_encoder_fp8_scale_path=str(scale_file),
).compute_hash()
assert base_hash != fp8_hash
assert fp8_hash != fp8_static_hash
@@ -0,0 +1,76 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ``system_fingerprint`` construction."""
from types import SimpleNamespace
import pytest
from vllm.entrypoints.openai import fingerprint as fp
def _cfg(tp=1, pp=1, dp=1, ep=False, digest="a3b21f94deadbeef"):
c = SimpleNamespace(
parallel_config=SimpleNamespace(
tensor_parallel_size=tp,
pipeline_parallel_size=pp,
data_parallel_size=dp,
enable_expert_parallel=ep,
)
)
c.compute_hash = lambda: digest # type: ignore[attr-defined]
return c
@pytest.fixture(autouse=True)
def _reset():
fp.set_default_fingerprint_mode("full")
yield
fp.set_default_fingerprint_mode("full")
def test_four_modes_produce_expected_shapes():
from vllm import __version__ as v
cfg = _cfg(tp=8, ep=True)
assert fp.build_system_fingerprint(cfg, "full") == (f"vllm-{v}-tp8-ep-a3b21f94")
assert fp.build_system_fingerprint(cfg, "hash") == f"vllm-{v}-a3b21f94"
assert fp.build_system_fingerprint(cfg, "custom", "my-fp") == "my-fp"
assert fp.build_system_fingerprint(cfg, "none") is None
def test_full_mode_emits_only_non_trivial_parallelism():
from vllm import __version__ as v
# Single-GPU: nothing between version and hash.
assert fp.build_system_fingerprint(_cfg(), "full") == f"vllm-{v}-a3b21f94"
# All parallelism axes.
assert (
fp.build_system_fingerprint(_cfg(tp=8, pp=2, dp=4, ep=True), "full")
== f"vllm-{v}-tp8-pp2-dp4-ep-a3b21f94"
)
def test_get_respects_set_default():
cfg = _cfg(tp=8)
full = fp.get_system_fingerprint(cfg)
assert full == fp.get_system_fingerprint(cfg)
fp.set_default_fingerprint_mode("hash")
hashed = fp.get_system_fingerprint(cfg)
assert hashed != full
assert "tp8" not in hashed
fp.set_default_fingerprint_mode("custom", "deploy-42")
assert fp.get_system_fingerprint(cfg) == "deploy-42"
fp.set_default_fingerprint_mode("none")
assert fp.get_system_fingerprint(cfg) is None
def test_compute_hash_failure_does_not_raise():
cfg = _cfg()
cfg.compute_hash = lambda: (_ for _ in ()).throw(RuntimeError("boom"))
assert fp.build_system_fingerprint(cfg, "full").endswith("-nohash")
assert fp.build_system_fingerprint(cfg, "hash").endswith("-nohash")
@@ -1,13 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
import weakref
import pytest
import torch
from tests.models.utils import softmax
from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
from vllm import LLM, ClassificationRequestOutput, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.tasks import PoolingTask
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
assert len(outputs[i].outputs.probs) == num_labels
@pytest.mark.skip_global_cleanup
def test_token_classify(llm: LLM, caplog_vllm):
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
assert "deprecated" in caplog_vllm.text
assert len(outputs) == 1
assert isinstance(outputs[0], PoolingRequestOutput)
assert outputs[0].prompt_token_ids == prompt_token_ids
assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
@pytest.mark.skip_global_cleanup
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
llm.score("ping", "pong", use_tqdm=False)
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "token_classify":
err_msg = "Try switching the model's pooling_task via.+"
else:
err_msg = "Embedding API is not supported by this model.+"
with pytest.raises(ValueError, match=err_msg):
@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
task = "token_classify"
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == 8
assert len(poolings.data[0].data[0]) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
},
)
assert response.json()["error"]["type"] == "BadRequestError"
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "token_classify":
err_msg = "Try switching the model's pooling_task via"
else:
err_msg = f"Unsupported task: {task!r}"
assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
import weakref
import pytest
@@ -38,11 +37,11 @@ def llm():
seed=0,
attention_config=attention_config,
)
assert embedding_size == llm.model_config.embedding_size
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
assert len(outputs[i].outputs.embedding) == embedding_size
@pytest.mark.skip_global_cleanup
def test_token_embed(llm: LLM, caplog_vllm):
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
assert "deprecated" in caplog_vllm.text
multi_vector = outputs[0].outputs.data
assert multi_vector.shape == (11, 384)
@pytest.mark.skip_global_cleanup
def test_pooling_params(llm: LLM):
def get_outputs(normalize):
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
)
@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
@pytest.mark.parametrize(
"task", ["token_classify", "classify", "token_embed", "plugin"]
)
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "token_embed":
err_msg = "Try switching the model's pooling_task via.+"
else:
err_msg = "Classification API is not supported by this model.+"
with pytest.raises(ValueError, match=err_msg):
+5 -22
View File
@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
task = "token_embed"
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == len(input_tokens)
assert len(poolings.data[0].data[0]) == 384
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
@pytest.mark.parametrize(
"task", ["classify", "token_classify", "token_embed", "plugin"]
)
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
assert response.json()["error"]["type"] == "BadRequestError"
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "token_embed":
err_msg = "Try switching the model's pooling_task via"
else:
err_msg = f"Unsupported task: {task!r}"
assert response.json()["error"]["message"].startswith(err_msg)
@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
assert len(poolings.data[0].data) == 1
@pytest.mark.asyncio
async def test_pooling_token_classify(server: RemoteOpenAIServer):
response = requests.post(
server.url_for("pooling"),
json={
"model": MODEL_NAME,
"task": "token_classify",
"input": input_text,
"encoding_format": "float",
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == len(input_tokens)
assert len(poolings.data[0].data[0]) == 1
@pytest.mark.asyncio
async def test_rerank_max_tokens_per_doc(
server: RemoteOpenAIServer,
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
@pytest.mark.asyncio
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
response = requests.post(
server.url_for("pooling"),
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
assert response.json()["error"]["type"] == "BadRequestError"
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "token_classify":
err_msg = "Try switching the model's pooling_task via"
else:
err_msg = f"Unsupported task: {task!r}"
assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
import weakref
import pytest
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
@pytest.mark.skip_global_cleanup
def test_score_api(llm: LLM):
err_msg = "Scoring API is only enabled for num_labels == 1."
err_msg = "This model does not support the Scoring API."
with pytest.raises(ValueError, match=err_msg):
llm.score("ping", "pong", use_tqdm=False)
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
if task == "classify":
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
llm.encode(prompt, pooling_task=task, use_tqdm=False)
assert "deprecated" in caplog_vllm.text
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "classify":
err_msg = "Try switching the model's pooling_task via.+"
else:
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
else:
err_msg = "Embedding API is not supported by this model.+"
err_msg = "Embedding API is not supported by this model.+"
with pytest.raises(ValueError, match=err_msg):
llm.encode(prompt, pooling_task=task, use_tqdm=False)
with pytest.raises(ValueError, match=err_msg):
llm.encode(prompt, pooling_task=task, use_tqdm=False)
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
"task": task,
},
)
assert response.json()["error"]["type"] == "BadRequestError"
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "classify":
err_msg = "Try switching the model's pooling_task via"
else:
err_msg = f"Unsupported task: {task!r}"
assert response.json()["error"]["message"].startswith(err_msg)
@@ -1,6 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging
import weakref
import pytest
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
if task == "embed":
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
llm.encode(prompt, pooling_task=task, use_tqdm=False)
assert "deprecated" in caplog_vllm.text
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "embed":
err_msg = "Try switching the model's pooling_task via.+"
else:
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
else:
err_msg = "Classification API is not supported by this model.+"
err_msg = "Classification API is not supported by this model.+"
with pytest.raises(ValueError, match=err_msg):
llm.encode(prompt, pooling_task=task, use_tqdm=False)
with pytest.raises(ValueError, match=err_msg):
llm.encode(prompt, pooling_task=task, use_tqdm=False)
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
"task": task,
},
)
assert response.json()["error"]["type"] == "BadRequestError"
if task == "plugin":
err_msg = "No IOProcessor plugin installed."
elif task == "embed":
err_msg = "Try switching the model's pooling_task via"
else:
err_msg = f"Unsupported task: {task!r}"
assert response.json()["error"]["message"].startswith(err_msg)
+279
View File
@@ -0,0 +1,279 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
import contextlib
import pytest
import torch
from vllm.triton_utils import HAS_TRITON
from vllm.utils.flashinfer import (
is_flashinfer_cudnn_fp8_prefill_attn_supported,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
def _has_flashinfer_cudnn() -> bool:
"""Check if FlashInfer cuDNN backend is available."""
try:
from flashinfer.prefill import (
cudnn_batch_prefill_with_kv_cache, # noqa: F401
)
return True
except ImportError:
return False
HEAD_DIMS = [72, 80]
SEQ_LENS = [256]
NUM_HEADS = [16]
@pytest.fixture
def _fp8_attention():
"""Create FP8-enabled MMEncoderAttention via config."""
from types import SimpleNamespace
from unittest.mock import patch
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.config.multimodal import MultiModalConfig
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
vllm_config = VllmConfig()
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
# MMEncoderAttention reads torch.get_default_dtype() during init
# to determine the output dtype. In real model loading this is bf16.
old_dtype = torch.get_default_dtype()
torch.set_default_dtype(torch.bfloat16)
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.model_executor.layers.attention.mm_encoder_attention"
".get_vit_attn_backend",
return_value=AttentionBackendEnum.FLASHINFER,
),
):
yield
torch.set_default_dtype(old_dtype)
def _build_cu_seqlens_and_meta(
seq_len: int,
num_heads: int,
head_dim: int,
fp8_padded_hidden_size: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
import numpy as np
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
AttentionBackendEnum.FLASHINFER,
cu_seqlens_np,
torch.device("cuda"),
)
max_seqlen = torch.tensor(
MMEncoderAttention.compute_max_seqlen(
AttentionBackendEnum.FLASHINFER, cu_seqlens_np
),
dtype=torch.int32,
)
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
AttentionBackendEnum.FLASHINFER,
cu_seqlens_np,
num_heads * head_dim,
1, # tp_size
torch.device("cuda"),
fp8_padded_hidden_size=fp8_padded_hidden_size,
)
return cu_seqlens, max_seqlen, sequence_lengths
@pytest.mark.skipif(
not (HAS_TRITON and _has_flashinfer_cudnn()),
reason="Triton and FlashInfer cuDNN required",
)
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
@pytest.mark.parametrize("seq_len", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
def test_fp8_attn_output_shape(
head_dim: int,
seq_len: int,
num_heads: int,
_fp8_attention,
) -> None:
"""Verify FP8 attention produces correct output shape after un-padding."""
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
from vllm.utils.math_utils import round_up
attn = None
with contextlib.suppress(ValueError, ImportError):
attn = MMEncoderAttention(
num_heads=num_heads,
head_size=head_dim,
prefix="visual.blocks.0.attn",
).to("cuda")
if attn is None or not attn.fp8_enabled:
pytest.skip("FP8 MMEncoderAttention not available")
assert attn is not None # mypy narrowing
# FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
)
q = torch.randn(
seq_len,
num_heads,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
k = torch.randn_like(q)
v = torch.randn_like(q)
output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
# Output should have original head_dim (un-padded)
assert output.shape[-1] == head_dim
assert output.dtype == torch.bfloat16
@pytest.mark.skipif(
not (HAS_TRITON and _has_flashinfer_cudnn()),
reason="Triton and FlashInfer cuDNN required",
)
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
@pytest.mark.parametrize("seq_len", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
def test_fp8_vs_bf16_close(
head_dim: int, seq_len: int, num_heads: int, _fp8_attention
) -> None:
"""FP8 attention output should be reasonably close to BF16 baseline."""
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
from vllm.utils.math_utils import round_up
torch.manual_seed(42)
q = torch.randn(
1,
seq_len,
num_heads,
head_dim,
device="cuda",
dtype=torch.bfloat16,
)
k = torch.randn_like(q)
v = torch.randn_like(q)
# FP8 path
attn_fp8 = None
with contextlib.suppress(ValueError, ImportError):
attn_fp8 = MMEncoderAttention(
num_heads=num_heads,
head_size=head_dim,
prefix="visual.blocks.0.attn",
).to("cuda")
if attn_fp8 is None or not attn_fp8.fp8_enabled:
pytest.skip("FP8 MMEncoderAttention not available")
assert attn_fp8 is not None # mypy narrowing
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
seq_len,
num_heads,
head_dim,
fp8_padded_hidden_size=fp8_padded_hidden_size,
)
out_fp8 = attn_fp8._forward_flashinfer(
q.clone(),
k.clone(),
v.clone(),
cu_seqlens,
max_seqlen,
seq_lengths,
)
# BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
# and calling the wrapper directly without FP8 quantization)
from vllm.model_executor.layers.attention.mm_encoder_attention import (
_get_flashinfer_workspace_buffer,
)
from vllm.v1.attention.ops.vit_attn_wrappers import (
vit_flashinfer_wrapper,
)
out_bf16 = vit_flashinfer_wrapper(
q=q.clone(),
k=k.clone(),
v=v.clone(),
scale=attn_fp8.scale,
workspace_buffer=_get_flashinfer_workspace_buffer(),
cu_seqlens=cu_seqlens,
max_seqlen=max_seqlen,
sequence_lengths=seq_lengths,
)
out_fp8_f = out_fp8.float()
out_bf16_f = out_bf16.float()
abs_diff = (out_fp8_f - out_bf16_f).abs()
abs_diff_flat = abs_diff.flatten()
# Relative diff (avoid division by zero)
denom = out_bf16_f.abs().clamp(min=1e-6)
rel_diff_flat = (abs_diff / denom).flatten()
cosine_sim = torch.nn.functional.cosine_similarity(
out_fp8_f.flatten().unsqueeze(0),
out_bf16_f.flatten().unsqueeze(0),
).item()
pcts = [50, 90, 95, 99, 99.9]
abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
print(f" cosine_sim={cosine_sim:.6f}")
print(
f" abs_diff: max={abs_diff_flat.max().item():.6f}, "
f"mean={abs_diff_flat.mean().item():.6f}, "
+ ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
)
print(
f" rel_diff: max={rel_diff_flat.max().item():.6f}, "
f"mean={rel_diff_flat.mean().item():.6f}, "
+ ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
)
assert abs_diff_flat.max().item() < 0.3, (
f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
)
assert abs_diff_flat.mean().item() < 0.03, (
f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
)
assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"
+124
View File
@@ -0,0 +1,124 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
import pytest
import torch
from vllm.platforms import current_platform
from vllm.triton_utils import HAS_TRITON
if HAS_TRITON:
from vllm.kernels.triton.qkv_padded_fp8_quant import (
quantize_fp8_pad_head_dim_triton,
)
HEAD_DIMS = [72, 80, 128]
SEQ_LENS = [64, 256]
NUM_HEADS = [16]
SCALES = [0.01, 0.1, 1.0]
def _naive_fp8_quantize(
tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
) -> torch.Tensor:
"""Reference FP8 quantization in PyTorch."""
fp8_dtype = current_platform.fp8_dtype()
fp8_max = torch.finfo(fp8_dtype).max
fp8_min = -fp8_max
x = tensor.float()
if not skip_scale:
x = x / scale.item()
x = x.clamp(fp8_min, fp8_max)
return x.to(fp8_dtype)
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
@pytest.mark.parametrize("seq_len", SEQ_LENS)
@pytest.mark.parametrize("num_heads", NUM_HEADS)
@pytest.mark.parametrize("scale_val", SCALES)
def test_quantize_contiguous(
head_dim: int, seq_len: int, num_heads: int, scale_val: float
) -> None:
"""Test quantization of contiguous 3D tensors."""
torch.manual_seed(42)
tensor = torch.randn(
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
)
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
1, 1, 1, 1
)
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
padded_dim = (head_dim + 15) // 16 * 16
assert result.shape == (seq_len, num_heads, padded_dim)
assert result.is_contiguous()
assert result.dtype == current_platform.fp8_dtype()
# Compare unpadded portion against reference
ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
# Padded region should be zero
if padded_dim > head_dim:
assert (result[:, :, head_dim:].float() == 0).all()
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
@pytest.mark.parametrize("head_dim", [72, 80])
def test_quantize_non_contiguous(head_dim: int) -> None:
"""Test quantization from non-contiguous QKV views (interleaved buffer)."""
seq_len, num_heads = 64, 16
# Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
qkv = torch.randn(
seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
)
# Q is every 3rd head slice - non-contiguous view
q = qkv[:, 0::3, :]
assert not q.is_contiguous()
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
result = quantize_fp8_pad_head_dim_triton(q, scale)
padded_dim = (head_dim + 15) // 16 * 16
assert result.shape == (seq_len, num_heads, padded_dim)
assert result.is_contiguous()
# Compare against contiguous reference
ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
def test_skip_scale() -> None:
"""Test skip_scale=True produces cast-only output (no division)."""
seq_len, num_heads, head_dim = 32, 8, 80
tensor = torch.randn(
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
)
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
# skip_scale should just cast, not divide
ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
# With scale != 1.0, skip and no-skip should differ
assert not torch.equal(result_skip.float(), result_noskip.float())
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
def test_4d_input() -> None:
"""Test that 4D input (B, S, H, D) is handled correctly."""
B, S, H, D = 2, 32, 8, 72
tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
padded_dim = (D + 15) // 16 * 16
assert result.shape == (B, S, H, padded_dim)
+251
View File
@@ -0,0 +1,251 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
import contextlib
import json
from types import SimpleNamespace
from unittest.mock import patch
import pytest
import torch
from vllm.model_executor.layers.attention.mm_encoder_attention import (
_FP8_AMAX_HISTORY_LEN,
_FP8_MAX,
)
from vllm.utils.flashinfer import (
is_flashinfer_cudnn_fp8_prefill_attn_supported,
)
LAYER_0 = "visual.blocks.0.attn.attn"
LAYER_1 = "visual.blocks.1.attn.attn"
NUM_HEADS = 16
HEAD_DIM = 72
@contextlib.contextmanager
def _build_attention(mm_config):
"""Yield an MMEncoderAttention with the given multimodal config.
The VllmConfig context stays active while the test runs so that
``get_multimodal_config()`` calls during the forward path resolve. Also
invokes ``process_weights_after_loading`` to simulate the model loader's
auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
"""
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
yield None
return
vllm_config = VllmConfig()
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.model_executor.layers.attention.mm_encoder_attention"
".get_vit_attn_backend",
return_value=AttentionBackendEnum.FLASHINFER,
),
):
attn = MMEncoderAttention(
num_heads=NUM_HEADS,
head_size=HEAD_DIM,
prefix=LAYER_0,
)
attn.process_weights_after_loading(torch.bfloat16)
yield attn
@pytest.fixture
def _make_attention():
"""Create an MMEncoderAttention with dynamic FP8 scaling."""
from vllm.config.multimodal import MultiModalConfig
with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
yield attn
@pytest.fixture
def _make_static_attention(tmp_path):
"""Create an MMEncoderAttention with static FP8 scales from a file."""
from vllm.config.multimodal import MultiModalConfig
scale_file = tmp_path / "scales.json"
scale_file.write_text(
json.dumps(
{
LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
}
)
)
with _build_attention(
MultiModalConfig(
mm_encoder_attn_dtype="fp8",
mm_encoder_fp8_scale_path=str(scale_file),
)
) as attn:
yield attn
def test_dynamic_scaling_updates_scales(_make_attention) -> None:
"""Verify that _record_amax_and_update_scales updates scale buffers."""
attn = _make_attention
if attn is None or not attn.fp8_enabled:
pytest.skip("FP8 attention not available (FlashInfer backend required)")
attn = attn.to("cuda")
S, H, D = 32, NUM_HEADS, HEAD_DIM
q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
attn._record_amax_and_update_scales(q, k, v)
expected_q_scale = 2.0 / _FP8_MAX
expected_k_scale = 3.0 / _FP8_MAX
expected_v_scale = 4.0 / _FP8_MAX
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
def test_circular_buffer_wraps(_make_attention) -> None:
"""Verify the amax circular buffer wraps at HISTORY_LEN."""
attn = _make_attention
if attn is None or not attn.fp8_enabled:
pytest.skip("FP8 attention not available (FlashInfer backend required)")
attn = attn.to("cuda")
S, H, D = 16, NUM_HEADS, HEAD_DIM
for i in range(_FP8_AMAX_HISTORY_LEN + 2):
mag = float(i + 1)
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
attn._record_amax_and_update_scales(q, k, v)
assert attn._fp8_amax_pos == 2
expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
expected_scale = expected_max / _FP8_MAX
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
def test_static_scales_loaded(_make_static_attention) -> None:
"""Verify static scales are loaded from the JSON file."""
attn = _make_static_attention
if attn is None or not attn.fp8_enabled:
pytest.skip("FP8 attention not available (FlashInfer backend required)")
assert attn.fp8_enabled
assert not attn._fp8_dynamic_scale
# Layer 0 scales (the layer this attention was created with).
assert attn._fp8_q_scale.item() == 224.0
assert attn._fp8_k_scale.item() == 198.0
assert attn._fp8_v_scale.item() == 210.0
assert not attn.skip_scale_q
assert not attn.skip_scale_k
assert not attn.skip_scale_v
# No amax history buffers for static scaling.
assert not hasattr(attn, "_fp8_q_amax")
def test_static_scales_missing_layer(tmp_path) -> None:
"""Verify error when requested layer is not in the scale file."""
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.config.multimodal import MultiModalConfig
from vllm.v1.attention.backends.registry import AttentionBackendEnum
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
pytest.skip("FlashInfer cuDNN not available")
scale_file = tmp_path / "wrong_layer.json"
scale_file.write_text(
json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
)
mm_config = MultiModalConfig(
mm_encoder_attn_dtype="fp8",
mm_encoder_fp8_scale_path=str(scale_file),
)
vllm_config = VllmConfig()
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
from vllm.model_executor.layers.attention.mm_encoder_attention import (
MMEncoderAttention,
)
with (
set_current_vllm_config(vllm_config),
patch(
"vllm.model_executor.layers.attention.mm_encoder_attention"
".get_vit_attn_backend",
return_value=AttentionBackendEnum.FLASHINFER,
),
):
attn = MMEncoderAttention(
num_heads=NUM_HEADS,
head_size=HEAD_DIM,
prefix=LAYER_0,
)
with pytest.raises(ValueError, match="scales not found for layer"):
attn.process_weights_after_loading(torch.bfloat16)
def test_dynamic_scales_auto_save(tmp_path) -> None:
"""Verify scales are saved to disk after the amax buffer fills."""
import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
from vllm.config.multimodal import MultiModalConfig
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
pytest.skip("FlashInfer cuDNN not available")
# Reset module-level state between runs (other tests may have left
# state behind after triggering a save).
_mod._fp8_scale_save_path = None
_mod._fp8_saved_scale_refs.clear()
save_file = tmp_path / "auto_scales.json"
with _build_attention(
MultiModalConfig(
mm_encoder_attn_dtype="fp8",
mm_encoder_fp8_scale_save_path=str(save_file),
)
) as attn:
if attn is None or not attn.fp8_enabled:
pytest.skip("FP8 attention not available")
attn = attn.to("cuda")
S, H, D = 16, NUM_HEADS, HEAD_DIM
# Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
for i in range(_FP8_AMAX_HISTORY_LEN):
mag = float(i + 1)
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
attn._record_amax_and_update_scales(q, k, v)
# File should have been written on the 16th call (buffer wrap).
assert save_file.is_file(), "Scale file was not saved"
scales = json.loads(save_file.read_text())
assert LAYER_0 in scales
assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
for val in scales[LAYER_0].values():
assert isinstance(val, float) and val > 0
# Path is cleared after the one-shot save fires.
assert _mod._fp8_scale_save_path is None
+4 -228
View File
@@ -3,11 +3,12 @@
"""
Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
Four test functions cover five paths:
Two paths tested:
A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
B) Indexer: head_dim=128 (all FP8), quant_block=128
C) DeepseekV4 Attention magnitude range: correctness across small/large values
D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
These serve as golden references for validating the future fused
compressor+quant+cache kernel.
"""
import math
@@ -20,12 +21,6 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
dequantize_and_gather_k_cache,
quantize_and_insert_k_cache,
)
from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
_fused_kv_compress_norm_rope_insert_indexer_attn,
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
)
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
@@ -314,222 +309,3 @@ def test_deepseek_v4_quant_magnitude_range():
f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
f"magnitude={magnitude:.4f}"
)
# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
#
# Both kernels share the same Triton signature; use_fp4 selects between them.
# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
# GPT-J RoPE → quant (MXFP4 or FP8) → paged cache insert.
def _reference_kv_compress_norm_rope(
state_cache: torch.Tensor,
block_table: torch.Tensor,
positions: torch.Tensor,
rms_weight: torch.Tensor,
cos_sin_cache: torch.Tensor,
compress_ratio: int = 1,
overlap: int = 0,
use_fp4: bool = False,
rms_eps: float = 1e-6,
fp8_max: float = 448.0,
):
"""Compress → RMSNorm → GPT-J RoPE → quantize.
Gathers (1+overlap)*compress_ratio state entries per output token, applies
per-element softmax over the scores, and computes the weighted kv sum.
Returns (quantized_values, scale) matching the kernel's output layout.
"""
device = state_cache.device
head_dim = rms_weight.shape[0]
rope_dim = cos_sin_cache.shape[-1]
state_block_size = state_cache.shape[1]
state_width = state_cache.shape[-1] // 2
nope_dim = head_dim - rope_dim
total = (1 + overlap) * compress_ratio
results = []
for pos in positions.tolist():
src = torch.arange(pos - total + 1, pos + 1, dtype=torch.int64, device=device)
valid = src >= 0
idx = src.clamp(min=0)
pages = block_table[0, idx // state_block_size]
offsets = idx % state_block_size
raw = state_cache[pages, offsets].float() # [total, state_dim]
# Group 0 (tokens 0..cr-1): kv[:H], score[SW:SW+H]
# Group 1 (tokens cr..2cr-1): kv[H:2H], score[SW+H:SW+2H]
if overlap:
sw = state_width
g0_kv = raw[:compress_ratio, :head_dim]
g1_kv = raw[compress_ratio:, head_dim : 2 * head_dim]
g0_scores = raw[:compress_ratio, sw : sw + head_dim]
g1_scores = raw[compress_ratio:, sw + head_dim : sw + 2 * head_dim]
kv = torch.cat([g0_kv, g1_kv])
scores = torch.cat([g0_scores, g1_scores])
else:
kv = raw[:, :head_dim]
scores = raw[:, state_width : state_width + head_dim]
scores[~valid] = float("-inf")
kv[~valid] = 0.0
weights = torch.softmax(scores, dim=0)
compressed = (kv * weights).sum(dim=0) # [H]
var = (compressed * compressed).mean()
normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
compressed_pos = (pos // compress_ratio) * compress_ratio
cos, sin = cos_sin_cache[compressed_pos].float().chunk(2)
nope, rope = normed.split([nope_dim, rope_dim])
rope = torch.stack(
[rope[0::2] * cos - rope[1::2] * sin, rope[1::2] * cos + rope[0::2] * sin],
dim=-1,
).reshape(rope_dim)
results.append(torch.cat([nope, rope]).to(state_cache.dtype))
result = torch.stack(results)
if use_fp4:
return quantize_to_mxfp4(result)
else:
pairs = [
_ue8m0_reference(result[t], head_dim, fp8_max) for t in range(len(result))
]
quants, scales = zip(*pairs)
return torch.stack(quants), torch.cat(scales)
@pytest.mark.parametrize("num_tokens", [1, 7, 32])
@pytest.mark.parametrize("kv_block_size", [16, 32])
@pytest.mark.parametrize("use_fp4", [False, True])
def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
"""Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
HEAD_DIM = 128
ROPE_DIM = 64
BLOCK_SIZE = 16
RMS_EPS = 1e-6
FP8_MAX = 448.0
device = "cuda"
torch.manual_seed(42)
compress_ratio = 4
if use_fp4:
TOKEN_STRIDE = HEAD_DIM // 2 # packed nibbles: 64 bytes
SCALE_DIM = HEAD_DIM // 32 # ue8m0 bytes: 4
QUANT_BLOCK = 32
kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
else:
TOKEN_STRIDE = HEAD_DIM # FP8 bytes: 128
SCALE_DIM = 4 # 1 float32: 4 bytes
QUANT_BLOCK = HEAD_DIM
kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
# overlap=1 whenever compress_ratio==4, matching DeepseekCompressor logic.
overlap = 1 if compress_ratio == 4 else 0
coff = 1 + overlap # multiplier for state_dim per entry
num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
state_cache = torch.randn(
num_pages,
BLOCK_SIZE,
2 * coff * HEAD_DIM, # kv_state + score_state, each coff*HEAD_DIM wide
dtype=torch.bfloat16,
device=device,
)
block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
positions = torch.arange(
compress_ratio - 1,
compress_ratio * num_tokens,
compress_ratio,
dtype=torch.int64,
device=device,
)
rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
cos_sin_cache = torch.randn(compress_ratio * num_tokens, ROPE_DIM, device=device)
kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
kv_cache = torch.zeros(
kv_n_blocks,
kv_block_size * (TOKEN_STRIDE + SCALE_DIM),
dtype=torch.uint8,
device=device,
)
kernel[(num_tokens,)](
state_cache,
state_cache.stride(0),
state_cache.stride(1),
token_to_req,
positions,
slot_mapping,
block_table,
block_table.stride(0),
BLOCK_SIZE,
rms_weight,
RMS_EPS,
cos_sin_cache,
cos_sin_cache.stride(0),
kv_cache,
slot_mapping,
kv_block_size,
HEAD_SIZE=HEAD_DIM,
TRITON_BLOCK_SIZE=HEAD_DIM,
STATE_WIDTH=coff * HEAD_DIM,
COMPRESS_RATIO=compress_ratio,
OVERLAP=overlap,
ROPE_HEAD_DIM=ROPE_DIM,
FP8_MAX=FP8_MAX,
QUANT_BLOCK=QUANT_BLOCK,
TOKEN_STRIDE=TOKEN_STRIDE,
SCALE_DIM=SCALE_DIM,
KV_BLOCK_STRIDE=kv_cache.stride(0),
num_warps=1,
)
k_quant, scale = _reference_kv_compress_norm_rope(
state_cache,
block_table,
positions,
rms_weight,
cos_sin_cache,
compress_ratio,
overlap,
use_fp4,
rms_eps=RMS_EPS,
fp8_max=FP8_MAX,
)
if use_fp4:
for i in range(num_tokens):
blk, pos = i // kv_block_size, i % kv_block_size
val_off = pos * TOKEN_STRIDE
fp4_actual = kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
assert torch.equal(k_quant[i], fp4_actual), (
f"token {i}: packed nibbles differ, "
f"{(k_quant[i] != fp4_actual).sum()} "
f"/ {TOKEN_STRIDE}"
)
scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
scale_actual = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
assert torch.equal(scale_actual, scale[i]), (
f"token {i}: ue8m0 {scale_actual.tolist()} != {scale[i].tolist()}"
)
else:
k_quant = k_quant.view(torch.uint8)
for i in range(num_tokens):
blk, pos = i // kv_block_size, i % kv_block_size
val_off = pos * TOKEN_STRIDE
assert torch.equal(
k_quant[i], kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
), f"token {i}: FP8 bytes differ"
scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
torch.float32
)
assert torch.equal(actual_scale, scale[i : i + 1]), (
f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
)
@@ -30,56 +30,6 @@ N_HEAD = 64
MAX_POS = 4096
def quantize_to_mxfp4(
x: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Reference MXFP4 quantization.
Args:
x: [..., head_dim] where head_dim is divisible by 32
Returns:
packed: [..., head_dim//2] uint8 2 E2M1 nibbles/byte, low nibble = even index
scales: [..., head_dim//32] uint8 1 ue8m0 byte
"""
MXFP4_BLOCK_SIZE = 32
orig_shape = x.shape
head_dim = orig_shape[-1]
n_blocks = head_dim // MXFP4_BLOCK_SIZE
x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE)
# Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127
# 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126))
log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0)
scale = log2_ratio.exp2()
ue8m0 = (log2_ratio + 127.0).to(torch.uint8) # [*, n_blocks]
# E2M1 round-to-nearest-even: midpoints round to the even code.
# E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00]
# boundaries: [ 0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00]
x_scaled = (x_f32 / scale).clamp(-6.0, 6.0)
abs_x = x_scaled.abs()
code = torch.zeros_like(abs_x, dtype=torch.int32)
code = torch.where(abs_x > 0.25, 1, code)
code = torch.where(abs_x >= 0.75, 2, code)
code = torch.where(abs_x > 1.25, 3, code)
code = torch.where(abs_x >= 1.75, 4, code)
code = torch.where(abs_x > 2.5, 5, code)
code = torch.where(abs_x >= 3.5, 6, code)
code = torch.where(abs_x > 5.0, 7, code)
sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8)
nibble = code.to(torch.uint8) | (sign << 3)
# Pack: even-index element → low nibble, odd-index → high nibble
nibble_flat = nibble.reshape(-1, head_dim)
packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous()
packed = packed.reshape(*orig_shape[:-1], head_dim // 2)
scales = ue8m0.view(*orig_shape[:-1], n_blocks)
return packed, scales
def _reference(
positions: torch.Tensor,
q: torch.Tensor,
@@ -87,7 +37,6 @@ def _reference(
weights: torch.Tensor,
softmax_scale: float,
head_scale: float,
use_fp4: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
q_rot = q.clone()
ops.rotary_embedding(
@@ -100,33 +49,22 @@ def _reference(
HEAD_DIM - ROPE_DIM, # rope_dim_offset → rotate the tail
False,
)
q_fp8, q_scale = per_token_group_quant_fp8(
q_rot.view(-1, HEAD_DIM).contiguous(),
HEAD_DIM,
use_ue8m0=True,
)
q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
q_scale = q_scale.view(-1, N_HEAD)
if use_fp4:
q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
# Pack 4 ue8m0 bytes into 1 int32
q_scale = ue8m0.view(torch.int32).squeeze(-1)
# FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
weights_out = weights.to(torch.float32) * softmax_scale * head_scale
return (q_packed, q_scale), weights_out
else:
q_fp8, q_scale = per_token_group_quant_fp8(
q_rot.view(-1, HEAD_DIM).contiguous(),
HEAD_DIM,
use_ue8m0=True,
)
q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
q_scale = q_scale.view(-1, N_HEAD)
weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
return q_fp8, weights_out
weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
return q_fp8, weights_out
@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("use_fp4", [False, True])
@torch.inference_mode()
def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
device = "cuda"
torch.manual_seed(0)
@@ -139,32 +77,21 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use
softmax_scale = HEAD_DIM**-0.5
head_scale = N_HEAD**-0.5
q_quant_ref, weights_ref = _reference(
positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
q_fp8_ref, weights_ref = _reference(
positions, q, cos_sin_cache, weights, softmax_scale, head_scale
)
q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
)
if use_fp4:
q_quant_ref, q_scale_ref = q_quant_ref
q_quant_fused, q_scale_fused = q_quant_fused
assert torch.equal(q_scale_ref, q_scale_fused), (
f"q_scale mismatch: "
f"{(q_scale_ref != q_scale_fused).sum().item()} "
f"/ {q_scale_ref.numel()} bytes differ"
)
# fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
ref_bits = q_quant_ref.view(torch.int8)
fused_bits = q_quant_fused.view(torch.int8)
ref_bits = q_fp8_ref.view(torch.int8)
fused_bits = q_fp8_fused.view(torch.int8)
assert torch.equal(ref_bits, fused_bits), (
f"q_quant_fused mismatch: "
f"q_fp8 mismatch: "
f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
)
assert weights_fused.dtype == torch.float32
assert torch.equal(weights_ref, weights_fused), (
f"weights mismatch: max abs diff "
f"{(weights_ref - weights_fused).abs().max().item()}"
@@ -6,6 +6,7 @@ from transformers import AutoModel
from tests.models.utils import check_embeddings_close
from vllm import TokensPrompt
from vllm.config import PoolerConfig
@pytest.mark.parametrize(
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
with vllm_runner(
model,
runner="pooling",
pooler_config=PoolerConfig(task="token_embed"),
max_model_len=128,
max_num_batched_tokens=chunk_size,
enforce_eager=True,
+46 -14
View File
@@ -3,7 +3,6 @@
import httpx
import openai
import pytest
import pytest_asyncio
import torch
from ....utils import RemoteOpenAIServer
@@ -25,29 +24,42 @@ sentences_2 = [
similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
lexical_score_reference = [0.19554901123046875, 0.0]
colbert_score_reference = [0.7797, 0.4620]
SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
def pooling_task(request):
yield request.param
@pytest.fixture(scope="module")
def server():
def server(pooling_task):
args = [
"--max-model-len",
str(MAX_MODEL_LEN),
"--hf-overrides",
'{"architectures": ["BgeM3EmbeddingModel"]}',
"--pooler-config.task",
pooling_task,
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server):
async with server.get_async_client() as async_client:
yield async_client
@pytest.mark.asyncio
async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
async def test_bge_m3_api_server_embedding(server, pooling_task):
client = server.get_async_client()
if pooling_task != "embed":
with pytest.raises(openai.InternalServerError):
await run_client_embeddings(
client,
MODEL_NAME,
sentences_1,
)
return
embeddings_list_1 = await run_client_embeddings(
client,
MODEL_NAME,
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(
@pytest.mark.asyncio
async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
client = server.get_async_client()
if pooling_task != "token_classify":
with pytest.raises(openai.BadRequestError):
await sparse_embeddings(client, sentences_1)
return
embeddings_1 = await sparse_embeddings(client, sentences_1)
embeddings_2 = await sparse_embeddings(client, sentences_2)
@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
@pytest.mark.asyncio
async def test_bge_m3_api_server_sparse_embedding_corner_case(
client: openai.AsyncOpenAI,
):
async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
if pooling_task != "token_classify":
return
client = server.get_async_client()
embeddings = await sparse_embeddings(client, ["Hi"])
assert len(embeddings) == 1
assert 2673 in embeddings[0]
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
@pytest.mark.asyncio
async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
async def test_bge_m3_api_server_multi_vector(server, pooling_task):
client = server.get_async_client()
if pooling_task != "token_embed":
with pytest.raises(openai.BadRequestError):
await client.post(
"../pooling",
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
cast_to=httpx.Response,
)
return
result_1 = await client.post(
"../pooling",
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
@@ -4,6 +4,7 @@ import pytest
import torch
from vllm import TokensPrompt
from vllm.config import PoolerConfig
@pytest.mark.parametrize(
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
max_model_len=128,
enforce_eager=True,
runner="pooling",
pooler_config=PoolerConfig(task="token_embed"),
enable_prefix_caching=True,
) as vllm_model:
pooling_outputs = vllm_model.llm.encode(
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
assert len(output.prompt_token_ids) == n
assert len(output.outputs.data) == n
assert output.num_cached_tokens == 0
# skip_reading_prefix_cache can still write to cache
# to accelerate following requests
pooling_outputs = vllm_model.llm.encode(
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
pooling_task="embed",
)
for n, output in zip(n_prompt_tokens, pooling_outputs):
assert len(output.prompt_token_ids) == n
assert output.num_cached_tokens > 0
@@ -5,6 +5,7 @@ import torch
from transformers import AutoModel
from tests.models.utils import check_embeddings_close
from vllm.config import PoolerConfig
@pytest.mark.parametrize(
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
with vllm_runner(
model,
runner="pooling",
pooler_config=PoolerConfig(task="token_embed"),
max_model_len=None,
) as vllm_model:
vllm_outputs = vllm_model.token_embed(example_prompts)
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
model,
max_model_len=512,
dtype=dtype,
pooler_config=PoolerConfig(use_activation=False),
pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
) as vllm_model:
wo_normalize = vllm_model.token_embed(example_prompts)
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
model,
max_model_len=512,
dtype=dtype,
pooler_config=PoolerConfig(use_activation=True),
pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
) as vllm_model:
w_normalize = vllm_model.token_embed(example_prompts)
+30
View File
@@ -594,6 +594,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
"MiMoV2FlashForCausalLM": _HfExamplesInfo(
"XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True
),
"MiMoV2ProForCausalLM": _HfExamplesInfo(
"XiaomiMiMo/MiMo-V2.5-Pro", trust_remote_code=True, is_available_online=False
),
"Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),
}
@@ -961,6 +964,18 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"PerceptronAI/Isaac-0.1",
trust_remote_code=True,
extras={"0.2-2B-Preview": "PerceptronAI/Isaac-0.2-2B-Preview"},
max_transformers_version="4.57",
transformers_version_reason={
"vllm": (
"Custom Isaac code is not compatible with Transformers v5. "
"The model should be upstreamed to Transformers for "
"long-term support."
),
"hf": (
"Isaac's remote model and processor code import or configure "
"APIs that changed in Transformers v5."
),
},
),
"InternS1ForConditionalGeneration": _HfExamplesInfo(
"internlm/Intern-S1",
@@ -1057,6 +1072,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"MiDashengLMModel": _HfExamplesInfo(
"mispeech/midashenglm-7b", trust_remote_code=True
),
"MiMoV2OmniForCausalLM": _HfExamplesInfo(
"XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False
),
"MiniCPMO": _HfExamplesInfo(
"openbmb/MiniCPM-o-2_6",
trust_remote_code=True,
@@ -1540,6 +1558,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
trust_remote_code=True,
speculative_model="XiaomiMiMo/MiMo-7B-RL",
),
"MiMoV2MTPModel": _HfExamplesInfo(
"XiaomiMiMo/MiMo-V2.5-Pro",
trust_remote_code=True,
speculative_model="XiaomiMiMo/MiMo-V2.5-Pro",
is_available_online=False,
),
"MiMoV2OmniMTPModel": _HfExamplesInfo(
"XiaomiMiMo/MiMo-V2.5-Omni",
trust_remote_code=True,
speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",
is_available_online=False,
),
"NemotronHMTPModel": _HfExamplesInfo(
"nvidia/Nemotron-Super-Placeholder",
speculative_model="nvidia/Nemotron-Super-Placeholder",
+16
View File
@@ -5,6 +5,8 @@ Unit tests for MultiModalRegistry.supports_multimodal_inputs and
Qwen2.5-VL visual component loading behavior.
"""
from types import SimpleNamespace
import pytest
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -32,3 +34,17 @@ def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected):
limit_mm_per_prompt=limit_mm_per_prompt,
)
assert MULTIMODAL_REGISTRY.supports_multimodal_inputs(ctx.model_config) is expected
def test_create_processor_error_uses_served_model_name():
model_config = SimpleNamespace(
is_multimodal_model=False,
model="/path/to/model/weights",
served_model_name="friendly-model-name",
)
with pytest.raises(
ValueError,
match="friendly-model-name is not a multimodal model",
):
MULTIMODAL_REGISTRY.create_processor(model_config)
+185
View File
@@ -0,0 +1,185 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for Cutlass W4A16 (Machete) kernel on Hopper.
Verifies that W4A16 quantized models loaded through vllm select the
MacheteLinearKernel on sm_90 GPUs, that weights are correctly repacked,
and that inference produces valid output.
Run `pytest tests/quantization/test_cutlass_w4a16.py`.
"""
import pytest
import torch
from vllm.platforms import current_platform
if not current_platform.has_device_capability(90):
pytest.skip(
"Machete W4A16 requires Hopper (sm_90).",
allow_module_level=True,
)
from vllm.model_executor.kernels.linear import (
MPLinearLayerConfig,
choose_mp_linear_kernel,
)
from vllm.model_executor.kernels.linear.mixed_precision import (
MacheteLinearKernel,
)
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
CompressedTensorsLinearMethod,
CompressedTensorsWNA16,
)
from vllm.scalar_type import scalar_types
@pytest.fixture(scope="function", autouse=True)
def enable_pickle(monkeypatch):
"""`LLM.apply_model` requires pickling a function."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
@pytest.mark.parametrize(
"act_type,weight_type,group_size,zero_points",
[
(torch.float16, scalar_types.uint4b8, 128, False),
(torch.bfloat16, scalar_types.uint4b8, 128, False),
(torch.float16, scalar_types.uint4, 128, True),
(torch.float16, scalar_types.uint4b8, -1, False),
],
ids=[
"fp16-gptq-g128",
"bf16-gptq-g128",
"fp16-awq-g128",
"fp16-channelwise",
],
)
def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points):
"""Verify choose_mp_linear_kernel picks MacheteLinearKernel."""
config = MPLinearLayerConfig(
full_weight_shape=(4096, 4096),
partition_weight_shape=(4096, 4096),
act_type=act_type,
weight_type=weight_type,
group_size=group_size,
zero_points=zero_points,
has_g_idx=False,
)
kernel = choose_mp_linear_kernel(config)
assert kernel is MacheteLinearKernel, (
f"Expected MacheteLinearKernel, got {kernel.__name__}"
)
@pytest.mark.parametrize(
"full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason",
[
((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"),
(
(4096, 4096),
(4096, 4096),
scalar_types.float6_e3m2f,
128,
False,
"Quant type",
),
((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"),
],
ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"],
)
def test_machete_rejects_invalid_config(
full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason
):
"""Verify Machete rejects unsupported configurations."""
config = MPLinearLayerConfig(
full_weight_shape=full_shape,
partition_weight_shape=part_shape,
act_type=torch.float16,
weight_type=weight_type,
group_size=group_size,
zero_points=False,
has_g_idx=has_g_idx,
)
can_impl, reason = MacheteLinearKernel.can_implement(config)
assert not can_impl
assert expected_reason in reason
def test_kernel_selection_with_disabled_machete(monkeypatch):
"""Verify kernel selection falls back when Machete is disabled."""
monkeypatch.setattr("vllm.envs.VLLM_DISABLED_KERNELS", ["MacheteLinearKernel"])
config = MPLinearLayerConfig(
full_weight_shape=(4096, 4096),
partition_weight_shape=(4096, 4096),
act_type=torch.float16,
weight_type=scalar_types.uint4b8,
group_size=128,
zero_points=False,
has_g_idx=False,
)
kernel = choose_mp_linear_kernel(config)
assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
@pytest.mark.parametrize(
"model_name",
[
"nm-testing/tinyllama-oneshot-w4a16-channel-v2",
"nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
],
)
def test_w4a16_machete_e2e(vllm_runner, model_name):
"""Load a W4A16 model, verify Machete kernel is used, and generate."""
with vllm_runner(model_name, enforce_eager=True, gpu_memory_utilization=0.5) as llm:
def check_model(model):
layer = model.model.layers[0]
qkv_proj = layer.self_attn.qkv_proj
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
assert isinstance(qkv_proj.scheme.kernel, MacheteLinearKernel), (
f"Expected MacheteLinearKernel on Hopper, "
f"got {type(qkv_proj.scheme.kernel).__name__}"
)
assert hasattr(qkv_proj, "weight_packed")
assert hasattr(qkv_proj, "weight_scale")
assert qkv_proj.weight_packed.dtype == torch.int32
llm.apply_model(check_model)
output = llm.generate_greedy("Hello my name is", max_tokens=10)
assert output
assert len(output[0][1]) > 0
def test_w4a16_machete_bfloat16_deterministic(vllm_runner):
"""Verify Machete works with bf16 activations and is deterministic."""
model_name = "nm-testing/tinyllama-oneshot-w4a16-channel-v2"
prompt = "The capital of France is"
with vllm_runner(
model_name,
enforce_eager=True,
dtype="bfloat16",
gpu_memory_utilization=0.5,
) as llm:
def check_kernel_type(model):
layer = model.model.layers[0]
scheme = layer.self_attn.qkv_proj.scheme
assert isinstance(scheme.kernel, MacheteLinearKernel), (
f"Expected MacheteLinearKernel with bf16, "
f"got {type(scheme.kernel).__name__}"
)
llm.apply_model(check_kernel_type)
out1 = llm.generate_greedy(prompt, max_tokens=10)
out2 = llm.generate_greedy(prompt, max_tokens=10)
assert out1[0][1] == out2[0][1], (
f"Non-deterministic: '{out1[0][1]}' vs '{out2[0][1]}'"
)
@@ -188,30 +188,6 @@ class TestExtractToolCalls:
"location": "NYC"
}
def test_type_conversion_in_non_streaming(self):
"""Non-streaming extraction must convert params using the tool schema."""
tool = ChatCompletionToolsParam(
function=FunctionDefinition(
name="toggle",
parameters={
"type": "object",
"properties": {
"enabled": {"type": "boolean"},
"count": {"type": "integer"},
},
},
),
)
parser = make_parser(tools=[tool])
model_output = build_tool_call("toggle", {"enabled": "true", "count": "42"})
result = parser.extract_tool_calls(model_output, None)
assert result.tools_called
assert len(result.tool_calls) == 1
args = json.loads(result.tool_calls[0].function.arguments)
assert args == {"enabled": True, "count": 42}
assert isinstance(args["enabled"], bool)
assert isinstance(args["count"], int)
# ---------------------------------------------------------------------------
# Tests: extract_tool_calls_streaming
-48
View File
@@ -2074,54 +2074,6 @@ def test_auto_fit_max_model_len_not_triggered():
assert vllm_config.model_config.max_model_len == 16
def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
"""Auto-fit must size max_model_len against the override-clamped pool, not
the raw `available_memory`. Without this, auto-fit could pick a
max_model_len that no longer fits once `num_gpu_blocks_override` is applied.
"""
model_config = ModelConfig(max_model_len=16384)
model_config.original_max_model_len = -1 # request auto-fit
vllm_config = VllmConfig(model_config=model_config)
# Cap the cache to 32 blocks regardless of available memory.
vllm_config.cache_config.num_gpu_blocks_override = 32
mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
kv_cache_specs = {
"layer_1": new_kv_cache_spec(), # block_size=16
"layer_2": new_kv_cache_spec(),
}
# Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384).
large_available_memory = mem_per_block_per_layer * 2 * 1024
get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
# 32 blocks * block_size 16 = 512 token slots, so max_model_len must
# auto-fit at or below that.
assert 0 < vllm_config.model_config.max_model_len <= 32 * 16
def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
"""Admission check must use the override-clamped pool size, not raw
`available_memory`. Without this, startup could accept a max_model_len
that does not actually fit in `num_gpu_blocks_override` blocks.
"""
model_config = ModelConfig(max_model_len=16384)
vllm_config = VllmConfig(model_config=model_config)
# 32 blocks is far too small for max_model_len=16384 (would need 1024).
vllm_config.cache_config.num_gpu_blocks_override = 32
mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
kv_cache_specs = {
"layer_1": new_kv_cache_spec(),
"layer_2": new_kv_cache_spec(),
}
# Plenty of raw memory: a bytes-only check against this would pass.
large_available_memory = mem_per_block_per_layer * 2 * 1024
with pytest.raises(ValueError, match="max seq len"):
get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
def test_unify_hybrid_kv_cache_specs():
# 1. has_full_attention and has_sliding_window
before_spec_1 = new_kv_cache_spec()
@@ -324,13 +324,10 @@ def run_test(
):
spec_decoding = spec_config is not None
cache_arg: dict[str, Any] = (
# Force preemptions: with 32 blocks the cache holds at most a single
# max-length request, so the ~34 concurrent prompts contend and trigger
# preemption. (Prompts here are << max_model_len, so dropping
# max_model_len from 4096 to 512 doesn't change generation behavior.)
dict(num_gpu_blocks_override=32, max_model_len=512)
# Force preemptions
dict(num_gpu_blocks_override=32)
if test_preemption
else dict(gpu_memory_utilization=0.9, max_model_len=4096)
else dict(gpu_memory_utilization=0.9)
)
spec_mml = (spec_config or {}).get("max_model_len")
spec_method = (spec_config or {}).get("method", "none")
@@ -346,6 +343,7 @@ def run_test(
with VllmRunner(
model,
max_model_len=4096,
enable_chunked_prefill=test_prefill_chunking,
# Force prefill chunking
max_num_batched_tokens=48 if test_prefill_chunking else None,
@@ -8,16 +8,11 @@ from unittest.mock import Mock
import pytest
from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
from vllm.reasoning import ReasoningParser
from vllm.v1.request import Request
from vllm.v1.structured_output import StructuredOutputManager
class MockReasoner:
def __init__(self, tokenizer):
self.is_reasoning_end = Mock(return_value=False)
self.is_reasoning_end_streaming = Mock(return_value=False)
class TestReasoningStructuredOutput:
"""Test reasoning-aware structured output functionality."""
@@ -55,6 +50,13 @@ class TestReasoningStructuredOutput:
config.speculative_config = None
return config
@pytest.fixture
def mock_reasoning_parser(self):
"""Create a mock ReasoningParser."""
parser = Mock(spec=ReasoningParser)
parser.is_reasoning_end = Mock(return_value=False)
return parser
@pytest.fixture
def mock_request_with_structured_output(self):
"""Create a mock request with structured output."""
@@ -62,8 +64,6 @@ class TestReasoningStructuredOutput:
request.structured_output_request = Mock()
request.structured_output_request.reasoning_ended = None
request.structured_output_request.grammar = Mock()
request.structured_output_request.reasoning_parser_kwargs = None
request.structured_output_request.reasoner = None
request.structured_output_request.grammar.is_terminated = Mock(
return_value=False
)
@@ -74,13 +74,6 @@ class TestReasoningStructuredOutput:
request.num_output_placeholders = 0
return request
@pytest.fixture
def manager_with_reasoner(self, mock_vllm_config):
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner_cls = MockReasoner
manager.tokenizer = Mock()
return manager
def test_should_fill_bitmask_with_enable_in_reasoning(
self, mock_vllm_config, mock_request_with_structured_output
):
@@ -96,17 +89,22 @@ class TestReasoningStructuredOutput:
def test_should_fill_bitmask_without_enable_in_reasoning(
self,
manager_with_reasoner,
mock_vllm_config,
mock_request_with_structured_output,
mock_reasoning_parser,
):
"""Test should_fill_bitmask when enable_in_reasoning is False."""
# Keep enable_in_reasoning as False (default)
config = manager_with_reasoner.vllm_config.structured_outputs_config
config = mock_vllm_config.structured_outputs_config
assert config.enable_in_reasoning is False
result = manager_with_reasoner.should_fill_bitmask(
mock_request_with_structured_output
)
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner = mock_reasoning_parser
# Mock reasoning not ended
mock_reasoning_parser.is_reasoning_end.return_value = False
result = manager.should_fill_bitmask(mock_request_with_structured_output)
# Should set reasoning_ended and return its value
assert (
@@ -120,92 +118,68 @@ class TestReasoningStructuredOutput:
):
"""Test should_fill_bitmask when no reasoner is configured."""
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner = None
result = manager.should_fill_bitmask(mock_request_with_structured_output)
# Should default to True when no reasoner
assert result is True
def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
self, mock_vllm_config, mock_request_with_structured_output
):
"""Test request-level parser kwargs override the default reasoner."""
class KwargReasoner:
def __init__(self, tokenizer, chat_template_kwargs=None):
self.chat_template_kwargs = chat_template_kwargs or {}
def is_reasoning_end(self, input_ids):
return not self.chat_template_kwargs.get("enable_thinking", False)
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner_cls = KwargReasoner
manager.tokenizer = Mock()
structured_req = mock_request_with_structured_output.structured_output_request
structured_req.reasoning_parser_kwargs = {
"chat_template_kwargs": {"enable_thinking": True}
}
result = manager.should_fill_bitmask(mock_request_with_structured_output)
assert result is False
assert (
mock_request_with_structured_output.structured_output_request.reasoner
is not None
)
def test_should_advance_with_enable_in_reasoning(
self,
manager_with_reasoner,
mock_vllm_config,
mock_request_with_structured_output,
mock_reasoning_parser,
):
"""Test should_advance when enable_in_reasoning is True."""
# Enable enable_in_reasoning
manager_with_reasoner.enable_in_reasoning = True
mock_vllm_config.structured_outputs_config.enable_in_reasoning = True
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner = mock_reasoning_parser
# Should always return True when enable_in_reasoning is enabled
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output
)
result = manager.should_advance(mock_request_with_structured_output)
assert result is True
def test_should_advance_reasoning_not_ended(
self,
manager_with_reasoner,
mock_vllm_config,
mock_request_with_structured_output,
mock_reasoning_parser,
):
"""Test should_advance when reasoning has not ended."""
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner = mock_reasoning_parser
# Set reasoning as not ended
(
mock_request_with_structured_output.structured_output_request
).reasoning_ended = False
mock_reasoning_parser.is_reasoning_end.return_value = False
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output
)
result = manager.should_advance(mock_request_with_structured_output)
# Should return False since reasoning hasn't ended
assert result is False
def test_should_advance_reasoning_just_ended(
self,
manager_with_reasoner,
mock_vllm_config,
mock_request_with_structured_output,
mock_reasoning_parser,
):
"""Test should_advance when reasoning ends in current step."""
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner = mock_reasoning_parser
# Set reasoning as not ended initially, but ends in this step
(
mock_request_with_structured_output.structured_output_request
).reasoning_ended = False
reasoner = MockReasoner(tokenizer=Mock())
reasoner.is_reasoning_end_streaming.return_value = True
structured_req = mock_request_with_structured_output.structured_output_request
structured_req.reasoner = reasoner
mock_reasoning_parser.is_reasoning_end.return_value = True
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output
)
result = manager.should_advance(mock_request_with_structured_output)
# Should set reasoning_ended to True but return False for this step
assert (
@@ -216,18 +190,20 @@ class TestReasoningStructuredOutput:
def test_should_advance_reasoning_already_ended(
self,
manager_with_reasoner,
mock_vllm_config,
mock_request_with_structured_output,
mock_reasoning_parser,
):
"""Test should_advance when reasoning has already ended."""
manager = StructuredOutputManager(mock_vllm_config)
manager.reasoner = mock_reasoning_parser
# Set reasoning as already ended
(
mock_request_with_structured_output.structured_output_request
).reasoning_ended = True
result = manager_with_reasoner.should_advance(
mock_request_with_structured_output
)
result = manager.should_advance(mock_request_with_structured_output)
# Should return True since reasoning has ended
assert result is True
+13 -1
View File
@@ -326,6 +326,10 @@ class ModelConfig:
mm_encoder_only: InitVar[bool | None] = None
mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None
mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = None
mm_encoder_attn_dtype: InitVar[str | None] = None
mm_encoder_fp8_scale_path: InitVar[str | None] = None
mm_encoder_fp8_scale_save_path: InitVar[str | None] = None
mm_encoder_fp8_scale_save_margin: InitVar[float | None] = None
interleave_mm_strings: InitVar[bool | None] = None
skip_mm_profiling: InitVar[bool | None] = None
video_pruning_rate: InitVar[float | None] = None
@@ -447,6 +451,10 @@ class ModelConfig:
mm_encoder_only: bool | None,
mm_encoder_tp_mode: MMEncoderTPMode | None,
mm_encoder_attn_backend: AttentionBackendEnum | str | None,
mm_encoder_attn_dtype: str | None,
mm_encoder_fp8_scale_path: str | None,
mm_encoder_fp8_scale_save_path: str | None,
mm_encoder_fp8_scale_save_margin: float | None,
interleave_mm_strings: bool | None,
skip_mm_profiling: bool | None,
video_pruning_rate: float | None,
@@ -513,6 +521,7 @@ class ModelConfig:
if dict_overrides:
self._apply_dict_overrides(hf_config, dict_overrides)
self.hf_text_config = get_hf_text_config(self.hf_config)
self.model_arch_config = self.get_model_arch_config()
self.attention_chunk_size = getattr(
self.hf_text_config, "attention_chunk_size", None
)
@@ -520,7 +529,6 @@ class ModelConfig:
self.hf_image_processor_config = get_hf_image_processor_config(
self.model, hf_token=self.hf_token, revision=self.revision
)
self.model_arch_config = self.get_model_arch_config()
architectures = self.architectures
registry = self.registry
@@ -643,6 +651,10 @@ class ModelConfig:
mm_encoder_only=mm_encoder_only,
mm_encoder_tp_mode=mm_encoder_tp_mode,
mm_encoder_attn_backend=mm_encoder_attn_backend,
mm_encoder_attn_dtype=mm_encoder_attn_dtype,
mm_encoder_fp8_scale_path=mm_encoder_fp8_scale_path,
mm_encoder_fp8_scale_save_path=mm_encoder_fp8_scale_save_path,
mm_encoder_fp8_scale_save_margin=mm_encoder_fp8_scale_save_margin,
interleave_mm_strings=interleave_mm_strings,
skip_mm_profiling=skip_mm_profiling,
video_pruning_rate=video_pruning_rate,
+51
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Literal, TypeAlias, TypedDict, final
from pydantic import ConfigDict, Field, field_validator, model_validator
@@ -158,6 +159,24 @@ class MultiModalConfig:
"""Optional override for the multi-modal encoder attention backend when
using vision transformers. Accepts any value from
`vllm.v1.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
mm_encoder_attn_dtype: Literal["fp8"] | None = None
"""Optional dtype override for ViT encoder attention. Set to `"fp8"` to
enable FP8 quantization via the FlashInfer cuDNN backend. When set to
`"fp8"` without a scale file, dynamic scaling is used automatically.
See docs/features/quantization/fp8_vit_attn.md for details."""
mm_encoder_fp8_scale_path: str | None = None
"""Path to a JSON file containing per-layer FP8 Q/K/V scales for ViT
encoder attention. When provided (with `mm_encoder_attn_dtype="fp8"`),
static scaling is used. When omitted, dynamic scaling is used."""
mm_encoder_fp8_scale_save_path: str | None = None
"""When set with dynamic FP8 scaling (`mm_encoder_attn_dtype="fp8"`
and no `mm_encoder_fp8_scale_path`), saves the calibrated scales to
this file after the amax history buffer is full. The saved file can
then be used as `mm_encoder_fp8_scale_path` in subsequent runs."""
mm_encoder_fp8_scale_save_margin: float = Field(default=1.5, gt=0.0)
"""Safety margin multiplied onto scales when auto-saving. A value > 1
leaves headroom so that inputs with larger activations than the
calibration set do not overflow FP8 range. Default 1.5."""
interleave_mm_strings: bool = False
"""Enable fully interleaved support for multimodal prompts, while using
--chat-template-content-format=string."""
@@ -233,6 +252,36 @@ class MultiModalConfig:
"'mm_shm_cache_max_object_size_mb' should only be set when "
"'mm_processor_cache_type' is 'shm'."
)
# Validate FP8 scale path combinations.
if self.mm_encoder_attn_dtype != "fp8" and (
self.mm_encoder_fp8_scale_path is not None
or self.mm_encoder_fp8_scale_save_path is not None
):
raise ValueError(
"'mm_encoder_fp8_scale_path' and "
"'mm_encoder_fp8_scale_save_path' require "
"'mm_encoder_attn_dtype' to be 'fp8'."
)
if (
self.mm_encoder_fp8_scale_path is not None
and self.mm_encoder_fp8_scale_save_path is not None
):
raise ValueError(
"'mm_encoder_fp8_scale_save_path' cannot be used with "
"'mm_encoder_fp8_scale_path' (saving requires dynamic scaling)."
)
# Validate file paths exist.
if self.mm_encoder_fp8_scale_path is not None:
scale_path = Path(self.mm_encoder_fp8_scale_path)
if not scale_path.is_file():
raise FileNotFoundError(f"FP8 scale file not found: {scale_path}")
if self.mm_encoder_fp8_scale_save_path is not None:
save_parent = Path(self.mm_encoder_fp8_scale_save_path).parent
if not save_parent.is_dir():
raise FileNotFoundError(
f"Parent directory for FP8 scale save path not found: {save_parent}"
)
return self
def compute_hash(self) -> str:
@@ -252,6 +301,8 @@ class MultiModalConfig:
if self.mm_encoder_attn_backend is not None
else None,
self.mm_encoder_tp_mode,
self.mm_encoder_attn_dtype,
self.mm_encoder_fp8_scale_path,
]
hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
return hash_str
+58 -6
View File
@@ -34,6 +34,7 @@ logger = init_logger(__name__)
MTPModelTypes = Literal[
"deepseek_mtp",
"mimo_mtp",
"mimo_v2_mtp",
"glm4_moe_mtp",
"glm4_moe_lite_mtp",
"glm_ocr_mtp",
@@ -63,7 +64,8 @@ SpeculativeMethod = Literal[
EagleModelTypes,
NgramGPUTypes,
]
RejectionSampleMethod = Literal["strict", "probabilistic", "synthetic"]
RejectionSampleMethod = Literal["standard", "synthetic"]
DraftSampleMethod = Literal["greedy", "gumbel"]
@config
@@ -183,11 +185,11 @@ class SpeculativeConfig:
"""Load config for the draft model. If not specified, will use the load
config from the target model."""
rejection_sample_method: RejectionSampleMethod = "strict"
"""Whether to use strict (target and draft sampled tokens match exactly)
or probabilistic rejection sampling. Both respect the target model
distribution, but the latter yields a higher acceptance rate at the cost
of more memory to cache draft logits."""
rejection_sample_method: RejectionSampleMethod = "standard"
"""The rejection sampling method to use. 'standard' uses probabilistic
rejection sampling (with or without cached draft logits, controlled by
draft_sample_method). 'synthetic' accepts draft tokens with a decaying
probability calibrated to synthetic_acceptance_rate."""
synthetic_acceptance_rates: list[float] | None = None
"""Per-position *unconditional* acceptance rates for synthetic rejection
@@ -248,6 +250,14 @@ class SpeculativeConfig:
)
return SpeculativeConfig._acceptance_length_to_rates(length, n)
draft_sample_method: DraftSampleMethod = "greedy"
"""How the draft model samples tokens. 'greedy' always picks the argmax
token, and the draft probabilities are treated as one-hot during rejection
sampling. 'gumbel' adds Gumbel noise for stochastic sampling, and the full
draft logits are used for the probability ratio test during rejection
sampling. This comes at the cost of additional GPU memory usage. This
parameter currently only applies to Model Runner V2."""
def compute_hash(self) -> str:
"""
WARNING: Whenever a new field is added to this config,
@@ -323,6 +333,48 @@ class SpeculativeConfig:
}
)
if (arch := hf_config.architectures[0]) in (
"MiMoV2ProForCausalLM",
"MiMoV2OmniForCausalLM",
):
from vllm.model_executor.models.mimo_v2_mtp import (
_MIMO_V2_PRO_NUM_MTP_LAYERS,
)
mtp_arch_maps = {
"MiMoV2ProForCausalLM": "MiMoV2MTPModel",
"MiMoV2OmniForCausalLM": "MiMoV2OmniMTPModel",
}
hf_config.model_type = "mimo_v2_mtp"
# vLLM currently supports only the first MiMo-V2 MTP layer.
n_predict = _MIMO_V2_PRO_NUM_MTP_LAYERS
hf_config.update(
{
"num_hidden_layers": 0,
"n_predict": n_predict,
"num_nextn_predict_layers": n_predict,
"architectures": [mtp_arch_maps[arch]],
}
)
if hf_config.architectures[0] == "MiMoV2FlashForCausalLM":
from vllm.model_executor.models.mimo_v2_mtp import (
_MIMO_V2_FLASH_NUM_MTP_LAYERS,
)
hf_config.model_type = "mimo_v2_mtp"
# vLLM currently supports only the first MiMo-V2 MTP layer.
n_predict = _MIMO_V2_FLASH_NUM_MTP_LAYERS
hf_config.update(
{
"num_hidden_layers": 0,
"n_predict": n_predict,
"num_nextn_predict_layers": n_predict,
"architectures": ["MiMoV2MTPModel"],
}
)
if hf_config.architectures[0] == "Glm4MoeForCausalLM":
hf_config.model_type = "glm4_moe_mtp"
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
-10
View File
@@ -1432,10 +1432,6 @@ class VllmConfig:
cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
range(256, max_graph_size + 1, 16))
`max_num_batched_tokens` is also appended to the list if it fits
within `max_cudagraph_capture_size`, so the max batch size is captured
even when off-stride.
In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
will be the final sizes to capture cudagraph (in ascending order).
@@ -1524,12 +1520,6 @@ class VllmConfig:
cudagraph_capture_sizes += list(
range(256, max_cudagraph_capture_size + 1, 16)
)
# ensure max_num_tokens is captured if within max capture size
if (
max_num_tokens <= max_cudagraph_capture_size
and max_num_tokens not in cudagraph_capture_sizes
):
cudagraph_capture_sizes.append(max_num_tokens)
# de-duplicate and sort the sizes
cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))
@@ -584,8 +584,6 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
top_k: int,
num_experts: int,
hidden_size: int,
dispatch_dtype_bytes_per_elem: int = 0,
dispatch_scale_bytes_per_token: int = 0,
):
"""Initialize the MoeAlltoAll workspace."""
if self.initialized:
@@ -616,13 +614,9 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
ep_config = MnnvlConfig(
comm_backend=CustomCommunicator(self.cpu_group),
)
if dispatch_dtype_bytes_per_elem == 0:
hidden_bytes = hidden_size // 2
else:
hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
total_dispatch_payload_size_per_token = (
hidden_bytes
+ dispatch_scale_bytes_per_token
hidden_size // 2 # nvfp4 hidden states
+ hidden_size // 16 # fp8 scaling factors
+ top_k * 4 # int32 topks ids
+ top_k * 4 # float32 topk weights
)
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Scheduler-side logic for the NIXL connector."""
import os
import threading
import time
from typing import TYPE_CHECKING, Any
@@ -88,6 +89,12 @@ class NixlConnectorScheduler:
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
logger.info("Hybrid Memory Allocator is enabled with NIXL")
if os.environ.get("VLLM_NIXL_ABORT_REQUEST_TIMEOUT") is not None:
logger.warning(
"VLLM_NIXL_ABORT_REQUEST_TIMEOUT is deprecated and will be "
"removed in release 0.22.0."
)
# Background thread for handling new handshake requests.
self._nixl_handshake_listener_t: threading.Thread | None = None
self._stop_event = threading.Event()
+28
View File
@@ -542,6 +542,14 @@ class EngineArgs:
mm_encoder_attn_backend: AttentionBackendEnum | str | None = (
MultiModalConfig.mm_encoder_attn_backend
)
mm_encoder_attn_dtype: str | None = MultiModalConfig.mm_encoder_attn_dtype
mm_encoder_fp8_scale_path: str | None = MultiModalConfig.mm_encoder_fp8_scale_path
mm_encoder_fp8_scale_save_path: str | None = (
MultiModalConfig.mm_encoder_fp8_scale_save_path
)
mm_encoder_fp8_scale_save_margin: float = (
MultiModalConfig.mm_encoder_fp8_scale_save_margin
)
io_processor_plugin: str | None = None
renderer_num_workers: int = 1
skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
@@ -1179,6 +1187,22 @@ class EngineArgs:
"--mm-encoder-attn-backend",
**multimodal_kwargs["mm_encoder_attn_backend"],
)
multimodal_group.add_argument(
"--mm-encoder-attn-dtype",
**multimodal_kwargs["mm_encoder_attn_dtype"],
)
multimodal_group.add_argument(
"--mm-encoder-fp8-scale-path",
**multimodal_kwargs["mm_encoder_fp8_scale_path"],
)
multimodal_group.add_argument(
"--mm-encoder-fp8-scale-save-path",
**multimodal_kwargs["mm_encoder_fp8_scale_save_path"],
)
multimodal_group.add_argument(
"--mm-encoder-fp8-scale-save-margin",
**multimodal_kwargs["mm_encoder_fp8_scale_save_margin"],
)
multimodal_group.add_argument(
"--interleave-mm-strings", **multimodal_kwargs["interleave_mm_strings"]
)
@@ -1517,6 +1541,10 @@ class EngineArgs:
mm_encoder_only=self.mm_encoder_only,
mm_encoder_tp_mode=self.mm_encoder_tp_mode,
mm_encoder_attn_backend=self.mm_encoder_attn_backend,
mm_encoder_attn_dtype=self.mm_encoder_attn_dtype,
mm_encoder_fp8_scale_path=self.mm_encoder_fp8_scale_path,
mm_encoder_fp8_scale_save_path=self.mm_encoder_fp8_scale_save_path,
mm_encoder_fp8_scale_save_margin=self.mm_encoder_fp8_scale_save_margin,
pooler_config=self.pooler_config,
generation_config=self.generation_config,
override_generation_config=self.override_generation_config,
-1
View File
@@ -78,7 +78,6 @@ class EngineClient(ABC):
priority: int = 0,
data_parallel_rank: int | None = None,
reasoning_ended: bool | None = None,
reasoning_parser_kwargs: dict[str, Any] | None = None,
) -> AsyncGenerator[RequestOutput, None]:
"""Generate outputs for a request."""
...
+5 -8
View File
@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
prompt_to_seq,
)
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
from vllm.tasks import PoolingTask
from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
from vllm.tokenizers import TokenizerLike
from vllm.usage.usage_lib import UsageContext
from vllm.utils.counter import Counter
@@ -1204,12 +1204,9 @@ class LLM:
f"Supported tasks: {self.supported_tasks}"
)
else:
logger.warning_once(
"Pooling multitask support is deprecated and will "
"be removed in v0.20. When the default pooling task is "
"not what you want, you need to manually specify it "
'via PoolerConfig(task="%s"). ',
pooling_task,
raise ValueError(
f"Try switching the model's pooling_task "
f'via `PoolerConfig(task="{pooling_task}")`'
)
if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
@@ -1412,7 +1409,7 @@ class LLM:
"pooling model."
)
score_type = self.model_config.score_type
score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
if (
score_type == "cross-encoder"
and getattr(self.model_config.hf_config, "num_labels", 0) != 1
@@ -317,4 +317,5 @@ class OpenAIServingChatBatch(OpenAIServingChat):
model=model_name,
choices=choices,
usage=usage,
system_fingerprint=self.system_fingerprint,
)
@@ -129,6 +129,9 @@ class ChatCompletionStreamResponse(OpenAIBaseModel):
model: str
choices: list[ChatCompletionResponseStreamChoice]
usage: UsageInfo | None = Field(default=None)
# Set only on the final chunk of a stream to mirror non-streaming responses
# without the per-chunk serialization overhead.
system_fingerprint: str | None = None
# not part of the OpenAI spec but for tracing the tokens
prompt_token_ids: list[int] | None = None
@@ -347,11 +347,6 @@ class OpenAIServingChat(OpenAIServing):
priority=request.priority,
data_parallel_rank=data_parallel_rank,
reasoning_ended=reasoning_ended,
reasoning_parser_kwargs={
"chat_template_kwargs": chat_template_kwargs,
}
if reasoning_parser
else None,
)
generators.append(generator)
@@ -1200,6 +1195,16 @@ class OpenAIServingChat(OpenAIServing):
choices=[choice_data],
model=model_name,
)
# Stamp the fingerprint on terminal chunks only (those with
# finish_reason set). When ``include_usage`` is on, the
# trailing usage chunk below overrides this as the true
# final message.
if (
not include_usage
and self.system_fingerprint is not None
and choice_data.finish_reason is not None
):
chunk.system_fingerprint = self.system_fingerprint
# handle usage stats if requested & if continuous
if include_continuous_usage:
@@ -1234,6 +1239,7 @@ class OpenAIServingChat(OpenAIServing):
choices=[],
model=model_name,
usage=final_usage,
system_fingerprint=self.system_fingerprint,
)
final_usage_data = final_usage_chunk.model_dump_json(
exclude_unset=True, exclude_none=True
@@ -1642,6 +1648,7 @@ class OpenAIServingChat(OpenAIServing):
model=model_name,
choices=choices,
usage=usage,
system_fingerprint=self.system_fingerprint,
prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
prompt_token_ids=(
final_res.prompt_token_ids if request.return_token_ids else None
+13 -1
View File
@@ -153,9 +153,21 @@ class BaseFrontendArgs:
"""If set to True, log the stack trace of error responses"""
tokens_only: bool = False
"""
If set to True, only enable the Tokens In<>Out endpoint.
If set to True, only enable the Tokens In<>Out endpoint.
This is intended for use in a Disaggregated Everything setup.
"""
fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full"
"""Controls the ``system_fingerprint`` field on responses.
- ``full`` (default): ``vllm-<version>[-<parallelism>]-<hash8>``. Encodes
server version, non-trivial parallelism degrees (tp/pp/dp/ep), and an
8-char config hash.
- ``hash``: ``vllm-<version>-<hash8>``. Parallelism stripped.
- ``custom``: emits the literal string from ``--fingerprint-value``.
- ``none``: the field is omitted (serialized as ``null``).
"""
fingerprint_value: str | None = None
"""Literal fingerprint string used when ``--fingerprint-mode=custom``."""
@classmethod
def _customize_cli_kwargs(
@@ -512,3 +512,6 @@ class CompletionStreamResponse(OpenAIBaseModel):
model: str
choices: list[CompletionResponseStreamChoice]
usage: UsageInfo | None = Field(default=None)
# Set only on the final chunk of a stream to mirror non-streaming responses
# without the per-chunk serialization overhead.
system_fingerprint: str | None = None
+12 -1
View File
@@ -383,6 +383,7 @@ class OpenAIServingCompletion(OpenAIServing):
chunk = CompletionStreamResponse(
id=request_id,
object="text_completion",
created=created_time,
model=model_name,
choices=[
@@ -401,6 +402,14 @@ class OpenAIServingCompletion(OpenAIServing):
)
],
)
# Stamp on terminal chunk only when no trailing usage chunk
# will follow (that one is the true final message).
if (
not include_usage
and self.system_fingerprint is not None
and finish_reason is not None
):
chunk.system_fingerprint = self.system_fingerprint
if include_continuous_usage:
prompt_tokens = num_prompt_tokens[prompt_idx]
completion_tokens = previous_num_tokens[i]
@@ -410,7 +419,7 @@ class OpenAIServingCompletion(OpenAIServing):
total_tokens=prompt_tokens + completion_tokens,
)
response_json = chunk.model_dump_json(exclude_unset=False)
response_json = chunk.model_dump_json(exclude_unset=True)
yield f"data: {response_json}\n\n"
total_prompt_tokens = sum(num_prompt_tokens)
@@ -433,6 +442,7 @@ class OpenAIServingCompletion(OpenAIServing):
model=model_name,
choices=[],
usage=final_usage_info,
system_fingerprint=self.system_fingerprint,
)
final_usage_data = final_usage_chunk.model_dump_json(
exclude_unset=False, exclude_none=True
@@ -562,6 +572,7 @@ class OpenAIServingCompletion(OpenAIServing):
model=model_name,
choices=choices,
usage=usage,
system_fingerprint=self.system_fingerprint,
kv_transfer_params=kv_transfer_params,
)
+13
View File
@@ -157,6 +157,19 @@ class OpenAIServing:
self.renderer = engine_client.renderer
self.input_processor = engine_client.input_processor
# Computed once at startup (cached by ``vllm_config`` identity) and
# stamped on non-streaming responses. Streaming chunks deliberately
# omit it to avoid per-chunk overhead.
from vllm.entrypoints.openai.fingerprint import get_system_fingerprint
try:
self.system_fingerprint: str | None = get_system_fingerprint(
engine_client.vllm_config
)
except Exception:
# Never fail server startup over the fingerprint.
self.system_fingerprint = None
async def beam_search(
self,
prompt: EngineInput,
+84
View File
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Build the ``system_fingerprint`` string returned by the OpenAI-compatible
server.
Four modes, configured via ``--fingerprint-mode``:
* ``full`` (default): ``vllm-<version>[-<parallelism>]-<hash8>`` encodes
server version, any non-trivial parallelism degree (tp/pp/dp/ep), and an
8-char prefix of ``vllm_config.compute_hash()`` (covers model identity,
quant config, speculative, attention backend, etc.).
* ``hash``: ``vllm-<version>-<hash8>`` parallelism stripped.
* ``custom``: user-provided literal via ``--fingerprint-value``.
* ``none``: the field is omitted (serialized as ``null``).
``get_system_fingerprint`` is only called at serving-class init (a handful
of times per server); each subclass caches the returned string on
``self.system_fingerprint``, so per-request cost is one attribute read.
"""
from __future__ import annotations
from typing import Any, Literal
FingerprintMode = Literal["full", "hash", "custom", "none"]
_DEFAULT_MODE: FingerprintMode = "full"
_CUSTOM_VALUE: str | None = None
def set_default_fingerprint_mode(
mode: FingerprintMode,
custom_value: str | None = None,
) -> None:
"""Configure the fingerprint mode for subsequent ``get_system_fingerprint``
calls. Called once at server startup."""
global _DEFAULT_MODE, _CUSTOM_VALUE
_DEFAULT_MODE = mode
_CUSTOM_VALUE = custom_value
def get_system_fingerprint(vllm_config: Any) -> str | None:
"""Return the fingerprint for ``vllm_config`` using the mode configured by
``set_default_fingerprint_mode``."""
return build_system_fingerprint(vllm_config, _DEFAULT_MODE, _CUSTOM_VALUE)
def build_system_fingerprint(
vllm_config: Any,
mode: FingerprintMode = "full",
custom_value: str | None = None,
) -> str | None:
if mode == "none":
return None
if mode == "custom":
return custom_value
from vllm import __version__ as vllm_version
try:
hash8 = vllm_config.compute_hash()[:8]
except Exception:
hash8 = "nohash"
if mode == "hash":
return f"vllm-{vllm_version}-{hash8}"
# mode == "full"
parts: list[str] = [f"vllm-{vllm_version}"]
pc = getattr(vllm_config, "parallel_config", None)
if pc is not None:
tp = getattr(pc, "tensor_parallel_size", 1)
if tp > 1:
parts.append(f"tp{tp}")
pp = getattr(pc, "pipeline_parallel_size", 1)
if pp > 1:
parts.append(f"pp{pp}")
dp = getattr(pc, "data_parallel_size", 1)
if dp > 1:
parts.append(f"dp{dp}")
if getattr(pc, "enable_expert_parallel", False):
parts.append("ep")
parts.append(hash8)
return "-".join(parts)
@@ -61,9 +61,17 @@ async def init_generate_state(
)
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
from vllm.entrypoints.openai.fingerprint import set_default_fingerprint_mode
from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
from vllm.entrypoints.serve.disagg.serving import ServingTokens
# Applied before any serving class is constructed so that each one picks
# up the chosen mode on its first cache miss.
set_default_fingerprint_mode(
getattr(args, "fingerprint_mode", "full"),
getattr(args, "fingerprint_value", None),
)
if args.tool_server == "demo":
tool_server: ToolServer | None = DemoToolServer()
assert isinstance(tool_server, DemoToolServer)
+1 -10
View File
@@ -472,13 +472,9 @@ class OpenAIServingResponses(OpenAIServing):
context = SimpleContext()
if self.parser and self.parser.reasoning_parser_cls is not None:
chat_template_kwargs = self._effective_chat_template_kwargs(request)
reasoning_parser_kwargs = {
"chat_template_kwargs": chat_template_kwargs,
}
reasoning_parser = self.parser.reasoning_parser_cls(
tokenizer,
chat_template_kwargs=chat_template_kwargs,
chat_template_kwargs=self._effective_chat_template_kwargs(request),
)
if (
isinstance(
@@ -501,9 +497,6 @@ class OpenAIServingResponses(OpenAIServing):
lora_request=lora_request,
priority=request.priority,
trace_headers=trace_headers,
reasoning_parser_kwargs=reasoning_parser_kwargs
if self.parser and self.parser.reasoning_parser_cls is not None
else None,
)
generators.append(generator)
@@ -650,7 +643,6 @@ class OpenAIServingResponses(OpenAIServing):
lora_request: LoRARequest | None = None,
priority: int = 0,
trace_headers: Mapping[str, str] | None = None,
reasoning_parser_kwargs: dict[str, Any] | None = None,
):
max_model_len = self.model_config.max_model_len
@@ -674,7 +666,6 @@ class OpenAIServingResponses(OpenAIServing):
lora_request=lora_request,
trace_headers=trace_headers,
priority=priority,
reasoning_parser_kwargs=reasoning_parser_kwargs,
)
async for res in generator:
+3 -12
View File
@@ -15,10 +15,7 @@ from starlette.datastructures import Headers
from vllm import PoolingParams, PoolingRequestOutput, envs
from vllm.config import VllmConfig
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import (
ChatTemplateConfig,
ChatTemplateContentFormatOption,
)
from vllm.entrypoints.chat_utils import ChatTemplateConfig
from vllm.entrypoints.logger import RequestLogger
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
@@ -48,9 +45,7 @@ class PoolingServingBase(ABC):
models: OpenAIServingModels,
*,
request_logger: RequestLogger | None,
chat_template: str | None = None,
chat_template_content_format: ChatTemplateContentFormatOption = "auto",
trust_request_chat_template: bool = False,
chat_template_config: ChatTemplateConfig,
return_tokens_as_token_ids: bool = False,
log_error_stack: bool = False,
):
@@ -63,11 +58,7 @@ class PoolingServingBase(ABC):
self.request_logger = request_logger
self.return_tokens_as_token_ids = return_tokens_as_token_ids
self.log_error_stack = log_error_stack
self.chat_template_config = ChatTemplateConfig(
chat_template=chat_template,
chat_template_content_format=chat_template_content_format,
trust_request_chat_template=trust_request_chat_template,
)
self.chat_template_config = chat_template_config
# Shared thread pool executor for preprocessing and postprocessing.
self._executor: Executor = models.renderer._executor
+33 -24
View File
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
from vllm.logger import init_logger
from vllm.plugins.io_processors import has_io_processor
from vllm.renderers import BaseRenderer
from vllm.tasks import POOLING_TASKS, SupportedTask
from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
from .base.io_processor import PoolingIOProcessor
from .utils import enable_scoring_api
@@ -43,23 +43,24 @@ def init_pooling_io_processors(
) -> dict[str, PoolingIOProcessor]:
model_config = vllm_config.model_config
processors: dict[str, type[PoolingIOProcessor]] = {}
pooling_task = model_config.get_pooling_task(supported_tasks)
if "classify" in supported_tasks:
if pooling_task == "classify":
from .classify.io_processor import ClassifyIOProcessor
processors["classify"] = ClassifyIOProcessor
if "token_classify" in supported_tasks:
if pooling_task == "token_classify":
from .classify.io_processor import TokenClassifyIOProcessor
processors["token_classify"] = TokenClassifyIOProcessor
if "embed" in supported_tasks:
if pooling_task == "embed":
from .embed.io_processor import EmbedIOProcessor
processors["embed"] = EmbedIOProcessor
if "token_embed" in supported_tasks:
if pooling_task == "token_embed":
from .embed.io_processor import TokenEmbedIOProcessor
processors["token_embed"] = TokenEmbedIOProcessor
@@ -71,15 +72,15 @@ def init_pooling_io_processors(
from .pooling.io_processor import PluginWithIOProcessorPlugins
processors["plugin"] = PluginWithIOProcessorPlugins
elif "plugin" in supported_tasks:
elif pooling_task == "plugin":
from .pooling.io_processor import PluginWithoutIOProcessorPlugins
processors["plugin"] = PluginWithoutIOProcessorPlugins
if enable_scoring_api(supported_tasks, model_config):
score_type = model_config.score_type
from .scoring.io_processor import ScoringIOProcessors
score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
if score_type is not None and score_type in ScoringIOProcessors:
processors[score_type] = ScoringIOProcessors[score_type]
@@ -140,6 +141,10 @@ def init_pooling_state(
request_logger: RequestLogger | None,
supported_tasks: tuple["SupportedTask", ...],
):
model_config = engine_client.model_config
if model_config is None:
return
from vllm.entrypoints.chat_utils import load_chat_template
from vllm.tasks import POOLING_TASKS
@@ -148,8 +153,14 @@ def init_pooling_state(
from .pooling.serving import ServingPooling
from .scoring.serving import ServingScores
model_config = engine_client.model_config
resolved_chat_template = load_chat_template(args.chat_template)
pooling_task = model_config.get_pooling_task(supported_tasks)
chat_template_config = ChatTemplateConfig(
chat_template=resolved_chat_template,
chat_template_content_format=args.chat_template_content_format,
trust_request_chat_template=args.trust_request_chat_template,
)
state.serving_pooling = (
(
@@ -158,9 +169,7 @@ def init_pooling_state(
state.openai_serving_models,
supported_tasks=supported_tasks,
request_logger=request_logger,
chat_template=resolved_chat_template,
chat_template_content_format=args.chat_template_content_format,
trust_request_chat_template=args.trust_request_chat_template,
chat_template_config=chat_template_config,
)
)
if any(t in supported_tasks for t in POOLING_TASKS)
@@ -171,11 +180,9 @@ def init_pooling_state(
engine_client,
state.openai_serving_models,
request_logger=request_logger,
chat_template=resolved_chat_template,
chat_template_content_format=args.chat_template_content_format,
trust_request_chat_template=args.trust_request_chat_template,
chat_template_config=chat_template_config,
)
if "embed" in supported_tasks
if pooling_task == "embed"
else None
)
state.serving_classification = (
@@ -183,21 +190,18 @@ def init_pooling_state(
engine_client,
state.openai_serving_models,
request_logger=request_logger,
chat_template=resolved_chat_template,
chat_template_content_format=args.chat_template_content_format,
trust_request_chat_template=args.trust_request_chat_template,
chat_template_config=chat_template_config,
)
if "classify" in supported_tasks
if pooling_task == "classify"
else None
)
state.serving_scores = (
ServingScores(
engine_client,
state.openai_serving_models,
supported_tasks=supported_tasks,
request_logger=request_logger,
chat_template=resolved_chat_template,
chat_template_content_format=args.chat_template_content_format,
trust_request_chat_template=args.trust_request_chat_template,
chat_template_config=chat_template_config,
enable_flash_late_interaction=getattr(
args, "enable_flash_late_interaction", True
),
@@ -214,7 +218,12 @@ def get_pooling_invocation_types(
# NOTE: Items defined earlier take higher priority
invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
if "embed" in supported_tasks:
if model_config is None:
return invocation_types
pooling_task = model_config.get_pooling_task(supported_tasks)
if pooling_task == "embed":
from .embed.api_router import create_embedding, embedding
from .embed.protocol import EmbeddingRequest
@@ -222,7 +231,7 @@ def get_pooling_invocation_types(
(EmbeddingRequest, (embedding, create_embedding)),
]
if "classify" in supported_tasks:
if pooling_task == "classify":
from .classify.api_router import classify, create_classify
from .classify.protocol import ClassificationRequest

Some files were not shown because too many files have changed in this diff Show More