Compare commits

..
Author SHA1 Message Date
khluuandClaude Opus 4.7 797cdc15f0 [CI] Use uv-installed standalone Python in vllm-base (drop deadsnakes)
The vllm-base stage installs python3.12 via the deadsnakes PPA, which
requires reaching launchpad.net (REST API) or ppa.launchpadcontent.net
(CDN). On the release-v2 build agents, both are unreachable
(Connection timed out to 185.125.190.80), causing the regular Ubuntu
22.04 release images to hang or fail at the apt-get update step.

Switch vllm-base to the same approach the `base` stage already uses:
uv installs a standalone python3.12 (via python-build-standalone,
hosted on astral.sh CDN) into /opt/venv, and we symlink it to
/usr/bin/python3. No Launchpad dependency.

Side benefits:
- Consistent python install path between `base` and `vllm-base`.
- Removes the software-properties-common / add-apt-repository step.
- Drops unused DEADSNAKES_MIRROR_URL / DEADSNAKES_GPGKEY_URL /
  GET_PIP_URL stage args.

Test branch: khluu/release-v0.20.1-uv-python (off releases/v0.20.1).

Signed-off-by: khluu <khluu000@gmail.com>

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-01 13:42:32 -07:00
khluuandClaude Opus 4.7 438f6548c1 [CI] Use deadsnakes mirror for non-Ubuntu24.04 release images
The 4 Ubuntu 22.04 release image builds (CUDA 12.9 / 13.0 x both archs)
hang for ~30min and intermittently fail because
`add-apt-repository ppa:deadsnakes/ppa` hits the launchpad.net REST API
to resolve PPA metadata, and that call has been timing out in CI:

    TimeoutError: [Errno 110] Connection timed out
    ...
    E: Unable to locate package python3.12

Plumb DEADSNAKES_MIRROR_URL and DEADSNAKES_GPGKEY_URL through to those
steps so the Dockerfile takes the mirror branch and writes the apt
source directly, skipping the flaky API call. The mirror points at the
public Launchpad CDN (same packages add-apt-repository would have
configured) and the GPG key comes from keyserver.ubuntu.com.

Ubuntu 24.04 builds are unaffected — they ship python3.12 in default
apt repos and never enter the deadsnakes branch.

Signed-off-by: khluu <khluu000@gmail.com>

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-01 13:25:05 -07:00
21 changed files with 517 additions and 492 deletions
+5 -42
View File
@@ -121,7 +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"
- "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_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 DEADSNAKES_MIRROR_URL=https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu --build-arg DEADSNAKES_GPGKEY_URL='https://keyserver.ubuntu.com/pks/lookup?op=get&search=0xF23C5A6CF475977595C89F51BA6932366A755776' --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"
@@ -134,7 +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"
- "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_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 DEADSNAKES_MIRROR_URL=https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu --build-arg DEADSNAKES_GPGKEY_URL='https://keyserver.ubuntu.com/pks/lookup?op=get&search=0xF23C5A6CF475977595C89F51BA6932366A755776' --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"
@@ -144,7 +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"
- "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_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 --build-arg DEADSNAKES_MIRROR_URL=https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu --build-arg DEADSNAKES_GPGKEY_URL='https://keyserver.ubuntu.com/pks/lookup?op=get&search=0xF23C5A6CF475977595C89F51BA6932366A755776' --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"
@@ -157,7 +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"
- "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_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 --build-arg DEADSNAKES_MIRROR_URL=https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu --build-arg DEADSNAKES_GPGKEY_URL='https://keyserver.ubuntu.com/pks/lookup?op=get&search=0xF23C5A6CF475977595C89F51BA6932366A755776' --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"
@@ -209,7 +209,6 @@ steps:
depends_on: ~
- label: "Build release image - x86_64 - CPU"
key: build-cpu-release-image-x86
depends_on:
- block-cpu-release-image-build
- input-release-version
@@ -228,8 +227,7 @@ steps:
depends_on: ~
- label: "Build release image - arm64 - CPU"
key: build-cpu-release-image-arm64
depends_on:
depends_on:
- block-arm64-cpu-release-image-build
- input-release-version
agents:
@@ -338,41 +336,6 @@ steps:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- block: "Publish release images to DockerHub"
key: block-publish-release-images
depends_on:
- create-multi-arch-manifest
- create-multi-arch-manifest-cuda-12-9
- create-multi-arch-manifest-ubuntu2404
- create-multi-arch-manifest-cuda-12-9-ubuntu2404
- build-rocm-release-image
- input-release-version
# Wait for CPU builds if their block steps were unblocked, so publish
# doesn't race the in-progress CPU build. allow_failure lets publish
# proceed when the operator legitimately leaves the CPU block steps
# unblocked or the CPU build fails.
- step: build-cpu-release-image-x86
allow_failure: true
- step: build-cpu-release-image-arm64
allow_failure: true
if: build.env("NIGHTLY") != "1"
- label: "Publish release images to DockerHub"
depends_on:
- block-publish-release-images
key: publish-release-images-dockerhub
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/publish-release-images.sh"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- group: "Publish wheels"
key: "publish-wheels"
steps:
+93 -1
View File
@@ -8,6 +8,8 @@ if [ -z "${RELEASE_VERSION}" ]; then
RELEASE_VERSION="1.0.0.dev"
fi
ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
To download the wheel (by commit):
\`\`\`
@@ -23,5 +25,95 @@ aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
\`\`\`
Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
To download and upload the image:
\`\`\`
# Download images:
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
# Tag and push images:
## CUDA
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai:latest-x86_64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker push vllm/vllm-openai:latest-x86_64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker push vllm/vllm-openai:latest-aarch64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker push vllm/vllm-openai:latest-aarch64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
## ROCm
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker push vllm/vllm-openai-rocm:latest
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
docker push vllm/vllm-openai-rocm:latest-base
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
## CPU
docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai-cpu:latest-x86_64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker push vllm/vllm-openai-cpu:latest-arm64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
# Create multi-arch manifest:
docker manifest rm vllm/vllm-openai:latest
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest push vllm/vllm-openai:latest
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
docker manifest rm vllm/vllm-openai:latest-cu129
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker manifest push vllm/vllm-openai:latest-cu129
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
docker manifest rm vllm/vllm-openai-cpu:latest || true
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker manifest push vllm/vllm-openai-cpu:latest
docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
\`\`\`
EOF
@@ -1,180 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Publish release Docker images from ECR to DockerHub.
# Pulls per-arch images, tags with latest and versioned tags, pushes them,
# then creates and pushes multi-arch manifests.
set -euo pipefail
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
exit 1
fi
COMMIT="$BUILDKITE_COMMIT"
ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
echo "========================================"
echo "Publishing release images v${RELEASE_VERSION}"
echo " Commit: ${COMMIT}"
echo " ROCm base cache key: ${ROCM_BASE_CACHE_KEY}"
echo "========================================"
# Login to ECR to pull staging images
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
# ---- CUDA (default: 13.0) ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai:latest-x86_64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker push vllm/vllm-openai:latest-aarch64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest rm vllm/vllm-openai:latest || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest push vllm/vllm-openai:latest
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
# ---- CUDA 12.9 ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker push vllm/vllm-openai:latest-x86_64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker push vllm/vllm-openai:latest-aarch64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker manifest rm vllm/vllm-openai:latest-cu129 || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker manifest push vllm/vllm-openai:latest-cu129
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
# ---- Ubuntu 24.04 (CUDA 13.0) ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
docker push vllm/vllm-openai:latest-x86_64-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
docker push vllm/vllm-openai:latest-aarch64-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true
docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
docker manifest push vllm/vllm-openai:latest-ubuntu2404
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404
# ---- Ubuntu 24.04 (CUDA 12.9) ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true
docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404
# ---- ROCm ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker push vllm/vllm-openai-rocm:latest
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
docker push vllm/vllm-openai-rocm:latest-base
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
# ---- CPU ----
# CPU images are behind separate block steps and may not have been built.
# All-or-nothing: inspect both arches first, then either publish everything
# (per-arch + multi-arch manifest) or skip everything. Publishing only one
# arch would leave `:latest-x86_64` pointing at the new release while the
# `:latest` multi-arch manifest still resolves to the previous release.
CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
CPU_X86_AVAILABLE=false
CPU_ARM_AVAILABLE=false
docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true
docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true
if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then
docker pull "${CPU_X86_TAG}"
docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64
docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai-cpu:latest-x86_64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker pull "${CPU_ARM_TAG}"
docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64
docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker push vllm/vllm-openai-cpu:latest-arm64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker manifest rm vllm/vllm-openai-cpu:latest || true
docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker manifest push vllm/vllm-openai-cpu:latest
docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then
echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)"
else
# Partial state: one arch built, the other did not. Fail loudly rather than
# ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch)
# disagree on which release they point at.
echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})."
echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags."
echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended."
exit 1
fi
echo ""
echo "Successfully published release images for v${RELEASE_VERSION}"
@@ -39,11 +39,10 @@ fi
set -x # avoid printing secrets above
# install twine and sdist build prerequisites from pypi
# install twine from pypi
python3 -m venv /tmp/vllm-release-env
source /tmp/vllm-release-env/bin/activate
pip install twine
pip install -r requirements/build/cuda.txt
python3 -m twine --version
# copy release wheels to local directory
+19 -6
View File
@@ -887,14 +887,27 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
uint32_t* shared_ordered =
reinterpret_cast<uint32_t*>(smem_raw + kFixedSmemLarge);
// RadixRowState for multi-CTA cooperative radix.
// Zero-initialization is done host-side via cudaMemsetAsync in topk.cu
// before launch — that gives a stream-ordered happens-before edge for all
// CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA
// __syncthreads) did not provide and which manifested as a race against
// CTA-1+'s first red_release on arrival_counter.
// RadixRowState for multi-CTA cooperative radix
RadixRowState* state = &params.row_states[group_id];
// -- Initialize RadixRowState (only needed if large rows exist) --
if (params.max_seq_len > RADIX_THRESHOLD) {
if (cta_in_group == 0) {
for (uint32_t buf = 0; buf < 3; buf++) {
for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
state->histogram[buf][i] = 0;
}
}
if (tx == 0) {
state->remaining_k = 0;
state->prefix = 0;
state->arrival_counter = 0;
state->output_counter = 0;
}
}
__syncthreads();
}
int barrier_phase = 0;
const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;
-32
View File
@@ -153,38 +153,6 @@ void launch_persistent_topk(const torch::Tensor& logits,
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
"workspace too small, need ", state_bytes, " bytes");
// Zero the per-group RadixRowState region before launch.
//
// Issued UNCONDITIONALLY so the memset is captured as its own node in
// the cudagraph (a separate cudaMemsetAsync node, sequenced before the
// persistent_topk_kernel launch on the same stream). The previous
// host-side guard `if (needs_cooperative)` was evaluated at capture time;
// when capture-time max_seq_len <= RADIX_THRESHOLD (always true under
// FULL_DECODE_ONLY with max_model_len < 32 K) the memset would NOT be
// captured, leaving the workspace state to accumulate across replays.
// That's a latent correctness bug if the runtime data ever takes the
// radix path, and removes one variable while debugging hangs in the
// decode/medium paths.
//
// Cost is sub-microsecond: state_bytes = num_groups * sizeof(RadixRowState)
// is ~3 KB per group, ~100 KB for the largest grids on this hardware.
//
// Why the memset is required (regardless of which path the kernel takes):
// 1. arrival_counter accumulates within a launch and is never reset,
// so a prior call leaves it at a large positive value. Without this
// reset, the very first wait_ge in the next call sees counter >>
// target and returns instantly, breaking the barrier.
// 2. The previous in-kernel init only ran in CTA-0 with intra-CTA
// __syncthreads(), so it had no happens-before edge to CTA-1+'s
// first red_release. cudaMemsetAsync is stream-ordered: the zero
// is globally visible before any CTA runs.
{
cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
state_bytes, stream);
TORCH_CHECK(mz_err == cudaSuccess,
"row_states memset failed: ", cudaGetErrorString(mz_err));
}
P::PersistentTopKParams params;
params.input = logits.data_ptr<float>();
params.output = output.data_ptr<int32_t>();
+14 -33
View File
@@ -478,9 +478,6 @@ FROM ${FINAL_BASE_IMAGE} AS vllm-base
ARG CUDA_VERSION
ARG PYTHON_VERSION
ARG DEADSNAKES_MIRROR_URL
ARG DEADSNAKES_GPGKEY_URL
ARG GET_PIP_URL
ENV DEBIAN_FRONTEND=noninteractive
WORKDIR /vllm-workspace
@@ -490,7 +487,10 @@ 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 system dependencies and Python via uv-managed standalone build.
# This avoids the deadsnakes PPA / Launchpad, which is not always reachable
# from CI build agents on Ubuntu 22.04 base images. Ubuntu 24.04 base images
# could use distro python3.12 directly, but using uv keeps both paths uniform.
RUN apt-get update -y \
&& apt-get install -y --no-install-recommends \
curl \
@@ -500,38 +500,19 @@ RUN apt-get update -y \
libxext6 \
libgl1 \
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/* \
&& curl -LsSf https://astral.sh/uv/install.sh | sh \
&& $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
&& rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
&& ln -s /opt/venv/bin/python3 /usr/bin/python3 \
&& ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
&& ln -s /opt/venv/bin/pip /usr/bin/pip \
&& python3 --version && python3 -m pip --version
# Activate venv: ensures `uv`, `pip`, `python` resolve to the /opt/venv install
ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
ENV VIRTUAL_ENV="/opt/venv"
# Install CUDA development tools for runtime JIT compilation
# (FlashInfer, DeepGEMM, EP kernels all require compilation at runtime)
RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
+228 -4
View File
@@ -3,12 +3,11 @@
"""
Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
Two paths tested:
Four test functions cover five paths:
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
These serve as golden references for validating the future fused
compressor+quant+cache kernel.
C) DeepseekV4 Attention magnitude range: correctness across small/large values
D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
"""
import math
@@ -21,6 +20,12 @@ 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):
@@ -309,3 +314,222 @@ 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,6 +30,56 @@ 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,
@@ -37,6 +87,7 @@ 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(
@@ -49,22 +100,33 @@ 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)
weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
return q_fp8, weights_out
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
@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):
def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
device = "cuda"
torch.manual_seed(0)
@@ -77,21 +139,32 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
softmax_scale = HEAD_DIM**-0.5
head_scale = N_HEAD**-0.5
q_fp8_ref, weights_ref = _reference(
positions, q, cos_sin_cache, weights, softmax_scale, head_scale
q_quant_ref, weights_ref = _reference(
positions, q, 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
q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
)
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_fp8_ref.view(torch.int8)
fused_bits = q_fp8_fused.view(torch.int8)
ref_bits = q_quant_ref.view(torch.int8)
fused_bits = q_quant_fused.view(torch.int8)
assert torch.equal(ref_bits, fused_bits), (
f"q_fp8 mismatch: "
f"q_quant_fused 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()}"
@@ -432,52 +432,3 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
)
== 15
)
def test_predictor_matches_allocator_blocks_calculation_with_admission_cap():
"""In forward steps, `get_num_blocks_to_allocate` must return exactly what
`allocate_new_blocks` will pull; otherwise `block_pool.get_new_blocks`
raises `ValueError: Cannot get N free blocks from the pool`.
"""
block_size = 2
sliding_window = 8 # 4-block live window
cap = sliding_window // block_size
spec = SlidingWindowSpec(
block_size=block_size,
num_kv_heads=1,
head_size=1,
dtype=torch.float32,
sliding_window=sliding_window,
)
block_pool = BlockPool(
num_gpu_blocks=100, enable_caching=True, hash_block_size=block_size
)
manager = SlidingWindowManager(
spec,
block_pool=block_pool,
enable_caching=False,
kv_cache_group_id=0,
max_admission_blocks_per_request=cap,
)
request_id = "req"
total_computed = 0
# Walk through request forward steps. Check num_blocks returned by
# `get_num_blocks_to_allocate` matches what `allocate_new_blocks` pulls
for num_tokens in (4, 8, 12, 16):
predicted = manager.get_num_blocks_to_allocate(
request_id=request_id,
num_tokens=num_tokens,
new_computed_blocks=[],
total_computed_tokens=total_computed,
num_tokens_main_model=num_tokens,
)
new_blocks = manager.allocate_new_blocks(
request_id, num_tokens=num_tokens, num_tokens_main_model=num_tokens
)
assert predicted == len(new_blocks), (
f"num_tokens={num_tokens}: predictor returned {predicted} "
f"but allocator pulled {len(new_blocks)}"
)
total_computed = num_tokens
+4 -4
View File
@@ -245,7 +245,7 @@ if TYPE_CHECKING:
VLLM_DEBUG_WORKSPACE: bool = False
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 4096
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
VLLM_USE_V2_MODEL_RUNNER: bool = False
VLLM_LOG_MODEL_INSPECTION: bool = False
@@ -1669,10 +1669,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
# tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs
# and overlap is a 5-45% win; above it the FP8 GEMM saturates the device
# and the cross-stream sync becomes pure overhead. Set to 0 to disable
# the multi-stream path entirely. See #PR 41526 for the empirical result
# for the default value of 1024 tokens.
# the multi-stream path entirely. Empirical crossover on B300 (148 SMs)
# is ~4096; B200 (132 SMs) is expected ~3072.
"VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int(
os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "1024")
os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "4096")
),
# Format for saving torch.compile cache artifacts
# - "binary": saves as binary file
@@ -93,7 +93,6 @@ def _moe_forward(
shared_experts_input: torch.Tensor | None,
input_ids: torch.Tensor | None,
layer_name: _layer_name_type,
hidden_dim_unpadded: int,
) -> torch.Tensor:
layer = get_layer_from_name(_resolve_layer_name(layer_name))
return layer.runner._forward_impl(
@@ -111,14 +110,7 @@ def _moe_forward_fake(
shared_experts_input: torch.Tensor | None,
input_ids: torch.Tensor | None,
layer_name: _layer_name_type,
hidden_dim_unpadded: int,
) -> torch.Tensor:
# `hidden_dim_unpadded > 0` only on the TRT-LLM MXFP4 path, where the
# real kernel writes narrower than `hidden_states.shape[-1]`. Plumbed
# as an op arg (not peeked from the layer registry) to keep the fake
# a pure shape function of its inputs and preserve subgraph dedup.
if hidden_dim_unpadded > 0:
return hidden_states.new_empty((*hidden_states.shape[:-1], hidden_dim_unpadded))
return torch.empty_like(hidden_states)
@@ -128,7 +120,6 @@ def _moe_forward_shared(
shared_experts_input: torch.Tensor | None,
input_ids: torch.Tensor | None,
layer_name: _layer_name_type,
hidden_dim_unpadded: int,
) -> tuple[torch.Tensor, torch.Tensor]:
layer = get_layer_from_name(_resolve_layer_name(layer_name))
return layer.runner._forward_impl(
@@ -146,17 +137,13 @@ def _moe_forward_shared_fake(
shared_experts_input: torch.Tensor | None,
input_ids: torch.Tensor | None,
layer_name: _layer_name_type,
hidden_dim_unpadded: int,
) -> tuple[torch.Tensor, torch.Tensor]:
# `fused_out`: see `_moe_forward_fake` for hidden_dim_unpadded semantics.
# `shared_out`: matches `shared_experts_input` if provided (latent MoE),
# else `hidden_states`.
if hidden_dim_unpadded > 0:
fused_out = hidden_states.new_empty(
(*hidden_states.shape[:-1], hidden_dim_unpadded)
)
else:
fused_out = torch.empty_like(hidden_states)
# Output shapes:
# - fused_out: same as hidden_states (routed experts use transformed size)
# - shared_out: same as shared_experts_input if provided, else same as
# hidden_states
# (For latent MoE: shared experts use original hidden_size, not latent size)
fused_out = torch.empty_like(hidden_states)
if shared_experts_input is not None:
shared_out = torch.empty_like(shared_experts_input)
else:
@@ -402,29 +389,6 @@ class MoERunner(MoERunnerInterface):
return "from_forward_context"
return self.layer_name
def _trtllm_mxfp4_unpadded_dim(self) -> int:
"""Return ``hidden_dim_unpadded`` when the active backend is TRT-LLM
MXFP4 (whose kernel writes narrower than the padded
``hidden_states.shape[-1]``), else 0. Other MXFP4 backends (notably
Cutlass MXFP4 MXFP8) write the full padded width, so
``moe_config.hidden_dim_unpadded`` alone is insufficient: it encodes
the model's logical hidden, not whether the kernel narrows. Computed
caller-side and passed as an op arg; doing the isinstance check
inside the fake would specialize per ``layer_name`` and break
subgraph dedup for identical-architecture models (e.g. Phi-MoE).
"""
from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
TrtLlmMxfp4ExpertsBase,
)
moe_kernel = getattr(self.quant_method, "moe_kernel", None)
fused_experts = getattr(
getattr(moe_kernel, "impl", None), "fused_experts", None
)
if isinstance(fused_experts, TrtLlmMxfp4ExpertsBase):
return self.moe_config.hidden_dim_unpadded or self.moe_config.hidden_dim
return 0
def _maybe_pad_hidden_states(
self,
shared_experts_input: torch.Tensor | None,
@@ -613,7 +577,6 @@ class MoERunner(MoERunnerInterface):
shared_experts_input,
input_ids,
self._encode_layer_name(),
self._trtllm_mxfp4_unpadded_dim(),
)
#
@@ -320,7 +320,7 @@ def sparse_attn_indexer(
num_rows = logits.shape[0]
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048):
if current_platform.is_cuda() and topk_tokens in (512, 2048):
workspace_manager = current_workspace_manager()
(topk_workspace,) = workspace_manager.get_simultaneous(
((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+6 -17
View File
@@ -715,15 +715,12 @@ class DeepseekV4MoE(nn.Module):
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.prefix = prefix
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
raise NotImplementedError(
"DeepSeek V4 MegaMoE currently requires expert parallel. "
"Enable it with --enable-expert-parallel, or pick a different "
"moe backend."
if vllm_config.parallel_config.enable_expert_parallel:
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
else:
self.use_mega_moe = False
self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
self.hidden_size = config.hidden_size
@@ -1227,15 +1224,7 @@ class DeepseekV4Model(nn.Module):
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.config = config
self.use_mega_moe = (
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
)
if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
raise NotImplementedError(
"DeepSeek V4 MegaMoE currently requires expert parallel. "
"Enable it with --enable-expert-parallel, or pick a different "
"moe backend."
)
self.vocab_size = config.vocab_size
self.hc_eps = config.hc_eps
self.hc_mult = config.hc_mult
@@ -1778,6 +1778,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
return None # If vision tower is skipped
if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
return None
if num_tokens > self.deepstack_input_embeds_num_tokens:
raise ValueError(
"Requested more deepstack tokens than available in buffer: "
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
)
# get deepstack_input_embeds from buffer, and clear the buffer
return IntermediateTensors(
@@ -1819,6 +1824,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
# clear deepstack_input_embeds in buffer
if num_tokens > 0:
if num_tokens > self.deepstack_input_embeds_num_tokens:
raise ValueError(
"Requested to clear more deepstack tokens than available in "
"buffer: "
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
)
for idx in range(self.deepstack_num_level):
self.deepstack_input_embeds[idx][:num_tokens].zero_()
self.deepstack_input_embeds_num_tokens = 0
+11
View File
@@ -1707,6 +1707,11 @@ class Qwen3VLForConditionalGeneration(
return None # If vision tower is skipped
if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
return None
if num_tokens > self.deepstack_input_embeds_num_tokens:
raise ValueError(
"Requested more deepstack tokens than available in buffer: "
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
)
# get deepstack_input_embeds from buffer, and clear the buffer
return IntermediateTensors(
@@ -1748,6 +1753,12 @@ class Qwen3VLForConditionalGeneration(
# clear deepstack_input_embeds in buffer
if num_tokens > 0:
if num_tokens > self.deepstack_input_embeds_num_tokens:
raise ValueError(
"Requested to clear more deepstack tokens than available in "
"buffer: "
f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
)
for idx in range(self.deepstack_num_level):
self.deepstack_input_embeds[idx][:num_tokens].zero_()
self.deepstack_input_embeds_num_tokens = 0
@@ -21,7 +21,7 @@ and N_QUANT_BLOCKS ue8m0 bytes.
from vllm.triton_utils import tl, triton
from .fused_indexer_q import _e2m1_nibble
from .fused_indexer_q import _fp32x2_to_fp4x2
# =============================================================================
@@ -566,18 +566,18 @@ def _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn(
tl.max(tl.abs(even_2d), axis=1),
tl.max(tl.abs(odd_2d), axis=1),
)
amax = tl.maximum(amax, 1e-4)
amax = tl.maximum(amax, 6.0 * (2**-126))
# ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored as (exp + 127) byte.
log2_ratio = tl.ceil(tl.log2(amax / 6.0))
log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0)))
log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
inv_scale = tl.exp2(-log2_ratio)
ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS]
inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1))
lo_nib = _e2m1_nibble(even_2d * inv_scale_col) # (N_BLOCKS, HALF_BLOCK) uint8
hi_nib = _e2m1_nibble(odd_2d * inv_scale_col)
packed = lo_nib | (hi_nib << 4)
packed = _fp32x2_to_fp4x2(
even_2d * inv_scale_col, odd_2d * inv_scale_col
) # (N_BLOCKS, HALF_BLOCK) uint8
packed_flat = tl.reshape(packed, (TOKEN_STRIDE,))
tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
@@ -24,36 +24,22 @@ def _get_cos_sin(
@triton.jit
def _e2m1_nibble(x):
"""Quantize fp32 x (already scale-divided) to E2M1 4-bit nibble in uint8.
Matches torch.bucketize with boundaries
[0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0] and right=False (each boundary
belongs to the lower bucket), plus sign bit."""
abs_x = tl.minimum(tl.abs(x), 6.0)
code = tl.where(
abs_x <= 0.25,
0.0,
tl.where(
abs_x <= 0.75,
1.0,
tl.where(
abs_x <= 1.25,
2.0,
tl.where(
abs_x <= 1.75,
3.0,
tl.where(
abs_x <= 2.5,
4.0,
tl.where(abs_x <= 3.5, 5.0, tl.where(abs_x <= 5.0, 6.0, 7.0)),
),
),
),
),
)
code_u8 = code.to(tl.uint8)
sign = ((x < 0) & (code_u8 != 0)).to(tl.uint8)
return code_u8 | (sign << 3)
def _fp32x2_to_fp4x2(x_lo, x_hi):
# NOTE: $1 is high nibble, $2 is low nibble
return tl.inline_asm_elementwise(
"""
{
.reg .b8 tmp;
cvt.rn.satfinite.e2m1x2.f32 tmp, $1, $2;
cvt.u32.u8 $0, tmp;
}
""",
constraints="=r,f,f",
args=[x_hi, x_lo],
dtype=tl.uint32,
is_pure=True,
pack=1,
).to(tl.uint8)
@triton.jit
@@ -65,17 +51,16 @@ def _quantize_mxfp4_pair(x_lo, x_hi):
- ue8m0 : scalar uint8 (block scale = 2^(ue8m0 - 127))
"""
amax = tl.maximum(tl.max(tl.abs(x_lo)), tl.max(tl.abs(x_hi)))
amax = tl.maximum(amax, 1e-4)
# 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
amax = tl.maximum(amax, 6.0 * (2**-126))
# ue8m0 block scale: 2^ceil(log2(amax/6.0)).
log2_ratio = tl.math.ceil(tl.math.log2(amax / 6.0))
log2_ratio = tl.math.ceil(tl.math.log2(amax * (1.0 / 6.0)))
log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
scale = tl.math.exp2(log2_ratio)
ue8m0 = (log2_ratio + 127.0).to(tl.uint8)
inv_scale = 1.0 / scale
lo_nib = _e2m1_nibble(x_lo * inv_scale)
hi_nib = _e2m1_nibble(x_hi * inv_scale)
packed = lo_nib | (hi_nib << 4)
packed = _fp32x2_to_fp4x2(x_lo * inv_scale, x_hi * inv_scale)
return packed, ue8m0
+1 -12
View File
@@ -85,7 +85,6 @@ class KVCacheCoordinator(ABC):
num_encoder_tokens: int,
total_computed_tokens: int,
num_tokens_main_model: int,
apply_admission_cap: bool = False,
) -> int:
"""
Get the number of blocks needed to be allocated for the request.
@@ -102,10 +101,6 @@ class KVCacheCoordinator(ABC):
num_tokens_main_model: The number of tokens for the main model (aka target
model in spec decode). w/o spec decode, it is num_tokens;
with spec decode, it is num_tokens - num_lookahead_tokens.
apply_admission_cap: If True, apply the recycling-aware
per-request admission cap (SWA / chunked-local). Set only by
the full-sequence admission gate; per-step allocation must
leave it False so the predictor matches `allocate_new_blocks`.
Returns:
The number of blocks to allocate.
@@ -116,12 +111,7 @@ class KVCacheCoordinator(ABC):
# For cross-attention, we issue a single static allocation
# of blocks based on the number of encoder input tokens.
num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
request_id,
num_encoder_tokens,
[],
0,
num_encoder_tokens,
apply_admission_cap=apply_admission_cap,
request_id, num_encoder_tokens, [], 0, num_encoder_tokens
)
else:
num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
@@ -130,7 +120,6 @@ class KVCacheCoordinator(ABC):
new_computed_blocks[i],
total_computed_tokens,
num_tokens_main_model,
apply_admission_cap=apply_admission_cap,
)
return num_blocks_to_allocate
-1
View File
@@ -257,7 +257,6 @@ class KVCacheManager:
num_encoder_tokens=num_encoder_tokens,
total_computed_tokens=total_computed_tokens,
num_tokens_main_model=full_num_tokens,
apply_admission_cap=True,
)
return num_blocks_to_allocate <= self.block_pool.get_num_free_blocks()
+1 -7
View File
@@ -92,7 +92,6 @@ class SingleTypeKVCacheManager(ABC):
new_computed_blocks: Sequence[KVCacheBlock],
total_computed_tokens: int,
num_tokens_main_model: int,
apply_admission_cap: bool = False,
) -> int:
"""
Get the number of blocks needed to be allocated for the request.
@@ -108,16 +107,13 @@ class SingleTypeKVCacheManager(ABC):
num_tokens_main_model: The number of tokens for the main model (aka target
model in spec decode). w/o spec decode, it is num_tokens;
with spec decode, it is num_tokens - num_lookahead_tokens.
apply_admission_cap: If True, clamp by `num_required_blocks` by
`_max_admission_blocks_per_request`for recycling-aware specs
(SWA, chunked-local).
Returns:
The number of blocks to allocate.
"""
num_required_blocks = cdiv(num_tokens, self.block_size)
if apply_admission_cap and self._max_admission_blocks_per_request is not None:
if self._max_admission_blocks_per_request is not None:
# Recycling-aware specs (SWA, chunked-local) cap the per-request
# reservation here so admission matches the startup pool sizer
# (`SlidingWindowSpec.max_admission_blocks_per_request` / its
@@ -897,7 +893,6 @@ class MambaManager(SingleTypeKVCacheManager):
new_computed_blocks: Sequence[KVCacheBlock],
total_computed_tokens: int,
num_tokens_main_model: int,
apply_admission_cap: bool = False,
) -> int:
assert isinstance(self.kv_cache_spec, MambaSpec)
if (
@@ -922,7 +917,6 @@ class MambaManager(SingleTypeKVCacheManager):
new_computed_blocks,
total_computed_tokens,
num_tokens_main_model,
apply_admission_cap=apply_admission_cap,
)
else:
# We don't allocate blocks for lookahead tokens in align mode, because if