[Bugfix] Fix DeepSeek V4 MTP HC state handling (#42320)

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
(cherry picked from commit f1cc7aad3c)

.buildkite/ci_config.yaml

@@ -8,6 +8,7 @@ run_all_patterns:
   - "CMakeLists.txt"
   - "requirements/common.txt"
   - "requirements/cuda.txt"
+  - "requirements/kv_connectors.txt"
   - "requirements/build/cuda.txt"
   - "requirements/test/cuda.txt"
   - "setup.py"

.buildkite/release-pipeline.yaml

@@ -28,6 +28,7 @@ steps:
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -41,6 +42,7 @@ steps:
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -54,6 +56,7 @@ steps:
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -67,6 +70,7 @@ steps:
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -80,6 +84,7 @@ steps:
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -93,6 +98,7 @@ steps:
           - "mkdir artifacts"
           - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
           - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -138,6 +144,7 @@ steps:
           # 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"
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
 
       - label: "Build release image - aarch64 - CUDA 13.0"
         depends_on: ~
@@ -160,6 +167,7 @@ steps:
               --progress plain \
               -f docker/Dockerfile .
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
 
       - label: "Build release image - x86_64 - CUDA 12.9"
         depends_on: ~
@@ -184,6 +192,7 @@ steps:
           # 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"
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
 
       - label: "Build release image - aarch64 - CUDA 12.9"
         depends_on: ~
@@ -205,6 +214,7 @@ steps:
               --progress plain \
               -f docker/Dockerfile .
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
 
       - label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
         depends_on: ~
@@ -231,6 +241,7 @@ steps:
           - "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"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
 
       - label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
         depends_on: ~
@@ -255,6 +266,7 @@ steps:
               --progress plain \
               -f docker/Dockerfile .
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
 
       - label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
         depends_on: ~
@@ -280,6 +292,7 @@ steps:
           - "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"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
 
       - label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
         depends_on: ~
@@ -303,6 +316,7 @@ steps:
               --progress plain \
               -f docker/Dockerfile .
           - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
 
       - block: "Build release image for x86_64 CPU"
         key: block-cpu-release-image-build
@@ -320,6 +334,7 @@ steps:
           - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
           - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -339,6 +354,7 @@ steps:
           - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
           - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
           - "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
         env:
           DOCKER_BUILDKIT: "1"
 
@@ -356,15 +372,7 @@ steps:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
           - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
-
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"'
-      - label: "Annotate release workflow - CUDA 13.0"
-        depends_on:
-          - create-multi-arch-manifest
-        id: annotate-release-workflow
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/annotate-release.sh"
 
       - label: "Create multi-arch manifest - CUDA 12.9"
         depends_on:
@@ -377,6 +385,7 @@ steps:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
           - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"'
 
       - label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
         depends_on:
@@ -389,6 +398,7 @@ steps:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
           - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"'
 
       - label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
         depends_on:
@@ -401,6 +411,7 @@ steps:
           - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
           - "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
           - "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
+          - 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"'
 
       - label: "Publish nightly multi-arch image to DockerHub"
         depends_on:
@@ -438,59 +449,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:
-      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
-        key: block-upload-release-wheels
-        depends_on:
-          - input-release-version
-          - build-wheels
-
-      - label: "Upload release wheels to PyPI"
-        depends_on:
-          - block-upload-release-wheels
-        id: upload-release-wheels
-        agents:
-          queue: small_cpu_queue_release
-        commands:
-          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
-
   # =============================================================================
   # ROCm Release Pipeline (x86_64 only)
   # =============================================================================
@@ -847,3 +805,60 @@ steps:
     env:
       DOCKER_BUILDKIT: "1"
       DOCKERHUB_USERNAME: "vllmbot"
+
+  # =============================================================================
+  # Publish to DockerHub and PyPI (at the end so all builds complete first)
+  # =============================================================================
+
+  - 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:
+      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
+        key: block-upload-release-wheels
+        depends_on:
+          - input-release-version
+          - build-wheels
+
+      - label: "Upload release wheels to PyPI"
+        depends_on:
+          - block-upload-release-wheels
+        id: upload-release-wheels
+        agents:
+          queue: small_cpu_queue_release
+        commands:
+          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"

.buildkite/scripts/annotate-build-artifact.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Append a build artifact line to the Buildkite annotation.
+# Usage: annotate-build-artifact.sh <label> <value>
+set -e
+echo "- **${1}**: \`${2}\`" | \
+  buildkite-agent annotate --append --style 'info' --context 'release-artifacts'

.buildkite/scripts/annotate-release.sh

@@ -1,27 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-# Get release version, default to 1.0.0.dev for nightly/per-commit builds
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null | sed 's/^v//')
-if [ -z "${RELEASE_VERSION}" ]; then
-  RELEASE_VERSION="1.0.0.dev"
-fi
-
-buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
-To download the wheel (by commit):
-\`\`\`
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
-
-(Optional) For CUDA 12.9:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
-
-(Optional) For CPU:
-aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
-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.
-EOF

.buildkite/scripts/upload-release-wheels-pypi.sh

@@ -39,11 +39,17 @@ fi
 
 set -x # avoid printing secrets above
 
-# install twine and sdist build prerequisites from pypi
+# install uv if not already available
-python3 -m venv /tmp/vllm-release-env
+if ! command -v uv &> /dev/null; then
+  curl -LsSf https://astral.sh/uv/install.sh | UV_VERSION=0.11.14 sh
+  export PATH="$HOME/.local/bin:$PATH"
+fi
+
+# install twine and sdist build prerequisites using uv with Python 3.12
+uv venv --python 3.12 /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
-pip install twine
+uv pip install twine
-pip install -r requirements/build/cuda.txt
+uv pip install -r requirements/build/cuda.txt
 python3 -m twine --version
 
 # copy release wheels to local directory

.buildkite/test_areas/kernels.yaml

@@ -105,7 +105,11 @@ steps:
   device: h100
   num_devices: 1
   source_file_dependencies:
+  - cmake/external_projects/deepgemm.cmake
   - tools/install_deepgemm.sh
+  - tools/build_deepgemm_C.py
+  - tools/setup_deepgemm_pythons.sh
+  - tools/check_wheel_deepgemm.py
   - vllm/utils/deep_gemm.py
   - vllm/model_executor/layers/fused_moe
   - vllm/model_executor/layers/quantization
@@ -115,6 +119,7 @@ steps:
   - tests/kernels/attention/test_deepgemm_attention.py
   - tests/quantization/test_cutlass_w4a16.py
   commands:
+    - python3 ../tools/check_wheel_deepgemm.py
     - 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

.github/workflows/scripts/build.sh

@@ -9,6 +9,9 @@ PATH=${cuda_home}/bin:$PATH
 LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
 
 # Install requirements
+if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
+    sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
+fi
 $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
 
 # Limit the number of parallel jobs to avoid OOM

benchmarks/attention_benchmarks/configs/mla_decode.yaml

@@ -53,6 +53,7 @@ backends:
   - FLASHINFER_MLA
   - FLASH_ATTN_MLA  # Hopper only
   - FLASHMLA        # Hopper only
+  - TOKENSPEED_MLA  # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)
 
 device: "cuda:0"
 repeats: 100

benchmarks/attention_benchmarks/configs/mla_prefill.yaml

@@ -3,6 +3,7 @@
 # Compares all available MLA prefill backends:
 #   FA backends:  fa2, fa3, fa4 (FlashAttention versions)
 #   Non-FA:       flashinfer, cudnn, trtllm (Blackwell-only, require flashinfer)
+#   CuTe DSL:     tokenspeed (Blackwell + R1 dims, requires tokenspeed_mla)
 #
 # Uses cutlass_mla as the decode backend for impl construction
 # (only the prefill path is exercised).
@@ -120,6 +121,7 @@ prefill_backends:
   - flashinfer
   - cudnn
   - trtllm
+  - tokenspeed
 
 device: "cuda:0"
 repeats: 20

benchmarks/attention_benchmarks/mla_runner.py

@@ -179,19 +179,27 @@ def create_minimal_vllm_config(
 
     if prefill_backend is not None:
         prefill_cfg = get_prefill_backend_config(prefill_backend)
-        if prefill_cfg["flash_attn_version"] is not None:
+        if prefill_cfg.get("mla_prefill_backend_enum") is not None:
-            vllm_config.attention_config.flash_attn_version = prefill_cfg[
+            # Registry-based backends bypass the deprecated boolean flags.
-                "flash_attn_version"
+            from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackendEnum
+
+            vllm_config.attention_config.mla_prefill_backend = MLAPrefillBackendEnum[
+                prefill_cfg["mla_prefill_backend_enum"]
             ]
-        vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
+        else:
-            "disable_flashinfer_prefill"
+            if prefill_cfg["flash_attn_version"] is not None:
-        ]
+                vllm_config.attention_config.flash_attn_version = prefill_cfg[
-        vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
+                    "flash_attn_version"
-            "use_cudnn_prefill"
+                ]
-        ]
+            vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
-        vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg[
+                "disable_flashinfer_prefill"
-            "use_trtllm_ragged_deepseek_prefill"
+            ]
-        ]
+            vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
+                "use_cudnn_prefill"
+            ]
+            vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = (
+                prefill_cfg["use_trtllm_ragged_deepseek_prefill"]
+            )
 
     return vllm_config
 
@@ -223,22 +231,17 @@ _PREFILL_BACKEND_CONFIG: dict[str, dict] = {
         "use_trtllm_ragged_deepseek_prefill": False,
     },
     "flashinfer": {
-        "flash_attn_version": None,
+        "mla_prefill_backend_enum": "FLASHINFER",
-        "disable_flashinfer_prefill": False,
-        "use_cudnn_prefill": False,
-        "use_trtllm_ragged_deepseek_prefill": False,
     },
     "cudnn": {
-        "flash_attn_version": None,
+        # cuDNN prefill backend was removed; AttentionConfig raises on use.
-        "disable_flashinfer_prefill": True,
+        "mla_prefill_backend_enum": "FLASHINFER",
-        "use_cudnn_prefill": True,
-        "use_trtllm_ragged_deepseek_prefill": False,
     },
     "trtllm": {
-        "flash_attn_version": None,
+        "mla_prefill_backend_enum": "TRTLLM_RAGGED",
-        "disable_flashinfer_prefill": True,
+    },
-        "use_cudnn_prefill": False,
+    "tokenspeed": {
-        "use_trtllm_ragged_deepseek_prefill": True,
+        "mla_prefill_backend_enum": "TOKENSPEED_MLA",
     },
 }
 
@@ -625,6 +628,21 @@ def _create_backend_impl(
     # Create mock layer
     layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)
 
+    # Attach a prefill backend (MLAAttention does this in __init__; the metadata
+    # builder reads layer.prefill_backend from static_forward_context).
+    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+
+    prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+    layer.prefill_backend = prefill_backend_cls(
+        num_heads=mla_dims["num_q_heads"],
+        scale=(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]) ** -0.5,
+        kv_lora_rank=mla_dims["kv_lora_rank"],
+        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
+        qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
+        v_head_dim=mla_dims["v_head_dim"],
+        vllm_config=vllm_config,
+    )
+
     # Create builder instance if needed
     builder_instance = None
     if builder_class:
@@ -961,19 +979,6 @@ def _run_mla_benchmark_batched(
     results = []
 
     with set_current_vllm_config(vllm_config):
-        # Clear cached prefill backend detection functions so they re-evaluate
-        # with the current VllmConfig. These are @functools.cache decorated and
-        # would otherwise return stale results from a previous backend's config.
-        from vllm.model_executor.layers.attention.mla_attention import (
-            use_cudnn_prefill,
-            use_flashinfer_prefill,
-            use_trtllm_ragged_deepseek_prefill,
-        )
-
-        use_flashinfer_prefill.cache_clear()
-        use_cudnn_prefill.cache_clear()
-        use_trtllm_ragged_deepseek_prefill.cache_clear()
-
         # Create backend impl, layer, builder, and indexer (reused across benchmarks)
         impl, layer, builder_instance, indexer = _create_backend_impl(
             backend_cfg,
@@ -985,36 +990,35 @@ def _run_mla_benchmark_batched(
             kv_cache_dtype=kv_cache_dtype,
         )
 
-        # Verify the actual prefill backend matches what was requested
+        # Verify the actual prefill backend matches what was requested. The
+        # selector + impl construction already raise on misuse; here we just
+        # check the resolved class against the requested name as a sanity guard.
         if prefill_backend is not None:
-            prefill_cfg = get_prefill_backend_config(prefill_backend)
+            expected_class = {
-            fa_version = prefill_cfg["flash_attn_version"]
+                "fa2": "FlashAttnPrefillBackend",
-
+                "fa3": "FlashAttnPrefillBackend",
-            if fa_version is not None:
+                "fa4": "FlashAttnPrefillBackend",
-                # FA backend: verify the impl's FA version
+                "flashinfer": "FlashInferPrefillBackend",
-                actual_fa_version = getattr(impl, "vllm_flash_attn_version", None)
+                "trtllm": "TrtllmRaggedPrefillBackend",
+                "tokenspeed": "TokenspeedMLAPrefillBackend",
+            }.get(prefill_backend)
+            actual_class = type(getattr(layer, "prefill_backend", None)).__name__
+            if expected_class and actual_class != expected_class:
+                raise RuntimeError(
+                    f"Prefill backend '{prefill_backend}' requested "
+                    f"{expected_class}, got {actual_class}. Check "
+                    f"attention_config plumbing or installed deps."
+                )
+            if prefill_backend in {"fa2", "fa3", "fa4"}:
+                fa_version = int(prefill_backend[2:])
+                actual_fa_version = getattr(
+                    layer.prefill_backend, "vllm_flash_attn_version", None
+                )
                 if actual_fa_version != fa_version:
                     raise RuntimeError(
                         f"Prefill backend '{prefill_backend}' requested FA "
-                        f"version {fa_version}, but the impl is using FA "
+                        f"version {fa_version}, got "
-                        f"version {actual_fa_version}. Check "
+                        f"{actual_fa_version} on {actual_class}."
-                        f"vllm/v1/attention/backends/fa_utils.py."
-                    )
-            else:
-                # Non-FA backend: verify the builder picked the right path
-                expected_flags = {
-                    "flashinfer": "_use_fi_prefill",
-                    "cudnn": "_use_cudnn_prefill",
-                    "trtllm": "_use_trtllm_ragged_prefill",
-                }
-                flag_name = expected_flags.get(prefill_backend)
-                if flag_name and not getattr(builder_instance, flag_name, False):
-                    raise RuntimeError(
-                        f"Prefill backend '{prefill_backend}' was requested "
-                        f"but the metadata builder did not enable it. This "
-                        f"usually means a dependency is missing (e.g., "
-                        f"flashinfer not installed) or the platform doesn't "
-                        f"support it."
                     )
 
         # Run each benchmark with the shared impl

cmake/external_projects/deepgemm.cmake

@@ -53,48 +53,67 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
 if(DEEPGEMM_ARCHS)
   message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
 
-  find_package(CUDAToolkit REQUIRED)
+  # Build _C once per interpreter in DEEPGEMM_PYTHON_INTERPRETERS (":"-
+  # separated paths) so the wheel imports cleanly on every supported Python.
+  # Unset → fall back to the build interpreter (editable / source builds).
+  # The compile is delegated to tools/build_deepgemm_C.py and always runs
+  # against the build interpreter's torch — target Pythons don't need torch.
+  # Note: empty-but-set env vars are still DEFINED in cmake; treat empty as
+  # unset so an empty interpreter list falls back to the build interpreter
+  # rather than silently skipping the per-Python build.
+  if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
+    string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
+  else()
+    set(_dg_pythons "${Python_EXECUTABLE}")
+  endif()
+  message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
 
-  #
+  # Header set fed to add_custom_command's DEPENDS so a header-only edit
-  # Build the _C pybind11 extension from DeepGEMM's C++ source.
+  # (in upstream DeepGEMM or its vendored cutlass/fmt) re-triggers the
-  # This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
+  # rebuild. add_custom_command does no implicit header scanning, unlike
-  #
+  # add_library.
-  Python_add_library(_deep_gemm_C MODULE WITH_SOABI
+  file(GLOB_RECURSE _dg_headers
-    "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")
+    "${deepgemm_SOURCE_DIR}/csrc/*.h"
+    "${deepgemm_SOURCE_DIR}/csrc/*.hpp"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
+    "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
 
-  # The pybind11 module name must be _C to match DeepGEMM's Python imports.
+  set(_dg_markers)
-  set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
+  set(_dg_seen_soabis)
-
+  foreach(_pybin IN LISTS _dg_pythons)
-  target_compile_definitions(_deep_gemm_C PRIVATE
+    execute_process(
-    "-DTORCH_EXTENSION_NAME=_C")
+      COMMAND "${_pybin}" -c
-
+        "import sysconfig; print(sysconfig.get_config_var('SOABI'))"
-  target_include_directories(_deep_gemm_C PRIVATE
+      OUTPUT_VARIABLE _dg_soabi
-    "${deepgemm_SOURCE_DIR}/csrc"
+      OUTPUT_STRIP_TRAILING_WHITESPACE
-    "${deepgemm_SOURCE_DIR}/deep_gemm/include"
+      COMMAND_ERROR_IS_FATAL ANY)
-    "${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
+    # Dedup so duplicate paths (or two paths resolving to the same CPython)
-    "${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
+    # don't register conflicting build rules.
-    "${deepgemm_SOURCE_DIR}/third-party/fmt/include")
+    if(_dg_soabi IN_LIST _dg_seen_soabis)
-
+      continue()
-  target_compile_options(_deep_gemm_C PRIVATE
+    endif()
-    $<$<COMPILE_LANGUAGE:CXX>:-O3>
+    list(APPEND _dg_seen_soabis "${_dg_soabi}")
-    $<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
+    set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
-    $<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
+    set(_dg_marker "${_dg_dir}/.built")
-
+    add_custom_command(
-  # torch_python is required because DeepGEMM uses pybind11 type casters
+      OUTPUT "${_dg_marker}"
-  # for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
+      COMMAND "${Python_EXECUTABLE}"
-  # use torch::Library custom ops.
+              "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
-  find_library(TORCH_PYTHON_LIBRARY torch_python
+              "${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
-    PATHS "${TORCH_INSTALL_PREFIX}/lib"
+      COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
-    REQUIRED)
+      DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
-
+              "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
-  target_link_libraries(_deep_gemm_C PRIVATE
+              ${_dg_headers}
-    torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
+      COMMENT "Building DeepGEMM _C for ${_pybin}"
-    CUDA::cudart CUDA::nvrtc)
+      VERBATIM)
-
+    list(APPEND _dg_markers "${_dg_marker}")
-  # Install the shared library into the vendored package directory
+    install(DIRECTORY "${_dg_dir}/"
-  install(TARGETS _deep_gemm_C
+      DESTINATION vllm/third_party/deep_gemm
-    LIBRARY DESTINATION vllm/third_party/deep_gemm
+      COMPONENT _deep_gemm_C
-    COMPONENT _deep_gemm_C)
+      FILES_MATCHING PATTERN "_C.cpython-*.so")
+  endforeach()
+  add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})
 
   #
   # Vendor DeepGEMM Python package files

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -156,6 +156,17 @@ inline int GetGroupsPerBlock(int64_t num_groups) {
   return 1;
 }
 
+// Largest divisor of padded_groups_per_row that is <= 16. ry = 16 / kx.
+inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) {
+  if (padded_groups_per_row % 16 == 0) {
+    return 16;
+  }
+  if (padded_groups_per_row % 8 == 0) {
+    return 8;
+  }
+  return 4;
+}
+
 void per_token_group_quant_8bit(const torch::stable::Tensor& input,
                                 torch::stable::Tensor& output_q,
                                 torch::stable::Tensor& output_s,
@@ -247,11 +258,11 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
 //
 // Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
 // THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
-template <typename T, typename DST_DTYPE, int GROUP_SIZE>
+template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
+          int kRowsPerBlock>
 __global__ void per_token_group_quant_8bit_packed_register_kernel(
     const T* __restrict__ input, void* __restrict__ output_q,
-    unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
+    unsigned int* __restrict__ output_s_packed, const int padded_groups_per_row,
-    const int groups_per_block, const int padded_groups_per_row,
     const int groups_per_row, const int mn, const int output_q_mn_extent,
     const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
     const float min_8bit, const float max_8bit) {
@@ -260,27 +271,25 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
   constexpr int VEC_SIZE = 32 / sizeof(T);  // 16 for bf16/fp16
   static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
                 "GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
-  // Each group's 8 threads must live in a single warp octet so the
-  // 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
-  // that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
-  // to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
-  // via num_threads = groups_per_block * THREADS_PER_GROUP).
   static_assert(32 % THREADS_PER_GROUP == 0,
                 "THREADS_PER_GROUP must divide warp size for the shuffle "
                 "mask to be valid");
+  static_assert(
+      kGroupsPerBlockX > 0 && (kGroupsPerBlockX & (kGroupsPerBlockX - 1)) == 0,
+      "kGroupsPerBlockX must be a positive power of 2");
+  static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
 
   const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
   const int lane_id = threadIdx.x % THREADS_PER_GROUP;
 
-  const int64_t block_group_id = blockIdx.x * groups_per_block;
+  const int sf_k_local = local_group_id % kGroupsPerBlockX;
-  const int64_t global_group_id = block_group_id + local_group_id;
+  const int row_local = local_group_id / kGroupsPerBlockX;
-  if (global_group_id >= num_groups_padded) {
+  const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
+
+  if (mn_idx >= tma_aligned_mn) {
     return;
   }
-
-  const int sf_k_idx =
-      static_cast<int>(global_group_id % padded_groups_per_row);
-  const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
   const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
 
   // Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -443,34 +452,53 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 
   constexpr int THREADS_PER_GROUP = 8;
   const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
-  const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
   const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
-  const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
+
+  STD_TORCH_CHECK(padded_groups_per_row % 4 == 0,
+                  "padded_groups_per_row=", padded_groups_per_row,
+                  " is not a multiple of 4.");
+  const int kx = GetGroupsPerBlockX(padded_groups_per_row);
+  const int ry = 16 / kx;
+  const int64_t blocks_x = padded_groups_per_row / kx;
+  const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
+  const int num_threads = (kx * ry) * THREADS_PER_GROUP;
+  // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
+  STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
+                      blocks_y <= static_cast<int64_t>(INT32_MAX),
+                  "per_token_group_quant_8bit_packed grid too large: (",
+                  blocks_x, ", ", blocks_y, ").");
 
   auto dst_type = output_q.scalar_type();
-  const int64_t num_blocks = num_groups_padded / groups_per_block;
-  const int num_threads = groups_per_block * THREADS_PER_GROUP;
-  // CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard
-  // against pathological inputs.
-  STD_TORCH_CHECK(num_blocks <= static_cast<int64_t>(INT32_MAX),
-                  "per_token_group_quant_8bit_packed grid too large: ",
-                  num_blocks, " blocks (max ", INT32_MAX, ").");
 
-#define LAUNCH_REG_KERNEL(T, DST_DTYPE)                                   \
+#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
-  do {                                                                    \
+  do {                                                                       \
-    dim3 grid(static_cast<unsigned int>(num_blocks));                     \
+    dim3 grid(static_cast<unsigned int>(blocks_x),                           \
-    dim3 block(num_threads);                                              \
+              static_cast<unsigned int>(blocks_y));                          \
-    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128>  \
+    dim3 block(num_threads);                                                 \
-        <<<grid, block, 0, stream>>>(                                     \
+    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
-            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
+                                                      RY>                    \
-            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),  \
+        <<<grid, block, 0, stream>>>(                                        \
-            num_groups_padded, groups_per_block,                          \
+            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),    \
-            static_cast<int>(padded_groups_per_row),                      \
+            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),     \
-            static_cast<int>(groups_per_row), static_cast<int>(mn),       \
+            static_cast<int>(padded_groups_per_row),                         \
-            static_cast<int>(output_q_mn_extent),                         \
+            static_cast<int>(groups_per_row), static_cast<int>(mn),          \
-            static_cast<int>(tma_aligned_mn), num_scale_elems,            \
+            static_cast<int>(output_q_mn_extent),                            \
-            static_cast<float>(eps), static_cast<float>(min_8bit),        \
+            static_cast<int>(tma_aligned_mn), num_scale_elems,               \
-            static_cast<float>(max_8bit));                                \
+            static_cast<float>(eps), static_cast<float>(min_8bit),           \
+            static_cast<float>(max_8bit));                                   \
+  } while (0)
+
+#define LAUNCH_REG_KERNEL(T, DST_DTYPE)                    \
+  do {                                                     \
+    if (kx == 16) {                                        \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1);         \
+    } else if (kx == 8) {                                  \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2);          \
+    } else if (kx == 4) {                                  \
+      LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4);          \
+    } else {                                               \
+      STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
+    }                                                      \
   } while (0)
 
   VLLM_STABLE_DISPATCH_HALF_TYPES(
@@ -488,6 +516,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
       }));
 
 #undef LAUNCH_REG_KERNEL
+#undef LAUNCH_REG_KERNEL_INST
 }
 
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,

docker/Dockerfile

@@ -199,7 +199,10 @@ COPY requirements/cuda.txt requirements/cuda.txt
 COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/root/.cache/uv \
-    if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
+    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
+    fi \
+    && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
         && uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
         --index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
@@ -301,6 +304,15 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         python3 use_existing_torch.py --prefix; \
     fi
 
+# Provision a bare interpreter for each CPython covered by `requires-python`
+# so DeepGEMM `_C` is built once per Python and bundled side-by-side in the
+# wheel; cmake reads DEEPGEMM_PYTHON_INTERPRETERS in deepgemm.cmake's
+# foreach loop. The matrix is derived from pyproject.toml.
+COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
+ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
+RUN --mount=type=cache,target=/root/.cache/uv \
+    tools/setup_deepgemm_pythons.sh > /tmp/dg_pythons.txt
+
 # Build the vLLM wheel
 # if USE_SCCACHE is set, use sccache to speed up compilation
 # AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
@@ -328,6 +340,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         && export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
         && export VLLM_MAIN_CUDA_VERSION="${VLLM_MAIN_CUDA_VERSION}" \
         && export VLLM_DOCKER_BUILD_CONTEXT=1 \
+        && export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) \
         && sccache --show-stats \
         && python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
         && sccache --show-stats; \
@@ -345,6 +358,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
         export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
         export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" && \
         export VLLM_DOCKER_BUILD_CONTEXT=1 && \
+        export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) && \
         python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
     fi
 
@@ -616,6 +630,9 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
 COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
 RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
+    fi && \
     uv pip install --system -r /tmp/requirements-cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
     rm /tmp/requirements-cuda.txt /tmp/common.txt
@@ -829,9 +846,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
                 libcublas-dev-${CUDA_VERSION_DASH} \
                 libcusolver-dev-${CUDA_VERSION_DASH}"; \
     if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
-        if [ "$CUDA_MAJOR" -ge 13 ]; then \
-            uv pip install --system nixl-cu13; \
-        fi; \
         uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
             # if the above fails, install from source
             apt-get update -y && \
@@ -841,6 +855,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
             # clean up -dev packages, keep runtime libraries
             rm -rf /var/lib/apt/lists/* \
         ); \
+        # Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is installed.
+        uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
     fi
 
 ENV VLLM_USAGE_SOURCE production-docker-image

docs/design/attention_backends.md

@@ -125,12 +125,13 @@ Priority is **1 = highest** (tried first).
 | Priority | Backend |
 | -------- | ------- |
 | 1 | `FLASHINFER_MLA` |
-| 2 | `CUTLASS_MLA` |
+| 2 | `TOKENSPEED_MLA` |
-| 3 | `FLASH_ATTN_MLA` |
+| 3 | `CUTLASS_MLA` |
-| 4 | `FLASHMLA` |
+| 4 | `FLASH_ATTN_MLA` |
-| 5 | `TRITON_MLA` |
+| 5 | `FLASHMLA` |
-| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
+| 6 | `TRITON_MLA` |
-| 7 | `FLASHMLA_SPARSE` |
+| 7 | `FLASHINFER_MLA_SPARSE`**\*** |
+| 8 | `FLASHMLA_SPARSE` |
 
 **Ampere/Hopper (SM 8.x-9.x):**
 
@@ -202,6 +203,7 @@ hardware and configuration.
 | `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
 | `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
+| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 
 > **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
 > On other GPUs, FlashAttention is used as the default.
@@ -222,5 +224,6 @@ MLA decode backends are selected using the standard
 | `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |

docs/features/mooncake_store_connector_usage.md

@@ -0,0 +1,161 @@
+# MooncakeStoreConnector Usage Guide
+
+MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:
+
+- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
+- **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
+- **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
+
+## Prerequisites
+
+### Install Mooncake
+
+Install mooncake through pip:
+
+```bash
+uv pip install mooncake-transfer-engine
+```
+
+Refer to the [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source.
+
+### Start the Mooncake Master Server
+
+The Mooncake master manages metadata and coordinates the distributed store. Start it before launching vLLM:
+
+```bash
+mooncake_master --port 50051
+```
+
+Default ports:
+
+- RPC: 50051
+
+Multiple vLLM instances can share the same master server.
+
+### Configure Mooncake
+
+Create a JSON configuration file (e.g., `mooncake_config.json`):
+
+```json
+{
+  "metadata_server": "P2PHANDSHAKE",
+  "master_server_address": "127.0.0.1:50051",
+  "global_segment_size": "80GB",
+  "local_buffer_size": "4GB",
+  "protocol": "rdma",
+  "device_name": ""
+}
+```
+
+- `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
+- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
+- `local_buffer_size`: Private buffer for this node's own operations (per GPU).
+
+Set the config path via environment variable:
+
+```bash
+export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
+```
+
+## Usage
+
+### Single-Node KV Cache Offloading
+
+Use MooncakeStoreConnector to offload KV cache to CPU memory, extending the effective cache size:
+
+```bash
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
+```
+
+### Disaggregated Prefill-Decode (XpYd)
+
+In disaggregated prefill-decode mode, use `MultiConnector` to combine `MooncakeConnector` (point-to-point KV transfer) with `MooncakeStoreConnector` (shared KV cache pool). This enables both direct P2P transfer between prefiller and decoder, and cross-instance prefix cache sharing via the distributed store.
+**Prefiller Node:**
+
+```bash
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8100 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_producer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_producer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_producer"
+                }
+            ]
+        }
+    }'
+```
+
+**Decoder Node:**
+
+```bash
+MOONCAKE_CONFIG_PATH=mooncake_config.json \
+VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
+vllm serve meta-llama/Llama-3.1-8B-Instruct \
+    --port 8200 \
+    --kv-transfer-config '{
+        "kv_connector": "MultiConnector",
+        "kv_role": "kv_consumer",
+        "kv_connector_extra_config": {
+            "connectors": [
+                {
+                    "kv_connector": "MooncakeConnector",
+                    "kv_role": "kv_consumer"
+                },
+                {
+                    "kv_connector": "MooncakeStoreConnector",
+                    "kv_role": "kv_consumer"
+                }
+            ]
+        }
+    }'
+```
+
+**Proxy:**
+
+A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.
+
+## Environment Variables
+
+| Variable | Description | Default |
+| --- | --- | --- |
+| `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
+| `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
+
+## KV Transfer Config
+
+### KV Role Options
+
+- **kv_producer**: For prefiller instances that store KV caches to the pool.
+- **kv_consumer**: For decoder instances that load KV caches from the pool.
+- **kv_both**: The instance both stores and loads KV caches. Use this for single-node CPU offloading.
+
+### kv_connector_extra_config
+
+- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
+- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
+- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
+- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
+
+## Notes
+
+### Cross-DP Prefix Cache Hits
+
+When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
+
+```bash
+PYTHONHASHSEED=0 vllm serve ...
+```
+
+Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.

requirements/cuda.txt

@@ -21,5 +21,8 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
 fastsafetensors >= 0.2.2
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
-nvidia-cutlass-dsl>=4.4.2
+nvidia-cutlass-dsl==4.4.2
 quack-kernels>=0.3.3
+
+# Tokenspeed_MLA for faster mla with spec decode
+tokenspeed-mla==0.1.2

requirements/kv_connectors.txt

@@ -1,5 +1,3 @@
 lmcache >= 0.3.9
-nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
+nixl >= 1.1.0 # Required for disaggregated prefill
-nixl-cu12 >= 0.7.1, <= 0.10.1
-nixl-cu13 >= 0.7.1, <= 0.10.1
 mooncake-transfer-engine >= 0.3.8

setup.py

@@ -969,6 +969,9 @@ def get_requirements() -> list[str]:
                 # vllm-flash-attn is built only for CUDA 12.x.
                 # Skip for other versions.
                 continue
+            if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
+                # [cu13] extra is the default; strip it on CUDA 12 builds.
+                req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
             modified_requirements.append(req)
         requirements = modified_requirements
     elif _is_hip():

tests/conftest.py

@@ -224,19 +224,28 @@ def init_test_http_connection():
 def dist_init():
     from tests.utils import ensure_current_vllm_config
 
-    temp_file = tempfile.mkstemp()[1]
+    # Close the fd returned by mkstemp; FileStore opens the path itself.
+    # Leaving it open leaks one FD per test and eventually exhausts the
+    # ulimit, causing FileStore's destructor to throw c10::DistStoreError
+    # ("Too many open files") during gc and abort the process.
+    fd, temp_file = tempfile.mkstemp()
+    os.close(fd)
 
-    with ensure_current_vllm_config():
+    try:
-        init_distributed_environment(
+        with ensure_current_vllm_config():
-            world_size=1,
+            init_distributed_environment(
-            rank=0,
+                world_size=1,
-            distributed_init_method=f"file://{temp_file}",
+                rank=0,
-            local_rank=0,
+                distributed_init_method=f"file://{temp_file}",
-            backend="nccl",
+                local_rank=0,
-        )
+                backend="nccl",
-        initialize_model_parallel(1, 1)
+            )
-        yield
+            initialize_model_parallel(1, 1)
-    cleanup_dist_env_and_memory()
+            yield
+        cleanup_dist_env_and_memory()
+    finally:
+        with contextlib.suppress(OSError):
+            os.unlink(temp_file)
 
 
 @pytest.fixture

tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py

@@ -66,6 +66,24 @@ def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Ten
     return reference_quant(ref_act_out, use_ue8m0)
 
 
+def reference_with_clamp(
+    x: torch.Tensor, use_ue8m0: bool, clamp_limit: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to
+    match the C++ compute() template, then run the standard silu_and_mul +
+    quant reference."""
+    N_2 = x.size(1) // 2
+    dtype = x.dtype
+    gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype)
+    up = (
+        x[..., N_2:]
+        .to(torch.float32)
+        .clamp(min=-clamp_limit, max=clamp_limit)
+        .to(dtype)
+    )
+    return reference(torch.cat([gate, up], dim=-1), use_ue8m0)
+
+
 @pytest.mark.parametrize("T", [128, 256, 512])
 @pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
 @pytest.mark.skipif(
@@ -89,3 +107,32 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):
 
     torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
     torch.testing.assert_close(output_scales, ref_output_scales)
+
+
+@pytest.mark.parametrize("T", [128, 256, 512])
+@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
+@pytest.mark.parametrize("clamp_limit", [7.0, 10.0])
+@pytest.mark.skipif(
+    current_platform.is_rocm(),
+    reason="ROCm does not support DeepGemm.",
+)
+def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float):
+    set_random_seed(42)
+
+    # Use a wide distribution so values routinely exceed both clamp limits and
+    # the clamp branch is actually exercised (uniform [0, 1) inputs would never
+    # trigger it).
+    input = torch.randn((T, N), dtype=torch.bfloat16, device="cuda") * 8.0
+
+    use_ue8m0 = is_deep_gemm_e8m0_used()
+
+    # Test
+    output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor(
+        input, use_ue8m0=use_ue8m0, clamp_limit=clamp_limit
+    )
+
+    # Reference
+    ref_output, ref_output_scales = reference_with_clamp(input, use_ue8m0, clamp_limit)
+
+    torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
+    torch.testing.assert_close(output_scales, ref_output_scales)

tests/v1/attention/test_mla_backends.py

@@ -20,17 +20,20 @@ from tests.v1.attention.utils import (
 from vllm import _custom_ops as ops
 from vllm.config.vllm import set_current_vllm_config
 from vllm.model_executor.layers.attention.mla_attention import (
+    MLAAttention,
     QueryLenSupport,
     _DecodeConcatQuantFP8,
 )
-from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
-from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+from vllm.v1.attention.backends.mla.prefill import (
+    MLAPrefillBackendEnum,
+    get_mla_prefill_backend,
+)
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
@@ -41,6 +44,7 @@ BACKENDS_TO_TEST = [
     AttentionBackendEnum.FLASH_ATTN_MLA,
     AttentionBackendEnum.FLASHINFER_MLA,
     AttentionBackendEnum.TRITON_MLA,
+    AttentionBackendEnum.TOKENSPEED_MLA,
 ]
 
 DEVICE_TYPE = current_platform.device_type
@@ -49,6 +53,7 @@ DEVICE_TYPE = current_platform.device_type
 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
     BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
     BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER_MLA)
+    BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
 
 # Remove FLASH_ATTN_MLA from the list if not supported
 if not flash_attn_supports_mla():
@@ -58,6 +63,22 @@ if not flash_attn_supports_mla():
 if not is_flashmla_dense_supported()[0]:
     BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)
 
+# Remove TOKENSPEED_MLA if the optional package is not installed
+if AttentionBackendEnum.TOKENSPEED_MLA in BACKENDS_TO_TEST:
+    try:
+        import tokenspeed_mla  # noqa: F401
+    except ImportError:
+        BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
+
+
+# Filtered per-test via validate_configuration (capability/deps/dims).
+PREFILL_BACKENDS_TO_TEST = [
+    MLAPrefillBackendEnum.FLASH_ATTN,
+    MLAPrefillBackendEnum.FLASHINFER,
+    MLAPrefillBackendEnum.TRTLLM_RAGGED,
+    MLAPrefillBackendEnum.TOKENSPEED_MLA,
+]
+
 
 SPEC_DECODE_BACKENDS = []
 for backend in BACKENDS_TO_TEST:
@@ -389,14 +410,18 @@ class MockSparseMLAAttentionLayer:
         return output
 
 
-class MockMLAAttentionLayer(AttentionLayerBase):
+class MockMLAAttentionLayer(MLAAttention):
     """A mock MLA attention layer for testing.
 
     This replicates the forward_impl logic from MLAAttention to allow
     testing MLA backends without the full layer infrastructure.
 
-    The W_UK_T and W_UV weight matrices are created on the layer (like in
+    Subclasses MLAAttention so that backends that filter
-    MLAAttention.process_weights_after_loading), not on the impl.
+    `static_forward_context` by `isinstance(layer, MLAAttention)` (e.g.
+    FlashInfer prefill, which reads sm_scale through that filter) see the
+    mock as a real MLA layer. MLAAttention.__init__ is intentionally
+    skipped — it would create its own impl/prefill_backend and self-register
+    in static_forward_context, which fights what the test sets up below.
     """
 
     def __init__(
@@ -412,6 +437,7 @@ class MockMLAAttentionLayer(AttentionLayerBase):
         q_scale: float,
         k_scale: float,
     ):
+        torch.nn.Module.__init__(self)
         self.impl = impl
         self.num_heads = num_heads
         self.qk_nope_head_dim = qk_nope_head_dim
@@ -562,11 +588,15 @@ def run_attention_backend(
     q_scale: float,
     k_scale: float,
     kv_cache_dtype: str = "auto",
+    prefill_backend: MLAPrefillBackendEnum | None = None,
 ) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
 
     builder_cls, impl_cls = try_get_attention_backend(backend)
 
+    # Force the prefill backend selection (None means auto-select).
+    vllm_config.attention_config.mla_prefill_backend = prefill_backend
+
     # Set the current vllm config so that get_current_vllm_config() works
     # in the backend implementations
     with set_current_vllm_config(vllm_config):
@@ -578,7 +608,11 @@ def run_attention_backend(
             vllm_config.parallel_config
         )
         head_size = vllm_config.model_config.get_head_size()
-        scale = 1.0 / (head_size**0.5)
+        # Production MLA passes 1/sqrt(qk_head_dim) (the prefill scale) to the
+        # impl and forwards the same value to the prefill backend. FLASHINFER
+        # prefill reads sm_scale back from impl.scale via global_hyperparameters
+        # at plan() time, so impl.scale must agree with prefill_backend.scale.
+        scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
         impl = impl_cls(
             num_heads=num_heads,
             head_size=head_size,
@@ -683,6 +717,7 @@ def run_attention_backend(
 @pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
 @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
+@pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST)
 def test_backend_correctness(
     default_vllm_config,
     dist_init,
@@ -693,6 +728,7 @@ def test_backend_correctness(
     kv_cache_dtype: str,
     q_scale: float,
     k_scale: float,
+    prefill_backend: MLAPrefillBackendEnum,
 ):
     """
     Test that all backends produce similar outputs to a reference implementation
@@ -729,6 +765,24 @@ def test_backend_correctness(
     if not backends_to_test:
         pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")
 
+    # Skip prefill backends that can't satisfy capability/deps/R1 constraints.
+    from vllm.v1.attention.backends.mla.prefill.selector import (
+        MLAPrefillSelectorConfig,
+    )
+
+    try:
+        prefill_invalid_reasons = prefill_backend.get_class().validate_configuration(
+            current_platform.get_device_capability(),
+            MLAPrefillSelectorConfig(dtype=torch.bfloat16, is_r1_compatible=True),
+        )
+    except ImportError:
+        prefill_invalid_reasons = ["ImportError"]
+    if prefill_invalid_reasons:
+        pytest.skip(
+            f"Prefill backend {prefill_backend.name} unavailable: "
+            f"{prefill_invalid_reasons}"
+        )
+
     batch_spec = BATCH_SPECS[batch_spec_name]
     is_spec_decode_test = batch_spec_name.startswith("spec_decode")
     unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test))
@@ -799,9 +853,13 @@ def test_backend_correctness(
     assert kv_lora_rank + qk_rope_head_dim == head_size, (
         f"MLA dimensions don't match: {total_head_size} != {head_size}"
     )
-    decode_scale = 1.0 / (total_head_size**0.5)
     qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
     prefill_scale = qk_head_dim**-0.5
+    # MLA reuses prefill_scale for the decode path: production sets
+    # impl.scale = 1/sqrt(qk_head_dim) and the decode kernels apply it even
+    # though the latent attention runs at head_size dimensions. Keeping the
+    # reference here in sync with run_attention_backend's impl.scale.
+    decode_scale = prefill_scale
 
     # 2. Generate data and compute SDPA reference output for MLA
     all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
@@ -1092,6 +1150,7 @@ def test_backend_correctness(
             qk_rope_head_dim,
             v_head_dim,
             mock_kv_b_proj,
+            prefill_backend=prefill_backend,
             q_scale=q_scale,
             k_scale=k_scale,
             kv_cache_dtype=kv_cache_dtype,

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock, patch
+
+from vllm.config import set_current_vllm_config
+from vllm.distributed.kv_events import BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorRole,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
+    connector,
+    worker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    MooncakeStoreConnectorMetadata,
+)
+from vllm.v1.outputs import KVConnectorOutput
+
+from .utils import create_vllm_config
+
+
+def _make_vllm_config():
+    return create_vllm_config(
+        kv_connector="MooncakeStoreConnector",
+        kv_role="kv_both",
+    )
+
+
+def _make_block_stored() -> BlockStored:
+    return BlockStored(
+        block_hashes=[b"hash"],
+        parent_block_hash=None,
+        token_ids=[1, 2, 3],
+        block_size=16,
+        lora_id=None,
+        medium="cpu",
+        lora_name=None,
+    )
+
+
+def test_scheduler_role_initializes_store_scheduler_only():
+    vllm_config = _make_vllm_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ) as mock_scheduler,
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker,
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
+
+    mock_scheduler.assert_called_once_with(vllm_config)
+    mock_worker.assert_not_called()
+    assert conn.connector_scheduler is mock_scheduler.return_value
+    assert conn.connector_worker is None
+
+
+def test_worker_role_initializes_store_worker_on_rank0():
+    vllm_config = _make_vllm_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ) as mock_scheduler,
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker,
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
+
+    mock_scheduler.assert_not_called()
+    mock_worker.assert_called_once_with(vllm_config)
+    assert conn.connector_scheduler is None
+    assert conn.connector_worker is mock_worker.return_value
+
+
+def test_worker_role_initializes_on_nonzero_rank():
+    vllm_config = _make_vllm_config()
+    vllm_config.parallel_config.rank = 1
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker,
+    ):
+        connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
+
+    mock_worker.assert_called_once_with(vllm_config)
+
+
+def test_lookup_rpc_path_uses_data_parallel_index_in_dense_dp():
+    vllm_config = _make_vllm_config()
+    vllm_config.parallel_config.data_parallel_rank = 0
+    vllm_config.parallel_config.data_parallel_index = 3
+
+    path = worker.get_zmq_rpc_path_lookup(vllm_config)
+
+    assert path.endswith("_dp_rank3")
+
+
+def test_lookup_rpc_path_uses_local_rank_when_local_engines_only():
+    vllm_config = _make_vllm_config()
+    vllm_config.parallel_config.data_parallel_index = 7
+    vllm_config.parallel_config.data_parallel_rank_local = 1
+    vllm_config.parallel_config.data_parallel_hybrid_lb = True
+
+    path = worker.get_zmq_rpc_path_lookup(vllm_config)
+
+    assert path.endswith("_dp_rank1")
+
+
+def test_worker_methods_delegate_to_store_worker():
+    vllm_config = _make_vllm_config()
+    kv_caches = {"layer0": MagicMock()}
+    metadata = MooncakeStoreConnectorMetadata(set(), set())
+    finished_req_ids = {"req-1"}
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker_cls,
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
+
+    worker_inst = mock_worker_cls.return_value
+    worker_inst.get_finished.return_value = ({"req-1"}, {"req-2"})
+    conn.bind_connector_metadata(metadata)
+
+    conn.register_kv_caches(kv_caches)
+    result = conn.get_finished(finished_req_ids)
+
+    worker_inst.register_kv_caches.assert_called_once_with(kv_caches)
+    worker_inst.get_finished.assert_called_once_with(finished_req_ids, metadata)
+    assert result == ({"req-1"}, {"req-2"})
+
+
+def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
+    vllm_config = _make_vllm_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker_cls,
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
+
+    mock_worker_cls.return_value.get_kv_events.return_value = []
+    assert conn.get_kv_connector_kv_cache_events() is None
+
+
+def test_get_kv_connector_kv_cache_events_wraps_worker_events():
+    vllm_config = _make_vllm_config()
+    event = _make_block_stored()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker_cls,
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
+
+    mock_worker_cls.return_value.get_kv_events.return_value = [event]
+    kv_events = conn.get_kv_connector_kv_cache_events()
+
+    assert isinstance(kv_events, connector.MooncakeStoreKVEvents)
+    assert kv_events.get_number_of_workers() == 1
+    assert kv_events.get_all_events() == [event]
+
+
+def test_prefer_cross_layer_blocks_from_config():
+    # Default: disabled
+    vllm_config = _make_vllm_config()
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ),
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
+    assert conn.prefer_cross_layer_blocks is False
+
+    # Enabled via config
+    vllm_config_enabled = create_vllm_config(
+        kv_connector="MooncakeStoreConnector",
+        kv_role="kv_both",
+        kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
+    )
+    with (
+        set_current_vllm_config(vllm_config_enabled),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ),
+    ):
+        conn_enabled = connector.MooncakeStoreConnector(
+            vllm_config_enabled, KVConnectorRole.SCHEDULER
+        )
+    assert conn_enabled.prefer_cross_layer_blocks is True
+
+
+def test_register_cross_layers_kv_cache_delegates_to_worker():
+    vllm_config = _make_vllm_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ) as mock_worker_cls,
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
+
+    fake_tensor = MagicMock()
+    fake_backend = MagicMock()
+    conn.register_cross_layers_kv_cache(fake_tensor, fake_backend)
+
+    worker_inst = mock_worker_cls.return_value
+    worker_inst.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
+
+
+def test_update_connector_output_and_take_events():
+    vllm_config = _make_vllm_config()
+    event = _make_block_stored()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ),
+    ):
+        conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
+
+    kv_events = connector.MooncakeStoreKVEvents(num_workers=1)
+    kv_events.add_events([event])
+    conn.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))
+
+    assert conn._kv_cache_events is kv_events
+    assert list(conn.take_events()) == [event]
+    assert conn._kv_cache_events is None

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
+    worker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    ReqMeta,
+)
+
+
+def _make_store_sending_thread(
+    store: MagicMock,
+) -> worker.KVCacheStoreSendingThread:
+    token_database = ChunkedTokenDatabase(
+        KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
+    )
+    token_database.set_kv_caches_base_addr([0x1000])
+    token_database.set_block_len([256])
+    thread = worker.KVCacheStoreSendingThread(
+        store=store,
+        token_database=token_database,
+        block_size=16,
+        tp_rank=0,
+        put_step=1,
+        kv_role="kv_producer",
+        ready_event=threading.Event(),
+    )
+    thread.request_queue.task_done = MagicMock()
+    return thread
+
+
+def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
+    return ReqMeta(
+        req_id=req_id,
+        token_len_chunk=32,
+        block_ids=[0, 1],
+        block_hashes=block_hashes,
+        can_save=True,
+        original_block_size=16,
+    )
+
+
+def test_store_sending_thread_skips_request_during_cpu_pressure():
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.side_effect = [
+        [-200, -200],
+        [256, 256],
+        [256, 256],
+    ]
+    thread = _make_store_sending_thread(store)
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
+
+    assert thread._store_pressure_active is True
+    assert "req-a" in thread._skip_store_requests
+    assert store.batch_put_from_multi_buffers.call_count == 1
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 1
+
+    thread.add_stored_request("req-b")
+    thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"]))
+
+    assert thread._store_pressure_active is False
+    assert "req-a" not in thread._skip_store_requests
+    assert store.batch_put_from_multi_buffers.call_count == 2
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 3
+
+
+def test_store_sending_thread_only_skips_on_no_available_handle():
+    store = MagicMock()
+    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
+    store.batch_put_from_multi_buffers.side_effect = [
+        [-500, -500],
+        [256, 256],
+    ]
+    thread = _make_store_sending_thread(store)
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
+
+    assert thread._store_pressure_active is False
+    assert "req-a" not in thread._skip_store_requests
+    assert store.batch_put_from_multi_buffers.call_count == 1
+
+    thread.add_stored_request("req-a")
+    thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
+
+    assert store.batch_put_from_multi_buffers.call_count == 2
+
+
+# ---------------------------------------------------------------------------
+# Helpers for register_kv_caches tests
+# ---------------------------------------------------------------------------
+
+
+def _auto_set_ready_event(*args, **kwargs):
+    """Side effect for mocked thread constructors that auto-sets ready_event."""
+    for arg in args:
+        if isinstance(arg, threading.Event):
+            arg.set()
+    for val in kwargs.values():
+        if isinstance(val, threading.Event):
+            val.set()
+    return MagicMock()
+
+
+def _make_bare_worker(
+    *,
+    num_gpu_blocks: int = 10,
+    block_size: int = 16,
+    kv_role: str = "kv_both",
+) -> worker.MooncakeStoreWorker:
+    """Construct a MooncakeStoreWorker via __new__, bypassing __init__.
+
+    Sets only the attributes that register_kv_caches() reads so we can
+    test the stride-based layout detection without a real
+    MooncakeDistributedStore.
+    """
+    w = object.__new__(worker.MooncakeStoreWorker)
+    w.cache_config = MagicMock()
+    w.cache_config.num_gpu_blocks = num_gpu_blocks
+    w.store = MagicMock()
+    w.store.register_buffer.return_value = 0
+    w.use_mla = False
+    w.token_database = ChunkedTokenDatabase(
+        KeyMetadata("test-model", 0, 0, 0, 0), block_size=block_size
+    )
+    w.kv_role = kv_role
+    w.block_size = block_size
+    w.tp_rank = 0
+    w.put_step = 1
+    w.enable_kv_events = False
+    w.kv_send_thread = None
+    w.kv_recv_thread = None
+    return w
+
+
+# ---------------------------------------------------------------------------
+# register_kv_caches tests
+# ---------------------------------------------------------------------------
+
+
+def test_register_kv_caches_blocks_first_single_segment():
+    """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
+    num_blocks = 10
+    page_size_elements = 64  # elements per block
+    w = _make_bare_worker(num_gpu_blocks=num_blocks)
+
+    # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
+    tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
+
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreSendingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreRecvingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+    ):
+        w.register_kv_caches({"layer0": tensor})
+
+    assert len(w.kv_caches_base_addr) == 1
+    assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
+
+    expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
+    assert len(w.block_len) == 1
+    assert w.block_len[0] == expected_block_len
+
+    w.store.register_buffer.assert_called_once_with(
+        tensor.untyped_storage().data_ptr(),
+        tensor.untyped_storage().nbytes(),
+    )
+
+
+def test_register_kv_caches_kv_first_two_segments():
+    """K/V-first layout (FlashAttn): two segments (K, V) per layer."""
+    num_blocks = 10
+    block_size_tokens = 16
+    num_kv_heads = 4
+    head_size = 8
+
+    w = _make_bare_worker(num_gpu_blocks=num_blocks)
+
+    # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
+    tensor = torch.zeros(
+        2,
+        num_blocks,
+        block_size_tokens,
+        num_kv_heads,
+        head_size,
+        dtype=torch.float16,
+    )
+
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreSendingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreRecvingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+    ):
+        w.register_kv_caches({"layer0": tensor})
+
+    # K/V-first: dim 0 has stride > page_size, so 2 segments
+    assert len(w.kv_caches_base_addr) == 2
+    assert len(w.block_len) == 2
+
+    el = tensor.element_size()
+    seg_stride = tensor.stride(0) * el  # stride of the K/V dim in bytes
+    base = tensor.untyped_storage().data_ptr()
+    assert w.kv_caches_base_addr[0] == base
+    assert w.kv_caches_base_addr[1] == base + seg_stride
+    assert w.block_len[0] == seg_stride // num_blocks
+    assert w.block_len[1] == seg_stride // num_blocks
+
+
+def test_register_kv_caches_cross_layer_single_segment():
+    """Cross-layer tensor: single segment with block_len = page_size * num_layers."""
+    num_blocks = 10
+    num_layers = 4
+    per_layer_page_elements = 64  # elements per layer per block
+
+    w = _make_bare_worker(num_gpu_blocks=num_blocks)
+
+    # Cross-layer blocks-first tensor: all layers packed into a single
+    # contiguous block.  Shape (num_blocks, num_layers * per_layer_page)
+    # mimics the physical layout after stride reordering.
+    total_page_elements = num_layers * per_layer_page_elements
+    tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
+
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreSendingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreRecvingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+    ):
+        # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
+        w.register_kv_caches({"__cross_layer__": tensor})
+
+    assert len(w.kv_caches_base_addr) == 1
+    assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
+
+    expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
+    # block_len should be per_layer_page_size * num_layers
+    assert (
+        expected_block_len
+        == num_layers * per_layer_page_elements * tensor.element_size()
+    )
+    assert len(w.block_len) == 1
+    assert w.block_len[0] == expected_block_len
+
+    # Also verify via register_cross_layers_kv_caches wrapper
+    w2 = _make_bare_worker(num_gpu_blocks=num_blocks)
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreSendingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "worker.KVCacheStoreRecvingThread",
+            side_effect=_auto_set_ready_event,
+        ),
+    ):
+        w2.register_cross_layers_kv_caches(tensor)
+
+    assert w2.kv_caches_base_addr == w.kv_caches_base_addr
+    assert w2.block_len == w.block_len

tools/build_deepgemm_C.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Build DeepGEMM's `_C` pybind11 extension for a target Python.
+
+Driven from `cmake/external_projects/deepgemm.cmake`. The driver is the
+build interpreter (which has torch); the *target* Python is only used for
+its header path and SOABI. This avoids needing torch installed in N venvs
+to produce N matching `.so` files.
+
+Usage: python build_deepgemm_C.py <DEEPGEMM_SRC_DIR> <OUTPUT_DIR> <TARGET_PY>
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+from torch.utils import cpp_extension
+
+if len(sys.argv) != 4:
+    sys.exit(f"usage: {sys.argv[0]} <SRC> <OUT> <TARGET_PY>")
+
+src = Path(sys.argv[1]).resolve()
+out = Path(sys.argv[2]).resolve()
+target_py = sys.argv[3]
+out.mkdir(parents=True, exist_ok=True)
+
+info = json.loads(
+    subprocess.check_output(
+        [
+            target_py,
+            "-c",
+            "import sysconfig, json; "
+            "print(json.dumps({k: sysconfig.get_config_var(k) "
+            "for k in ('EXT_SUFFIX', 'INCLUDEPY')}))",
+        ]
+    ).decode()
+)
+
+cuda_home = cpp_extension.CUDA_HOME
+if cuda_home is None:
+    sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C")
+# CCCL lives outside the standard CUDAToolkit search, mirroring DeepGEMM's
+# own setup.py.
+includes = [
+    info["INCLUDEPY"],
+    f"{cuda_home}/include",
+    f"{cuda_home}/include/cccl",
+    str(src / "csrc"),
+    str(src / "deep_gemm/include"),
+    str(src / "third-party/cutlass/include"),
+    str(src / "third-party/cutlass/tools/util/include"),
+    str(src / "third-party/fmt/include"),
+    *cpp_extension.include_paths(device_type="cuda"),
+]
+
+cmd = [
+    os.environ.get("CXX", "g++"),
+    "-shared",
+    "-fPIC",
+    "-std=c++20",
+    "-O3",
+    "-g0",
+    "-Wno-psabi",
+    "-Wno-deprecated-declarations",
+    "-DTORCH_API_INCLUDE_EXTENSION_H",
+    "-DTORCH_EXTENSION_NAME=_C",
+    f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}",
+    *(f"-I{p}" for p in includes),
+    str(src / "csrc/python_api.cpp"),
+    *(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")),
+    f"-L{cuda_home}/lib64",
+    "-ltorch",
+    "-ltorch_python",
+    "-ltorch_cpu",
+    "-ltorch_cuda",
+    "-lc10",
+    "-lc10_cuda",
+    "-lcudart",
+    "-lnvrtc",
+    "-o",
+    str(out / f"_C{info['EXT_SUFFIX']}"),
+]
+print("[build_deepgemm_C] " + " ".join(cmd), flush=True)
+subprocess.check_call(cmd)

tools/check_wheel_deepgemm.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython
+covered by `requires-python`. Fails closed if a Python's `.so` is missing
+from the wheel — i.e. the regression that surfaced in #41476/#41512.
+
+Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm
+kernel tests in .buildkite/test_areas/kernels.yaml.
+"""
+
+import importlib.util
+import os
+import sys
+from pathlib import Path
+
+import regex as re
+import tomllib
+
+SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-")
+
+
+def required_pythons() -> list[str]:
+    pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml"
+    spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"]
+    m = re.match(r">=3\.(\d+),<3\.(\d+)", spec)
+    if not m:
+        sys.exit(f"unexpected requires-python format: {spec!r}")
+    return [f"3.{v}" for v in range(int(m[1]), int(m[2]))]
+
+
+spec = importlib.util.find_spec("vllm.third_party.deep_gemm")
+if spec is None or spec.origin is None:
+    sys.exit("vllm.third_party.deep_gemm not importable; is vllm installed?")
+pkg_dir = Path(spec.origin).parent
+
+found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))}
+required = required_pythons()
+missing = [v for v in required if v not in found]
+print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}")
+sys.exit(1 if missing else 0)

tools/setup_deepgemm_pythons.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# Provision bare Python interpreters for the DeepGEMM `_C` per-Python build
+# and print a colon-separated list of their paths to stdout.
+#
+# Each target Python only needs a working interpreter — torch is not
+# installed since `tools/build_deepgemm_C.py` runs from the build interpreter.
+# uv re-uses any matching system Python and downloads a managed build
+# otherwise.
+#
+# Usage:
+#   export DEEPGEMM_PYTHON_INTERPRETERS=$(tools/setup_deepgemm_pythons.sh)
+#   python setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
+#
+# With no args, expands to every CPython covered by `requires-python` in
+# pyproject.toml. Pass explicit versions (e.g. `3.10 3.11`) to override.
+#
+# Skip this script if you don't have uv: set DEEPGEMM_PYTHON_INTERPRETERS
+# directly to existing interpreter paths. Editable / single-Python builds
+# don't need the env var at all (cmake falls back to the build interpreter).
+#
+# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv).
+set -euo pipefail
+
+if [ "$#" -eq 0 ]; then
+  # Derive the matrix from `requires-python = ">=3.X,<3.Y"` in pyproject.toml.
+  pyproject="$(dirname "$0")/../pyproject.toml"
+  spec=$(grep -E '^requires-python' "$pyproject" \
+         | grep -oE '>=3\.[0-9]+,<3\.[0-9]+')
+  lo=${spec#>=3.}; lo=${lo%%,*}
+  hi=${spec##*<3.}
+  set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./')
+fi
+
+prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}"
+mkdir -p "$prefix"
+
+paths=""
+for V in "$@"; do
+  venv="$prefix/$V"
+  # Force a managed (uv-downloaded) Python so dev headers are bundled.
+  # System Pythons on the build base may lack headers (manylinux's
+  # /opt/python/cpXY-cpXY are off PATH; an apt-installed python3.X often
+  # has no -dev), and the per-Python build needs Python.h.
+  [ -x "$venv/bin/python" ] || \
+    uv venv --python "$V" "$venv" --python-preference only-managed --seed \
+      >/dev/null
+  paths="$paths:$venv/bin/python"
+done
+echo "${paths#:}"

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -197,6 +197,11 @@ KVConnectorFactory.register_connector(
     "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
     "MooncakeConnector",
 )
+KVConnectorFactory.register_connector(
+    "MooncakeStoreConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector",
+    "MooncakeStoreConnector",
+)
 KVConnectorFactory.register_connector(
     "FlexKVConnectorV1",
     "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py

@@ -8,6 +8,7 @@ import uvicorn
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
+from vllm.config import ParallelConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
 from vllm.logger import init_logger
 
@@ -16,6 +17,15 @@ WorkerAddr = str
 logger = init_logger(__name__)
 
 
+def get_mooncake_dp_engine_index(parallel_config: ParallelConfig) -> int:
+    """Return the per-engine DP index used for Mooncake side channels."""
+    if parallel_config.local_engines_only:
+        assert parallel_config.data_parallel_rank_local is not None
+        return parallel_config.data_parallel_rank_local
+
+    return parallel_config.data_parallel_index
+
+
 class RegisterWorkerPayload(BaseModel):
     engine_id: EngineId
     dp_rank: int

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""MooncakeStoreConnector - KV cache connector using MooncakeDistributedStore.
+
+Unlike MooncakeConnector which does direct P2P transfer, this connector
+uses MooncakeDistributedStore as a shared KV cache pool. Both producer
+and consumer instances read/write KV to/from the store independently,
+enabling prefix caching via hash-based deduplication.
+"""
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import (
+    KVCacheEvent,
+    KVConnectorKVEvents,
+    KVEventAggregator,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+)
+from vllm.forward_context import ForwardContext
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.request import Request
+
+from .data import MooncakeStoreConnectorMetadata
+from .scheduler import MooncakeStoreScheduler
+from .worker import MooncakeStoreWorker
+
+logger = init_logger(__name__)
+
+
+class MooncakeStoreKVEvents(KVConnectorKVEvents):
+    """KV event aggregation for MooncakeStoreConnector."""
+
+    def __init__(self, num_workers: int) -> None:
+        self._aggregator = KVEventAggregator(num_workers)
+
+    def add_events(self, events: list[KVCacheEvent]) -> None:
+        self._aggregator.add_events(events)
+
+    def aggregate(self) -> "MooncakeStoreKVEvents":
+        common_events = self._aggregator.get_common_events()
+        self._aggregator.clear_events()
+        self._aggregator.add_events(common_events)
+        self._aggregator.reset_workers()
+        return self
+
+    def increment_workers(self, count: int = 1) -> None:
+        self._aggregator.increment_workers(count)
+
+    def get_all_events(self) -> list[KVCacheEvent]:
+        return self._aggregator.get_all_events()
+
+    def get_number_of_workers(self) -> int:
+        return self._aggregator.get_number_of_workers()
+
+    def clear_events(self) -> None:
+        self._aggregator.clear_events()
+        self._aggregator.reset_workers()
+
+    def __repr__(self) -> str:
+        return f"<MooncakeStoreKVEvents events={self.get_all_events()}>"
+
+
+class MooncakeStoreConnector(KVConnectorBase_V1):
+    """KV connector using MooncakeDistributedStore as shared KV pool."""
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        extra_config = self._kv_transfer_config.kv_connector_extra_config
+        return (
+            str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
+            == "true"
+        )
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: KVCacheConfig | None = None,
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            role=role,
+            kv_cache_config=kv_cache_config,  # type: ignore[arg-type]
+        )
+        assert vllm_config.kv_transfer_config is not None
+        self.kv_role = vllm_config.kv_transfer_config.kv_role
+        self._kv_cache_events: MooncakeStoreKVEvents | None = None
+
+        self.connector_scheduler: MooncakeStoreScheduler | None = None
+        self.connector_worker: MooncakeStoreWorker | None = None
+
+        if role == KVConnectorRole.SCHEDULER:
+            self.connector_scheduler = MooncakeStoreScheduler(vllm_config)
+        else:
+            self.connector_worker = MooncakeStoreWorker(vllm_config)
+
+    # ============================================================
+    # Scheduler-side methods
+    # ============================================================
+
+    def get_num_new_matched_tokens(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> tuple[int, bool]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.get_num_new_matched_tokens(
+            request, num_computed_tokens
+        )
+
+    def update_state_after_alloc(
+        self,
+        request: Request,
+        blocks: KVCacheBlocks,
+        num_external_tokens: int,
+    ):
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.update_state_after_alloc(
+            request, blocks, num_external_tokens
+        )
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.build_connector_meta(scheduler_output)
+
+    def request_finished(
+        self,
+        request: Request,
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, block_ids)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        kv_cache_events = connector_output.kv_cache_events
+        if not kv_cache_events or not isinstance(
+            kv_cache_events, MooncakeStoreKVEvents
+        ):
+            return
+
+        if self._kv_cache_events is None:
+            self._kv_cache_events = kv_cache_events
+        else:
+            self._kv_cache_events.add_events(kv_cache_events.get_all_events())
+            self._kv_cache_events.increment_workers(
+                kv_cache_events.get_number_of_workers()
+            )
+
+    def take_events(self) -> Iterable[KVCacheEvent]:
+        if self._kv_cache_events is not None:
+            self._kv_cache_events.aggregate()
+            yield from self._kv_cache_events.get_all_events()
+            self._kv_cache_events.clear_events()
+            self._kv_cache_events = None
+
+    # ============================================================
+    # Worker-side methods
+    # ============================================================
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        assert self.connector_worker is not None
+        self.connector_worker.register_kv_caches(kv_caches)
+
+    def register_cross_layers_kv_cache(
+        self, kv_cache: torch.Tensor, attn_backend: type
+    ):
+        assert self.connector_worker is not None
+        self.connector_worker.register_cross_layers_kv_caches(kv_cache)
+
+    def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
+        # No-op: loads are issued in get_finished() for compute overlap.
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        # No layerwise support - no-op
+        return
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        # No layerwise support - no-op
+        return
+
+    def wait_for_save(self):
+        # No-op: stores are issued in get_finished() for compute overlap.
+        pass
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        assert self.connector_worker is not None
+        metadata = self._get_connector_metadata()
+        assert isinstance(metadata, MooncakeStoreConnectorMetadata)
+        return self.connector_worker.get_finished(finished_req_ids, metadata)
+
+    def get_kv_connector_kv_cache_events(
+        self,
+    ) -> MooncakeStoreKVEvents | None:
+        assert self.connector_worker is not None
+        events = self.connector_worker.get_kv_events()
+        if not events:
+            return None
+
+        kv_events = MooncakeStoreKVEvents(num_workers=1)
+        kv_events.add_events(events)
+        return kv_events

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""Data classes for MooncakeStoreConnector."""
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+)
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.kv_cache_utils import BlockHash
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class KeyMetadata:
+    """Metadata for constructing pool keys."""
+
+    model_name: str
+    tp_rank: int
+    pcp_rank: int
+    dcp_rank: int
+    pp_rank: int
+
+
+@dataclass(order=True)
+class PoolKey:
+    """Key for addressing KV cache blocks in the distributed store."""
+
+    key_metadata: KeyMetadata
+    chunk_hash: str
+
+    def __hash__(self):
+        return hash(
+            (
+                self.key_metadata.model_name,
+                self.key_metadata.tp_rank,
+                self.key_metadata.pcp_rank,
+                self.key_metadata.dcp_rank,
+                self.key_metadata.pp_rank,
+                self.chunk_hash,
+            )
+        )
+
+    def to_string(self) -> str:
+        return (
+            f"{self.key_metadata.model_name}"
+            f"@tp_rank:{self.key_metadata.tp_rank}"
+            f"@pcp{self.key_metadata.pcp_rank}"
+            f"@dcp{self.key_metadata.dcp_rank}"
+            f"@pp_rank:{self.key_metadata.pp_rank}"
+            f"@{self.chunk_hash}"
+        )
+
+
+class ChunkedTokenDatabase:
+    """Maps token positions to store keys and GPU memory addresses."""
+
+    def __init__(self, metadata: KeyMetadata, block_size: int):
+        self.metadata = metadata
+        self.block_size = block_size
+        self.kv_caches_base_addr: list[int] = []
+        self.block_len: list[int] = []
+
+    def _make_key_by_hash(self, chunk_hash: str) -> PoolKey:
+        return PoolKey(self.metadata, chunk_hash)
+
+    def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):
+        self.kv_caches_base_addr = kv_caches_base_addr
+
+    def set_block_len(self, block_len: list[int]):
+        for length in block_len:
+            if length % self.block_size != 0:
+                raise ValueError(f"block_len {length} % {self.block_size} != 0")
+        self.block_len = block_len
+
+    def prepare_value(
+        self, start: int, end: int, block_ids: list[int]
+    ) -> tuple[list[int], list[int], int]:
+        """Compute memory addresses and sizes for a token range.
+
+        Returns:
+            (addr_list, size_list, block_id)
+        """
+        addr_list = []
+        size_list = []
+        block_id = block_ids[start // self.block_size]
+        length = len(self.block_len)
+        for index, base_addr in enumerate(self.kv_caches_base_addr):
+            addr = base_addr + block_id * self.block_len[index % length]
+            size = self.block_len[index % length] // self.block_size * (end - start)
+            addr_list.append(addr)
+            size_list.append(size)
+        return addr_list, size_list, block_id
+
+    def process_tokens(
+        self,
+        token_len: int,
+        block_hashes: list[BlockHash] | list[str],
+        mask_num: int = 0,
+    ) -> Iterable[tuple[int, int, PoolKey]]:
+        """Process tokens and yield (start_idx, end_idx, pool_key) tuples.
+
+        Args:
+            token_len: Total number of tokens.
+            block_hashes: Block hashes for each block.
+            mask_num: Number of tokens to skip from the beginning.
+        """
+        if not block_hashes:
+            return
+        if not isinstance(block_hashes[0], str):
+            block_hashes = [
+                h.hex()  # type: ignore[union-attr]
+                for h in block_hashes
+            ]
+        for chunk_id, hash_val in enumerate(block_hashes):
+            start_idx = chunk_id * self.block_size
+            if start_idx >= token_len:
+                break
+            end_idx = min(start_idx + self.block_size, token_len)
+            if start_idx < mask_num:
+                continue
+            else:
+                yield (
+                    start_idx,
+                    end_idx,
+                    self._make_key_by_hash(
+                        hash_val  # type: ignore[arg-type]
+                    ),
+                )
+
+
+@dataclass
+class LoadSpec:
+    """Specification for loading KV cache from external store."""
+
+    vllm_cached_tokens: int
+    kvpool_cached_tokens: int
+    can_load: bool
+    token_len: int = 0
+
+
+@dataclass
+class RequestTracker:
+    """Tracks per-request state across scheduler ticks."""
+
+    req_id: str
+    token_len: int
+    allocated_block_ids: list[int]
+    num_saved_tokens: int = 0
+    token_ids: list[int] | None = None
+    # Snapshot of the prefill range length at tracker creation time.
+    # For a fresh request this is len(prompt). For a resumed-from-preemption
+    # request it includes previously-generated tokens, which are re-prefilled.
+    prefill_end_tokens: int = 0
+
+    def update(
+        self,
+        new_block_ids: tuple[list[int], ...] | list[int],
+    ) -> None:
+        if len(new_block_ids) == 0:
+            new_block_ids = []
+        elif isinstance(new_block_ids, tuple):
+            new_block_ids = new_block_ids[0]
+        elif isinstance(new_block_ids, list):
+            pass
+        else:
+            raise ValueError(f"Unsupported new_block_ids type {type(new_block_ids)}")
+        self.allocated_block_ids.extend(new_block_ids)
+
+
+@dataclass
+class ReqMeta:
+    """Per-request metadata for store put/get operations."""
+
+    req_id: str
+    token_len_chunk: int
+    block_ids: list[int]
+    block_hashes: list[BlockHash]
+
+    can_save: bool | None = None
+    load_spec: LoadSpec | None = None
+    is_last_chunk: bool | None = None
+    current_event: torch.cuda.Event | None = None
+
+    token_ids: list[int] | None = None
+    original_block_size: int | None = None
+
+    @staticmethod
+    def from_request_tracker(
+        tracker: RequestTracker,
+        block_size: int,
+        load_spec: LoadSpec | None = None,
+        skip_save: bool | None = False,
+        block_hashes: list[BlockHash] | None = None,
+        is_last_chunk: bool | None = None,
+        discard_partial_chunks: bool = True,
+        original_block_size: int | None = None,
+    ) -> "ReqMeta | None":
+        """Create ReqMeta from a RequestTracker."""
+        if block_hashes is None:
+            block_hashes = []
+        input_token_len = tracker.token_len
+
+        chunk_boundary = (
+            cdiv(tracker.num_saved_tokens + 1, block_size) * block_size
+            if discard_partial_chunks
+            else 0
+        )
+        num_tokens_to_save = (
+            (input_token_len // block_size * block_size)
+            if discard_partial_chunks
+            else input_token_len
+        )
+
+        skip_save = skip_save or num_tokens_to_save < chunk_boundary
+        if skip_save and load_spec is None:
+            return None
+
+        if not skip_save:
+            tracker.num_saved_tokens = num_tokens_to_save
+
+        token_ids = None
+        if tracker.token_ids:
+            token_ids = tracker.token_ids
+
+        if load_spec is not None and load_spec.can_load:
+            logger.debug(
+                "Scheduled to load %d tokens for request %s",
+                load_spec.kvpool_cached_tokens,
+                tracker.req_id,
+            )
+        else:
+            load_spec = None
+
+        logger.debug(
+            "request:%s, meta save spec:%s, meta load spec:%s",
+            tracker.req_id,
+            not skip_save,
+            load_spec,
+        )
+        return ReqMeta(
+            req_id=tracker.req_id,
+            token_len_chunk=num_tokens_to_save,
+            block_ids=tracker.allocated_block_ids,
+            can_save=not skip_save,
+            load_spec=load_spec,
+            block_hashes=block_hashes,
+            is_last_chunk=is_last_chunk,
+            token_ids=token_ids,
+            original_block_size=original_block_size,
+        )
+
+
+class MooncakeStoreConnectorMetadata(KVConnectorMetadata):
+    """Metadata passed from scheduler to worker."""
+
+    def __init__(
+        self,
+        unfinished_request_ids: set[str],
+        preempted_req_ids: set[str],
+    ):
+        self.requests: list[ReqMeta] = []
+        self.unfinished_request_ids = unfinished_request_ids
+        self.preempted_req_ids = preempted_req_ids
+
+    def add_request(self, req_meta: ReqMeta) -> None:
+        self.requests.append(req_meta)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

@@ -0,0 +1,380 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""Scheduler-side logic for MooncakeStoreConnector."""
+
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    LoadSpec,
+    MooncakeStoreConnectorMetadata,
+    ReqMeta,
+    RequestTracker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (  # noqa: E501
+    LookupKeyClient,
+)
+from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
+from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+def _new_req_prefill_tokens(request: NewRequestData) -> list[int]:
+    """Tokens this prefill will compute KV for.
+
+    Under the v2 model runner, resumed-from-preemption requests appear in
+    ``scheduled_new_reqs`` with ``prefill_token_ids`` set to the request's full
+    token list (prompt + previously-generated). For all other cases this falls
+    back to the original prompt.
+    """
+    if request.prefill_token_ids is not None:
+        return request.prefill_token_ids
+    assert request.prompt_token_ids is not None
+    return request.prompt_token_ids
+
+
+class MooncakeStoreScheduler:
+    """Scheduler-side component for MooncakeStoreConnector."""
+
+    def __init__(self, vllm_config: VllmConfig):
+        assert vllm_config.kv_transfer_config is not None
+        self.kv_role = vllm_config.kv_transfer_config.kv_role
+        self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
+            "load_async", True
+        )
+        self.client = LookupKeyClient(vllm_config)
+
+        self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
+        self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+        self.original_block_size = vllm_config.cache_config.block_size
+        self._block_size = vllm_config.cache_config.block_size
+        if self.pcp_size > 1:
+            self._block_size *= self.pcp_size
+        if self.dcp_size > 1:
+            self._block_size *= self.dcp_size
+
+        self._discard_partial_chunks = (
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "discard_partial_chunks", True
+            )
+        )
+
+        # Per-request state
+        self.load_specs: dict[str, LoadSpec] = {}  # to be loaded
+        self._request_trackers: dict[str, RequestTracker] = {}  # scheduled new requests
+        self._preempted_req_ids: set[str] = set()  # preempted requests
+        self._unfinished_requests: dict[str, tuple[Request, list[int]]] = {}
+        self._unfinished_request_ids: set[str] = set()
+
+    def get_num_new_matched_tokens(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> tuple[int, bool]:
+        """Check for external KV cache hit."""
+        # Look up against the full prefill range, not just the prompt.
+        if self._discard_partial_chunks:
+            token_len = request.num_tokens // self._block_size * self._block_size
+        else:
+            token_len = request.num_tokens
+
+        if token_len < self._block_size:
+            return 0, False
+
+        num_external_hit_tokens = self.client.lookup(token_len, request.block_hashes)
+
+        if num_external_hit_tokens == request.num_tokens:
+            num_external_hit_tokens -= 1
+
+        if num_external_hit_tokens < num_computed_tokens:
+            need_to_allocate = 0
+        else:
+            need_to_allocate = num_external_hit_tokens - num_computed_tokens
+
+        logger.debug(
+            "Reqid: %s, Total tokens %d, kvpool hit tokens: %d, need to load: %d",
+            request.request_id,
+            request.num_tokens,
+            num_external_hit_tokens,
+            need_to_allocate,
+        )
+
+        if need_to_allocate <= 0:
+            return 0, False
+
+        self.load_specs[request.request_id] = LoadSpec(
+            vllm_cached_tokens=num_computed_tokens,
+            kvpool_cached_tokens=num_external_hit_tokens,
+            can_load=False,
+        )
+
+        return need_to_allocate, self.load_async
+
+    def update_state_after_alloc(
+        self,
+        request: Request,
+        blocks: KVCacheBlocks,
+        num_external_tokens: int,
+    ):
+        """Update state after block allocation."""
+        local_block_ids: list[int] = []
+        if num_external_tokens > 0:
+            local_block_ids = blocks.get_block_ids()[0]
+
+        self._unfinished_requests[request.request_id] = (request, local_block_ids)
+        self._unfinished_request_ids.add(request.request_id)
+
+        if request.request_id not in self.load_specs:
+            return
+
+        if num_external_tokens == 0:
+            self.load_specs[request.request_id].can_load = False
+            return
+
+        assert (
+            num_external_tokens > 0
+            and num_external_tokens
+            == self.load_specs[request.request_id].kvpool_cached_tokens
+            - self.load_specs[request.request_id].vllm_cached_tokens
+        ), (
+            f"Mismatch in number of tokens: {num_external_tokens} vs "
+            f"{self.load_specs[request.request_id].kvpool_cached_tokens} - "
+            f"{self.load_specs[request.request_id].vllm_cached_tokens}"
+            f" for request {request.request_id}"
+        )
+
+        self.load_specs[request.request_id].can_load = True
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        """Build connector metadata for this scheduler step."""
+        force_skip_save = self.kv_role == "kv_consumer"
+
+        for finished_req_id in scheduler_output.finished_req_ids:
+            self.load_specs.pop(finished_req_id, None)
+            self._request_trackers.pop(finished_req_id, None)
+            self._unfinished_requests.pop(finished_req_id, None)
+            self._unfinished_request_ids.discard(finished_req_id)
+            self._preempted_req_ids.discard(finished_req_id)
+
+        preempted_ids = scheduler_output.preempted_req_ids or set()
+        self._preempted_req_ids.update(preempted_ids)
+        for req_id in preempted_ids:
+            self._request_trackers.pop(req_id, None)
+            self._unfinished_requests.pop(req_id, None)
+
+        meta = MooncakeStoreConnectorMetadata(
+            self._unfinished_request_ids,
+            preempted_ids,
+        )
+
+        # Handle new requests
+        for request in scheduler_output.scheduled_new_reqs:
+            load_spec = self.load_specs.pop(request.req_id, None)
+            num_tokens_to_compute = (
+                request.num_computed_tokens
+                + scheduler_output.num_scheduled_tokens[request.req_id]
+            )
+            assert request.req_id in self._unfinished_requests
+            request_tuple = self._unfinished_requests.get(request.req_id)
+            request_real = request_tuple[0]  # type: ignore[index]
+
+            if not isinstance(request.block_ids[0], list):
+                unfolded_block_ids = request.block_ids.copy()
+            else:
+                # TODO: support HMA
+                unfolded_block_ids = request.block_ids[0].copy()
+
+            prefill_tokens = _new_req_prefill_tokens(request)
+            request_tracker = RequestTracker(
+                req_id=request.req_id,
+                token_len=num_tokens_to_compute,
+                allocated_block_ids=unfolded_block_ids,
+                num_saved_tokens=0,
+                token_ids=prefill_tokens[:num_tokens_to_compute],
+                prefill_end_tokens=len(prefill_tokens),
+            )
+            self._request_trackers[request.req_id] = request_tracker
+
+            last_chunk_tokens_num = (
+                (len(prefill_tokens) // self._block_size * self._block_size)
+                if self._discard_partial_chunks
+                else len(prefill_tokens)
+            )
+
+            req_meta = ReqMeta.from_request_tracker(
+                request_tracker,
+                self._block_size,
+                load_spec=load_spec,
+                skip_save=force_skip_save,
+                block_hashes=request_real.block_hashes,
+                is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num),
+                discard_partial_chunks=self._discard_partial_chunks,
+                original_block_size=self.original_block_size,
+            )
+            if req_meta is not None:
+                meta.add_request(req_meta)
+
+        # Handle cached (running, or MRV1 resumed-from-preemption) requests
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        if not force_skip_save:
+            for i, req_id in enumerate(cached_reqs.req_ids):
+                new_block_ids = cached_reqs.new_block_ids[i]
+                if not new_block_ids:
+                    continue
+
+                req_meta = None
+                if req_id in self._preempted_req_ids:
+                    # Resumed after preemption
+                    if isinstance(new_block_ids, tuple):
+                        block_ids_list = new_block_ids[0].copy()
+                    else:
+                        block_ids_list = new_block_ids.copy()
+                    self._preempted_req_ids.discard(req_id)
+                    load_spec = self.load_specs.pop(req_id, None)
+                    request_tuple = self._unfinished_requests.get(req_id)
+                    request_real = request_tuple[0]  # type: ignore[index]
+                    num_tokens_to_compute = (
+                        request_real.num_computed_tokens
+                        + scheduler_output.num_scheduled_tokens[req_id]
+                    )
+                    # On resume, the request re-prefills prompt + previously
+                    # generated tokens (all_token_ids).
+                    prefill_tokens = list(request_real.all_token_ids)
+                    request_tracker = RequestTracker(
+                        req_id=req_id,
+                        token_len=num_tokens_to_compute,
+                        allocated_block_ids=block_ids_list,
+                        num_saved_tokens=0,
+                        token_ids=prefill_tokens[:num_tokens_to_compute].copy(),
+                        prefill_end_tokens=len(prefill_tokens),
+                    )
+                    self._request_trackers[req_id] = request_tracker
+
+                    last_chunk_tokens_num = (
+                        (len(prefill_tokens) // self._block_size * self._block_size)
+                        if self._discard_partial_chunks
+                        else len(prefill_tokens)
+                    )
+                    req_meta = ReqMeta.from_request_tracker(
+                        request_tracker,
+                        self._block_size,
+                        load_spec=load_spec,
+                        skip_save=force_skip_save,
+                        block_hashes=request_real.block_hashes,
+                        is_last_chunk=(
+                            request_tracker.token_len >= last_chunk_tokens_num
+                        ),
+                        discard_partial_chunks=self._discard_partial_chunks,
+                        original_block_size=self.original_block_size,
+                    )
+                else:
+                    # Decode/chunked request
+                    request_tracker = self._request_trackers[req_id]
+                    num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
+                    req_tuple = self._unfinished_requests.get(req_id)
+                    if req_tuple:
+                        unfinished_req = req_tuple[0]
+                        num_current_tokens = request_tracker.token_len
+                        new_token_ids = unfinished_req.all_token_ids[
+                            num_current_tokens : num_current_tokens + num_new_tokens
+                        ]
+                        request_tracker.token_len += len(new_token_ids)
+                    else:
+                        raise ValueError(
+                            f"Request {req_id} is not in _unfinished_requests"
+                        )
+                    num_computed_token = cached_reqs.num_computed_tokens[i]
+                    # Use the tracker's snapshot of the prefill range so resumed
+                    # requests keep saving past the original prompt boundary.
+                    prefill_end = request_tracker.prefill_end_tokens
+                    if num_computed_token >= prefill_end:
+                        continue
+                    request_tracker.update(new_block_ids)
+
+                    last_chunk_tokens_num = (
+                        (prefill_end // self._block_size * self._block_size)
+                        if self._discard_partial_chunks
+                        else prefill_end
+                    )
+                    req_meta = ReqMeta.from_request_tracker(
+                        request_tracker,
+                        self._block_size,
+                        load_spec=None,
+                        skip_save=force_skip_save,
+                        block_hashes=unfinished_req.block_hashes,
+                        is_last_chunk=(
+                            request_tracker.token_len >= last_chunk_tokens_num
+                        ),
+                        discard_partial_chunks=self._discard_partial_chunks,
+                        original_block_size=self.original_block_size,
+                    )
+
+                if req_meta is not None:
+                    meta.add_request(req_meta)
+
+        # Handle requests with pending load specs not yet scheduled
+        request_ids = [req.req_id for req in scheduler_output.scheduled_new_reqs]
+        for request_id, (
+            unfinished_req,
+            block_ids,
+        ) in self._unfinished_requests.items():
+            if request_id not in request_ids and request_id not in cached_reqs.req_ids:
+                load_spec = self.load_specs.pop(request_id, None)
+                if not load_spec:
+                    continue
+                num_tokens_to_compute = load_spec.kvpool_cached_tokens
+                if (num_tokens_to_compute % self._block_size != 0) and (
+                    num_tokens_to_compute == unfinished_req.num_tokens - 1
+                ):
+                    num_tokens_to_compute = num_tokens_to_compute + 1
+                request_tracker = RequestTracker(
+                    req_id=request_id,
+                    token_len=num_tokens_to_compute,
+                    allocated_block_ids=block_ids,
+                    num_saved_tokens=0,
+                )
+                self._request_trackers[request_id] = request_tracker
+                req_meta = ReqMeta.from_request_tracker(
+                    request_tracker,
+                    self._block_size,
+                    load_spec=load_spec,
+                    skip_save=None,
+                    block_hashes=unfinished_req.block_hashes,
+                    discard_partial_chunks=self._discard_partial_chunks,
+                )
+                if req_meta is not None:
+                    meta.add_request(req_meta)
+
+        return meta
+
+    def request_finished(
+        self,
+        request: Request,
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """Determine whether to delay freeing blocks for async save."""
+        if self.kv_role == "kv_consumer":
+            return False, None
+        tracker = self._request_trackers.get(request.request_id)
+        assert tracker is not None
+        if tracker.num_saved_tokens <= 0:
+            return False, None
+        delay_free_blocks = len(block_ids) > 0
+        if delay_free_blocks:
+            logger.debug(
+                "Delaying free of %d blocks for request %s",
+                len(block_ids),
+                request.request_id,
+            )
+        return delay_free_blocks, None

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -0,0 +1,979 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# The transfer-thread scaffolding (KVTransferThread, KVCacheStoreSendingThread,
+# KVCacheStoreRecvingThread) is adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""Worker-side logic for MooncakeStoreConnector.
+
+Includes the store worker, transfer threads, lookup server,
+and MooncakeDistributedStore integration.
+"""
+
+import json
+import os
+import queue
+import threading
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any
+
+import regex as re
+import torch
+import zmq
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.distributed import (
+    get_dcp_group,
+    get_pcp_group,
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.distributed.kv_events import BlockStored
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
+    get_mooncake_dp_engine_index,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    MooncakeStoreConnectorMetadata,
+    ReqMeta,
+)
+from vllm.logger import init_logger
+from vllm.utils.network_utils import get_ip, make_zmq_socket
+from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
+
+logger = init_logger(__name__)
+
+DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
+DEFAULT_LOCAL_BUFFER_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
+MOONCAKE_NO_AVAILABLE_HANDLE = -200
+
+
+@dataclass
+class MooncakeStoreConfig:
+    """Configuration for MooncakeDistributedStore."""
+
+    metadata_server: str
+    global_segment_size: int
+    local_buffer_size: int
+    protocol: str
+    device_name: str
+    master_server_address: str
+
+    @staticmethod
+    def from_file(file_path: str) -> "MooncakeStoreConfig":
+        with open(file_path) as file:
+            config = json.load(file)
+        return MooncakeStoreConfig(
+            metadata_server=config.get("metadata_server", ""),
+            global_segment_size=_parse_size(
+                config.get("global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE)
+            ),
+            local_buffer_size=_parse_size(
+                config.get("local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE)
+            ),
+            protocol=config.get("protocol", "rdma"),
+            device_name=config.get("device_name", ""),
+            master_server_address=config.get("master_server_address", ""),
+        )
+
+    @staticmethod
+    def load_from_env() -> "MooncakeStoreConfig":
+        config_path = os.getenv("MOONCAKE_CONFIG_PATH")
+        if not config_path:
+            raise ValueError(
+                "The environment variable 'MOONCAKE_CONFIG_PATH' is not set."
+            )
+        return MooncakeStoreConfig.from_file(config_path)
+
+
+def _parse_size(value: Any) -> int:
+    """Parse storage size strings with units: GB, MB, KB, B."""
+    if isinstance(value, int):
+        return value
+    if not isinstance(value, str):
+        try:
+            return int(value)
+        except (TypeError, ValueError) as e:
+            raise TypeError(f"Unsupported type for size: {type(value)}") from e
+
+    cleaned = value.strip().lower()
+    if not cleaned:
+        raise ValueError("Size cannot be empty.")
+
+    unit_multipliers = {
+        "gb": 1024**3,
+        "mb": 1024**2,
+        "kb": 1024,
+        "b": 1,
+    }
+    match = re.match(r"^\s*([\d.]+)\s*(gb|mb|kb|b)?\s*$", cleaned)
+    if not match:
+        raise ValueError(f"Invalid format: '{value}'")
+
+    number_str = match.group(1)
+    unit = match.group(2) or "b"
+    multiplier = unit_multipliers[unit]
+
+    try:
+        numeric_value = float(number_str)
+    except ValueError as exc:
+        raise ValueError(f"Invalid numeric value '{number_str}' in: '{value}'") from exc
+    return int(numeric_value * multiplier)
+
+
+# ============================================================
+# Transfer Threads
+# ============================================================
+
+
+class KVTransferThread(threading.Thread):
+    """Base class for async KV cache transfer threads."""
+
+    def __init__(
+        self,
+        store: Any,
+        token_database: ChunkedTokenDatabase,
+        block_size: int,
+        tp_rank: int,
+        ready_event: threading.Event,
+        name: str,
+    ):
+        super().__init__(daemon=True, name=name)
+        self.store = store
+        self.ready_event = ready_event
+        self.block_size = block_size
+        self.tp_rank = tp_rank
+        self.token_database = token_database
+        self.done_task_lock = threading.Lock()
+        self.request_queue: queue.Queue[Any] = queue.Queue()
+        self.finished_requests: set[str] = set()
+        self.kv_event_lock = threading.Lock()
+        self.kv_events: list[BlockStored] = []
+
+    def add_request(self, request: ReqMeta) -> None:
+        self.request_queue.put(request)
+
+    def get_and_clear_finished_requests(self) -> set[str]:
+        with self.done_task_lock:
+            finished = self.finished_requests.copy()
+            self.finished_requests.clear()
+        return finished
+
+    def set_finished_request(self, req_id: str):
+        with self.done_task_lock:
+            self.finished_requests.add(req_id)
+
+    def run(self):
+        self.ready_event.set()
+        while True:
+            try:
+                request_data = self.request_queue.get()
+                if request_data is None:
+                    logger.warning("Received a None request!")
+                    self.request_queue.task_done()
+                    continue
+                self._handle_request(request_data)
+            except Exception as e:
+                logger.error("Error in %s: %s", self.name, e)
+
+    def _handle_request(self, req_meta: Any):
+        pass
+
+    def update_kv_event(self, events: list[BlockStored]):
+        with self.kv_event_lock:
+            self.kv_events.extend(events)
+
+    def get_kv_events(self) -> list[BlockStored]:
+        with self.kv_event_lock:
+            events = self.kv_events.copy()
+            self.kv_events.clear()
+        return events
+
+
+class KVCacheStoreSendingThread(KVTransferThread):
+    """Background thread for storing KV cache blocks to the store."""
+
+    def __init__(
+        self,
+        store: Any,
+        token_database: ChunkedTokenDatabase,
+        block_size: int,
+        tp_rank: int,
+        put_step: int,
+        kv_role: str,
+        ready_event: threading.Event,
+        enable_kv_event: bool = False,
+    ):
+        super().__init__(
+            store,
+            token_database,
+            block_size,
+            tp_rank,
+            ready_event,
+            name="KVCacheStoreSendingThread",
+        )
+        self.put_step = put_step
+        self.kv_role = kv_role
+        self.stored_requests: defaultdict[str, int] = defaultdict(int)
+        self.enable_kv_event = enable_kv_event
+
+        # Pause store requests when CPU offloading is under pressure.
+        self._store_pressure_active = False
+        self._skip_store_requests: set[str] = set()
+
+    def add_stored_request(self, req_id: str):
+        with self.done_task_lock:
+            self.stored_requests[req_id] += 1
+
+    def dec_stored_request(self, req_id: str):
+        with self.done_task_lock:
+            if req_id in self.stored_requests:
+                self.stored_requests[req_id] -= 1
+
+    def delete_finished_stored_request(self, req_id: str):
+        with self.done_task_lock:
+            if req_id in self.stored_requests:
+                del self.stored_requests[req_id]
+            self._skip_store_requests.discard(req_id)
+
+    def _should_skip_request(self, req_id: str) -> bool:
+        with self.done_task_lock:
+            return self._store_pressure_active and req_id in self._skip_store_requests
+
+    def _mark_request_skipped_for_pressure(self, req_id: str) -> bool:
+        with self.done_task_lock:
+            already_skipped = req_id in self._skip_store_requests
+            self._store_pressure_active = True
+            self._skip_store_requests.add(req_id)
+        return already_skipped
+
+    def _clear_store_pressure(self) -> bool:
+        with self.done_task_lock:
+            if not self._store_pressure_active and not self._skip_store_requests:
+                return False
+            self._store_pressure_active = False
+            self._skip_store_requests.clear()
+        return True
+
+    def _handle_request(self, req_meta: ReqMeta):
+        token_len = req_meta.token_len_chunk
+        block_ids = req_meta.block_ids
+        req_id = req_meta.req_id
+        current_event = req_meta.current_event
+
+        if req_id not in self.stored_requests:
+            self.request_queue.task_done()
+            return
+        if self._should_skip_request(req_id):
+            logger.debug(
+                "Skipping Mooncake store for request %s while CPU offloading "
+                "is under pressure",
+                req_id,
+            )
+            self.dec_stored_request(req_id)
+            self.request_queue.task_done()
+            return
+
+        starts = []
+        ends = []
+        keys = []
+        block_hashes: list[BlockHash] = []
+        for index, (start, end, key) in enumerate(
+            self.token_database.process_tokens(token_len, req_meta.block_hashes)
+        ):
+            starts.append(start)
+            ends.append(end)
+            keys.append(key.to_string())
+            block_hashes.append(req_meta.block_hashes[index])
+
+        # Apply put_step striding for TP
+        starts = starts[self.tp_rank % self.put_step :: self.put_step]
+        ends = ends[self.tp_rank % self.put_step :: self.put_step]
+        keys = keys[self.tp_rank % self.put_step :: self.put_step]
+        block_hashes = block_hashes[self.tp_rank % self.put_step :: self.put_step]
+
+        if not keys:
+            self.dec_stored_request(req_id)
+            return
+
+        # Check which blocks already exist (dedup)
+        exists_states = self.store.batch_is_exist(keys)
+        missing_indices = [i for i, exists in enumerate(exists_states) if exists != 1]
+
+        if not missing_indices:
+            self.dec_stored_request(req_id)
+            return
+
+        starts = [starts[i] for i in missing_indices]
+        ends = [ends[i] for i in missing_indices]
+        keys = [keys[i] for i in missing_indices]
+        block_hashes = [block_hashes[i] for i in missing_indices]
+
+        logger.debug(
+            "Storing KV cache for %d out of %d blocks "
+            "(missing_count=%d) for request %s",
+            len(keys),
+            token_len // self.block_size,
+            len(missing_indices),
+            req_id,
+        )
+
+        addrs = []
+        sizes = []
+        stored_events: list[BlockStored] = []
+        prev_key = None
+        new_block_hashes = [maybe_convert_block_hash(bh) for bh in block_hashes]
+
+        for index, start in enumerate(starts):
+            addr, size, _ = self.token_database.prepare_value(
+                start, ends[index], block_ids
+            )
+            addrs.append(addr)
+            sizes.append(size)
+
+            if self.enable_kv_event:
+                token_ids = (
+                    req_meta.token_ids[start : ends[index]]
+                    if req_meta.token_ids is not None
+                    else None
+                )
+                stored_event = BlockStored(
+                    block_hashes=[new_block_hashes[index]],
+                    parent_block_hash=prev_key,
+                    token_ids=token_ids,
+                    block_size=req_meta.original_block_size,
+                    lora_id=None,
+                    medium="cpu",
+                    lora_name=None,
+                )
+                stored_events.append(stored_event)
+                prev_key = new_block_hashes[index]
+
+        if current_event is not None:
+            current_event.synchronize()
+
+        try:
+            res = self.store.batch_put_from_multi_buffers(keys, addrs, sizes)
+            failed = [i for i, v in enumerate(res) if v < 0]
+            if failed:
+                # Compute total bytes attempted for this batch
+                total_bytes = sum(sum(s) if isinstance(s, list) else s for s in sizes)
+                failed_codes = set(res[i] for i in failed)
+                logger.warning(
+                    "batch_put failed: %d/%d keys failed "
+                    "(codes=%s, batch_bytes=%d, num_keys=%d), "
+                    "first_key=%s",
+                    len(failed),
+                    len(keys),
+                    failed_codes,
+                    total_bytes,
+                    len(keys),
+                    keys[0] if keys else "N/A",
+                )
+                if (
+                    MOONCAKE_NO_AVAILABLE_HANDLE in failed_codes
+                    and not self._mark_request_skipped_for_pressure(req_id)
+                ):
+                    logger.warning(
+                        "Detected Mooncake CPU offloading pressure "
+                        "(NO_AVAILABLE_HANDLE); skipping future store "
+                        "batches for request %s until a later store "
+                        "batch succeeds",
+                        req_id,
+                    )
+            elif self._clear_store_pressure():
+                logger.info(
+                    "Mooncake CPU offloading pressure cleared after a "
+                    "successful store batch"
+                )
+        except Exception as e:
+            logger.error("Failed to put key %s, error: %s", keys, e)
+
+        if self.enable_kv_event and stored_events:
+            self.update_kv_event(stored_events)
+
+        self.dec_stored_request(req_id)
+        self.request_queue.task_done()
+
+
+class KVCacheStoreRecvingThread(KVTransferThread):
+    """Background thread for loading KV cache blocks from the store."""
+
+    def __init__(
+        self,
+        store: Any,
+        token_database: ChunkedTokenDatabase,
+        block_size: int,
+        tp_rank: int,
+        ready_event: threading.Event,
+    ):
+        super().__init__(
+            store,
+            token_database,
+            block_size,
+            tp_rank,
+            ready_event,
+            name="KVCacheStoreRecvingThread",
+        )
+
+    def _handle_request(self, req_meta: ReqMeta):
+        token_len = req_meta.load_spec.token_len  # type: ignore[union-attr]
+        req_id = req_meta.req_id
+        mask_num = (
+            req_meta.load_spec.vllm_cached_tokens  # type: ignore[union-attr]
+            // self.block_size
+            * self.block_size
+        )
+
+        addr_list = []
+        size_list = []
+        key_list = []
+        for start, end, key in self.token_database.process_tokens(
+            token_len, req_meta.block_hashes, mask_num
+        ):
+            addr, size, _ = self.token_database.prepare_value(
+                start, end, req_meta.block_ids
+            )
+            key_list.append(key.to_string())
+            addr_list.append(addr)
+            size_list.append(size)
+
+        # Rotate lists by tp_rank for load balancing
+        key_list_c = (
+            key_list[self.tp_rank % len(key_list) :]
+            + key_list[: self.tp_rank % len(key_list)]
+        )
+        addr_list_c = (
+            addr_list[self.tp_rank % len(addr_list) :]
+            + addr_list[: self.tp_rank % len(addr_list)]
+        )
+        size_list_c = (
+            size_list[self.tp_rank % len(size_list) :]
+            + size_list[: self.tp_rank % len(size_list)]
+        )
+
+        try:
+            res = self.store.batch_get_into_multi_buffers(
+                key_list_c, addr_list_c, size_list_c
+            )
+            failed = [
+                (key, value)
+                for key, value in zip(key_list_c, res, strict=True)
+                if value < 0
+            ]
+            if failed:
+                logger.warning(
+                    "Failed to get %d Mooncake keys (batch_keys=%d, first_failures=%s)",
+                    len(failed),
+                    len(key_list_c),
+                    failed[:3],
+                )
+        except Exception as e:
+            logger.warning(
+                "Failed to get Mooncake batch %s, error: %s",
+                key_list_c[:3],
+                e,
+            )
+
+        self.set_finished_request(req_id)
+        self.request_queue.task_done()
+
+
+# ============================================================
+# Store Worker
+# ============================================================
+
+
+class MooncakeStoreWorker:
+    """Worker-side component for MooncakeStoreConnector."""
+
+    def __init__(self, vllm_config: VllmConfig):
+        try:
+            from mooncake.store import MooncakeDistributedStore  # type: ignore
+        except ImportError as e:
+            raise ImportError(
+                "Please install mooncake by following the instructions at "
+                "https://github.com/kvcache-ai/Mooncake/blob/main/doc/"
+                "en/build.md to run vLLM with MooncakeStoreConnector."
+            ) from e
+
+        model_config = vllm_config.model_config
+        parallel_config = vllm_config.parallel_config
+
+        self.dp_rank = get_mooncake_dp_engine_index(parallel_config)
+        self.tp_rank = get_tensor_model_parallel_rank()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.pp_size = parallel_config.pipeline_parallel_size
+        self.pp_rank = (parallel_config.rank // self.tp_size) % self.pp_size
+
+        self.pcp_size = get_pcp_group().world_size
+        self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0
+        self.dcp_size = get_dcp_group().world_size
+        self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0
+
+        assert vllm_config.kv_transfer_config is not None
+        self.kv_role = vllm_config.kv_transfer_config.kv_role
+        self.load_async = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
+            "load_async", True
+        )
+        self.cache_config = vllm_config.cache_config
+        self.original_block_size = self.cache_config.block_size
+        self.block_size = self.cache_config.block_size
+        if self.pcp_size > 1:
+            self.block_size *= self.pcp_size
+        if self.dcp_size > 1:
+            self.block_size *= self.dcp_size
+        self.num_layers = model_config.get_num_layers(parallel_config)
+
+        self.use_mla = False
+        if (
+            hasattr(model_config, "use_mla")
+            and isinstance(model_config.use_mla, bool)
+            and model_config.use_mla
+        ):
+            self.use_mla = True
+
+        if self.use_mla:
+            self.num_kv_head = 1
+        else:
+            self.num_kv_head = model_config.get_total_num_kv_heads()
+
+        if self.num_kv_head < self.tp_size:
+            self.put_step = self.tp_size // self.num_kv_head
+            self.head_or_tp_rank = self.tp_rank // self.put_step
+        else:
+            self.head_or_tp_rank = self.tp_rank
+            self.put_step = 1
+
+        self.metadata = KeyMetadata(
+            model_name=model_config.model.rstrip("/").split("/")[-1],
+            tp_rank=self.head_or_tp_rank,
+            pcp_rank=self.pcp_rank,
+            dcp_rank=self.dcp_rank,
+            pp_rank=self.pp_rank,
+        )
+
+        self.token_database = ChunkedTokenDatabase(self.metadata, self.block_size)
+
+        # Initialize MooncakeDistributedStore with its own TransferEngine
+        store_config = MooncakeStoreConfig.load_from_env()
+        self.store = MooncakeDistributedStore()
+
+        local_seg = get_ip()
+        config_dict = {
+            "local_hostname": local_seg,
+            "metadata_server": store_config.metadata_server,
+            "global_segment_size": str(store_config.global_segment_size),
+            "local_buffer_size": str(store_config.local_buffer_size),
+            "protocol": store_config.protocol,
+            "rdma_devices": store_config.device_name,
+            "master_server_addr": store_config.master_server_address,
+        }
+        ret = self.store.setup(config_dict)
+        if ret != 0:
+            msg = "Initialize MooncakeDistributedStore failed."
+            logger.error(msg)
+            raise RuntimeError(msg)
+
+        kv_event_config = vllm_config.kv_events_config
+        self.enable_kv_events = False
+        if kv_event_config and kv_event_config.enable_kv_cache_events:
+            self.enable_kv_events = True
+
+        self.kv_send_thread: KVCacheStoreSendingThread | None = None
+        self.kv_recv_thread: KVCacheStoreRecvingThread | None = None
+        self.finished_store_req: set[str] = set()
+
+        # Start lookup server on rank 0 for scheduler-side prefix queries
+        self.lookup_server: LookupKeyServer | None = None
+        if vllm_config.parallel_config.rank == 0:
+            self.lookup_server = LookupKeyServer(self, vllm_config)
+
+    def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
+        """Register a cross-layers KV cache tensor.
+
+        Wraps the unified tensor in a single-entry dict so that the
+        existing stride-based logic in register_kv_caches() produces
+        the correct single-segment result (block_len = page_size * num_layers).
+        """
+        self.register_kv_caches({"__cross_layer__": kv_cache})
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """Register KV cache tensors and start transfer threads."""
+        # TODO(yifan): we haven't supported HMA yet.
+        first_kv_cache = next(iter(kv_caches.values()))
+
+        # num_blocks from cache_config is authoritative (set after
+        # profiling, before KV cache allocation).
+        assert self.cache_config.num_gpu_blocks is not None
+        self.num_blocks = self.cache_config.num_gpu_blocks
+
+        # Detect the KV cache memory layout using the stride-based
+        # approach from simple_kv_offload/worker.py.
+        #
+        # The physical layout varies across attention backends:
+        #   FlashAttn/ROCm : (2, num_blocks, ...) → K/V outermost
+        #   FlashInfer/MLA : (num_blocks, ...)    → blocks outermost
+        #
+        # We derive page_size_bytes = storage.nbytes() // num_blocks,
+        # then classify dims: any dim whose byte-stride exceeds
+        # page_size_bytes must be an outer segment dim (e.g. the K/V
+        # dim of size 2).  For those backends we register each segment
+        # (K, V) as a separate base-address so that the per-block
+        # offset arithmetic in prepare_value() stays correct.
+        storage = first_kv_cache.untyped_storage()
+        el = first_kv_cache.element_size()
+        page_size_bytes = storage.nbytes() // self.num_blocks
+        outer_dims = [
+            d
+            for d in range(first_kv_cache.ndim)
+            if first_kv_cache.stride(d) * el > page_size_bytes
+        ]
+
+        # Register buffers with the store (deduplicate shared storages)
+        # and record per-segment base addresses for every layer.
+        seen_ptrs: set[int] = set()
+        self.kv_caches_base_addr: list[int] = []
+        self.block_len: list[int] = []
+
+        for cache in kv_caches.values():
+            cache_storage = cache.untyped_storage()
+            base_addr = cache_storage.data_ptr()
+            region_len = cache_storage.nbytes()
+
+            if base_addr not in seen_ptrs:
+                seen_ptrs.add(base_addr)
+                ret = self.store.register_buffer(base_addr, region_len)
+                if ret != 0:
+                    logger.error(
+                        "register_buffer failed for addr %#x len %d: %d",
+                        base_addr,
+                        region_len,
+                        ret,
+                    )
+
+            if not outer_dims:
+                # Blocks-first layout (FlashInfer / MLA): one segment.
+                self.kv_caches_base_addr.append(base_addr)
+                self.block_len.append(page_size_bytes)
+            else:
+                # K/V-first layout (FlashAttn / ROCm): split segments.
+                seg_stride = cache.stride(outer_dims[0]) * el
+                for idx in range(cache.shape[outer_dims[0]]):
+                    self.kv_caches_base_addr.append(base_addr + idx * seg_stride)
+                    self.block_len.append(seg_stride // self.num_blocks)
+
+        logger.info(
+            "Registering KV_Caches. use_mla: %s, shape %s, "
+            "num_blocks: %d, block_len: %s, "
+            "per_key_bytes: %d, "
+            "num_segments: %d",
+            self.use_mla,
+            first_kv_cache.shape,
+            self.num_blocks,
+            list(set(self.block_len)),
+            sum(self.block_len),
+            len(self.kv_caches_base_addr),
+        )
+
+        self.token_database.set_kv_caches_base_addr(self.kv_caches_base_addr)
+        self.token_database.set_block_len(self.block_len)
+
+        # Start transfer threads
+        if self.kv_role in ["kv_producer", "kv_both"]:
+            ready_event_sending = threading.Event()
+            self.kv_send_thread = KVCacheStoreSendingThread(
+                self.store,
+                self.token_database,
+                self.block_size,
+                self.tp_rank,
+                self.put_step,
+                self.kv_role,
+                ready_event_sending,
+                self.enable_kv_events,
+            )
+            self.kv_send_thread.start()
+
+        ready_event_recving = threading.Event()
+        self.kv_recv_thread = KVCacheStoreRecvingThread(
+            self.store,
+            self.token_database,
+            self.block_size,
+            self.tp_rank,
+            ready_event_recving,
+        )
+        self.kv_recv_thread.start()
+        ready_event_recving.wait()
+
+    def start_load_kv(
+        self,
+        metadata: MooncakeStoreConnectorMetadata,
+    ):
+        """No-op: loads are issued in get_finished() for overlap."""
+        pass
+
+    def wait_for_save(
+        self,
+        metadata: MooncakeStoreConnectorMetadata,
+    ):
+        """No-op: stores are issued in get_finished() for overlap."""
+        pass
+
+    def get_finished(
+        self,
+        finished_req_ids: set[str],
+        meta: MooncakeStoreConnectorMetadata,
+    ) -> tuple[set[str], set[str]]:
+        """Issue all I/O and get completed send/recv request IDs.
+
+        All load and store I/O requests are issued here (after model
+        compute is launched on the compute stream) for better
+        compute-I/O overlap.
+        """
+        # Issue async loads
+        for request in meta.requests:
+            load_spec = request.load_spec
+            if load_spec is None or not load_spec.can_load:
+                continue
+
+            token_len = request.token_len_chunk
+            if (load_spec.kvpool_cached_tokens % self.block_size != 0) and (
+                load_spec.kvpool_cached_tokens == token_len - 1
+            ):
+                token_len = load_spec.kvpool_cached_tokens + 1
+            else:
+                token_len = load_spec.kvpool_cached_tokens
+            load_spec.token_len = token_len
+
+            assert self.kv_recv_thread is not None
+            self.kv_recv_thread.add_request(request)
+
+        assert self.load_async, "load_async must be True for better performance."
+        # Issue stores with CUDA event synchronization
+        if self.kv_role in ["kv_producer", "kv_both"]:
+            current_event = None
+            for request in meta.requests:
+                if request.can_save:
+                    current_event = torch.cuda.Event()
+                    current_event.record()
+                    break
+
+            for request in meta.requests:
+                if not request.can_save:
+                    continue
+                request.current_event = current_event
+                assert self.kv_send_thread is not None
+                self.kv_send_thread.add_stored_request(request.req_id)
+                self.kv_send_thread.add_request(request)
+
+        # Check completion of previously queued transfers
+        done_sending = (
+            self._get_and_clear_finished_sending(finished_req_ids, meta)
+            if self.kv_role in ["kv_producer", "kv_both"]
+            else set()
+        )
+
+        done_recving = (
+            self.kv_recv_thread.get_and_clear_finished_requests()
+            if self.load_async and self.kv_recv_thread is not None
+            else set()
+        )
+
+        logger.debug(
+            "Completed send: %d, recv: %d, tp_rank: %d",
+            len(done_sending),
+            len(done_recving),
+            self.tp_rank,
+        )
+        return done_sending, done_recving
+
+    def _get_and_clear_finished_sending(
+        self,
+        finished_req_ids: set[str],
+        meta: MooncakeStoreConnectorMetadata,
+    ) -> set[str]:
+        assert self.kv_send_thread is not None
+        finished_sending: set[str] = set()
+
+        for req_id in meta.preempted_req_ids:
+            self.kv_send_thread.delete_finished_stored_request(req_id)
+
+        for req_id in self.kv_send_thread.stored_requests.copy():
+            if (
+                self.kv_send_thread.stored_requests[req_id] == 0
+                and req_id in self.finished_store_req
+            ):
+                self.finished_store_req.remove(req_id)
+                finished_sending.add(req_id)
+                self.kv_send_thread.delete_finished_stored_request(req_id)
+
+        for req_id in finished_req_ids:
+            req_remain_jobs = self.kv_send_thread.stored_requests.get(req_id)
+            if req_remain_jobs == 0:
+                finished_sending.add(req_id)
+                self.kv_send_thread.delete_finished_stored_request(req_id)
+            elif req_remain_jobs is not None:
+                self.finished_store_req.add(req_id)
+
+        return finished_sending
+
+    def lookup(
+        self,
+        token_len: int,
+        block_hashes: list[BlockHash],
+    ) -> int:
+        """Check how many prefix tokens exist in the store.
+
+        Checks across all TP ranks and PP ranks.
+        """
+        end = 0
+        keys: list[str] = []
+        try:
+            starts: list[int] = []
+            for start, end, key in self.token_database.process_tokens(
+                token_len, block_hashes
+            ):
+                keys.append(key.to_string())
+                starts.append(start)
+
+            # Expand keys for all TP ranks
+            multi_tp_keys = keys[:]
+            for i in range(1, min(self.tp_size, self.num_kv_head)):
+                for item in keys:
+                    new_str = item.replace("@tp_rank:0", f"@tp_rank:{i}", 1)
+                    multi_tp_keys.append(new_str)
+
+            # Expand keys for all PP ranks
+            pp_base_keys = multi_tp_keys.copy()
+            for i in range(1, self.pp_size):
+                for item in pp_base_keys:
+                    new_str = item.replace("@pp_rank:0", f"@pp_rank:{i}", 1)
+                    multi_tp_keys.append(new_str)
+
+            res = self.store.batch_is_exist(multi_tp_keys)
+
+            num_block = len(keys)
+            multi_tp_values = [
+                res[i * num_block : (i + 1) * num_block]
+                for i in range(min(self.tp_size, self.num_kv_head) * self.pp_size)
+            ]
+            index = self._find_min_first_non_one_index(multi_tp_values)
+            if index != -1:
+                return starts[index]
+        except Exception as e:
+            logger.error("Remote connection failed in lookup: %s", e)
+            return 0
+        return end
+
+    @staticmethod
+    def _find_min_first_non_one_index(
+        arr: list[list[int]],
+    ) -> int:
+        try:
+            return min(idx for row in arr for idx, val in enumerate(row) if val != 1)
+        except ValueError:
+            return -1
+
+    def get_kv_events(self) -> list[BlockStored]:
+        if self.enable_kv_events and self.kv_send_thread is not None:
+            return self.kv_send_thread.get_kv_events()
+        return []
+
+
+# ============================================================
+# Lookup Key Server
+# ============================================================
+
+
+class LookupKeyServer:
+    """ZMQ server on worker rank 0 for handling prefix lookup queries."""
+
+    def __init__(
+        self,
+        store_worker: MooncakeStoreWorker,
+        vllm_config: VllmConfig,
+    ):
+        self.decoder = MsgpackDecoder()
+        self.ctx = zmq.Context()  # type: ignore[attr-defined]
+        socket_path = get_zmq_rpc_path_lookup(vllm_config)
+        self._ipc_path = socket_path.removeprefix("ipc://")
+        if os.path.exists(self._ipc_path):
+            os.unlink(self._ipc_path)
+        self.socket = make_zmq_socket(
+            self.ctx,
+            socket_path,
+            zmq.REP,  # type: ignore[attr-defined]
+            bind=True,
+        )
+
+        self.store_worker = store_worker
+        self.running = True
+
+        def process_request():
+            while self.running:
+                all_frames = self.socket.recv_multipart(copy=False)
+                token_len = int.from_bytes(all_frames[0], byteorder="big")
+                hash_frames = all_frames[1:]
+                hashes_str = self.decoder.decode(hash_frames)
+                result = self.store_worker.lookup(token_len, hashes_str)
+                response = result.to_bytes(4, "big")
+                self.socket.send(response)
+
+        self.thread = threading.Thread(target=process_request, daemon=True)
+        self.thread.start()
+
+    def close(self):
+        self.socket.close(linger=0)
+        if os.path.exists(self._ipc_path):
+            os.unlink(self._ipc_path)
+
+
+# ============================================================
+# Lookup Key Client
+# ============================================================
+
+
+class LookupKeyClient:
+    """ZMQ client for querying prefix cache hits from worker."""
+
+    def __init__(self, vllm_config: VllmConfig):
+        self.encoder = MsgpackEncoder()
+        self.ctx = zmq.Context()  # type: ignore[attr-defined]
+        socket_path = get_zmq_rpc_path_lookup(vllm_config)
+        self.socket = make_zmq_socket(
+            self.ctx,
+            socket_path,
+            zmq.REQ,  # type: ignore[attr-defined]
+            bind=False,
+        )
+
+    def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
+        hash_strs = [h.hex() for h in block_hashes]
+        hash_frames = self.encoder.encode(hash_strs)
+        token_len_bytes = token_len.to_bytes(4, byteorder="big")
+        all_frames = [token_len_bytes] + list(hash_frames)
+        self.socket.send_multipart(all_frames, copy=False)
+        resp = self.socket.recv()
+        result = int.from_bytes(resp, "big")
+        return result
+
+    def close(self):
+        self.socket.close(linger=0)
+
+
+def get_zmq_rpc_path_lookup(vllm_config: VllmConfig) -> str:
+    """Construct IPC path for ZMQ lookup socket."""
+    dp_rank = get_mooncake_dp_engine_index(vllm_config.parallel_config)
+    base_url = envs.VLLM_RPC_BASE_PATH
+    rpc_port = 0
+    assert vllm_config.kv_transfer_config is not None
+    extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
+    if "lookup_rpc_port" in extra_config:
+        rpc_port = extra_config["lookup_rpc_port"]
+    uid = os.getuid()
+    logger.debug("Base URL: %s, RPC Port: %s, UID: %s", base_url, rpc_port, uid)
+    return f"ipc://{base_url}/lookup_rpc_port_{rpc_port}_uid{uid}_dp_rank{dp_rank}"

vllm/model_executor/layers/attention/mla_attention.py

@@ -1367,6 +1367,7 @@ def backend_supports_prefill_query_quantization() -> bool:
     return backend_cls.get_name() in (
         "FLASHINFER",
         "TRTLLM_RAGGED",
+        "TOKENSPEED_MLA",
     )
 
 

vllm/model_executor/layers/fused_moe/config.py

@@ -598,6 +598,7 @@ def fp8_w8a8_moe_quant_config(
     a2_gscale: torch.Tensor | None = None,
     g1_alphas: torch.Tensor | None = None,
     g2_alphas: torch.Tensor | None = None,
+    gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for fp8 activations and fp8 weights.
@@ -617,6 +618,7 @@ def fp8_w8a8_moe_quant_config(
         per_act_token_quant=per_act_token_quant,
         per_out_ch_quant=per_out_ch_quant,
         block_shape=block_shape,
+        gemm1_clamp_limit=gemm1_clamp_limit,
     )
 
 
@@ -743,6 +745,7 @@ def mxfp4_w4a8_moe_quant_config(
     w1_bias: torch.Tensor | None = None,
     w2_bias: torch.Tensor | None = None,
     block_shape: list[int] | None = None,
+    gemm1_clamp_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for fp8 activations and mxfp4 weights.
@@ -752,6 +755,7 @@ def mxfp4_w4a8_moe_quant_config(
         _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None),
         _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
         _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
+        gemm1_clamp_limit=gemm1_clamp_limit,
     )
 
 

vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py

@@ -128,6 +128,8 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
         assert not quant_config.per_act_token_quant
         assert not quant_config.per_out_ch_quant
 
+        self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit
+
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
@@ -209,6 +211,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
                     input=input,
                     output_q=output,
                     group_size=block_k,
+                    clamp_limit=self.gemm1_clamp_limit,
                 )
             act_out = torch.empty(
                 (M_sum, activation_out_dim), dtype=input.dtype, device=input.device
@@ -228,6 +231,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
                 input=input,
                 output=output,
                 use_ue8m0=use_ue8m0,
+                clamp_limit=self.gemm1_clamp_limit,
             )
 
         # 3. fallback path for non-SiLU activations in non‑UE8M0 cases.
@@ -437,6 +441,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
                 input=input,
                 output=output,
                 use_ue8m0=use_ue8m0,
+                clamp_limit=self.gemm1_clamp_limit,
             )
 
         act_out = torch.empty(

vllm/model_executor/layers/fused_moe/experts/marlin_moe.py

@@ -413,6 +413,7 @@ def batched_fused_marlin_moe(
     is_k_full: bool = True,
     output: torch.Tensor | None = None,
     inplace: bool = False,
+    clamp_limit: float | None = None,
 ) -> torch.Tensor:
     """
     This function massages the inputs so the batched hidden_states can be
@@ -535,6 +536,7 @@ def batched_fused_marlin_moe(
         intermediate_cache2=intermediate_cache2,
         output=output.view(-1, K) if output is not None else output,
         is_k_full=is_k_full,
+        clamp_limit=clamp_limit,
     )
 
     output = output.view(B, BATCH_TOKENS_MAX, K)
@@ -768,6 +770,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
                 sort_indices2=self.w2_g_idx_sort_indices,
                 is_k_full=self.is_k_full,
                 input_dtype=self.input_dtype,
+                clamp_limit=self.gemm1_clamp_limit,
             )
             return
 
@@ -970,4 +973,5 @@ class BatchedMarlinExperts(MarlinExpertsBase):
             sort_indices1=self.w13_g_idx_sort_indices,
             sort_indices2=self.w2_g_idx_sort_indices,
             is_k_full=self.is_k_full,
+            clamp_limit=self.gemm1_clamp_limit,
         )

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -513,6 +513,7 @@ def make_fp8_moe_quant_config(
     block_shape: list[int] | None = None,
     per_act_token_quant: bool = False,
     per_out_ch_quant: bool = False,
+    swiglu_limit: float | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Create FusedMoEQuantConfig for the specified FP8 Backend.
@@ -556,6 +557,7 @@ def make_fp8_moe_quant_config(
             a2_gscale=(1.0 / a2_scale),
             g1_alphas=(w1_scale * a1_scale).squeeze(),
             g2_alphas=(w2_scale * a2_scale).squeeze(),
+            gemm1_clamp_limit=swiglu_limit,
         )
     # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to
     # _mxfp8_e4m3_quantize rather than standard FP8 block quantization.
@@ -570,6 +572,7 @@ def make_fp8_moe_quant_config(
             a2_scale=a2_scale,
             block_shape=block_shape,
             is_scale_swizzled=False,
+            gemm1_clamp_limit=swiglu_limit,
         )
 
     # All other backends use normal config.
@@ -581,6 +584,7 @@ def make_fp8_moe_quant_config(
         block_shape=block_shape,
         per_act_token_quant=per_act_token_quant,
         per_out_ch_quant=per_out_ch_quant,
+        gemm1_clamp_limit=swiglu_limit,
     )
 
 

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -1433,6 +1433,7 @@ def make_mxfp4_moe_quant_config(
             w1_bias=w1_bias,
             w2_bias=w2_bias,
             block_shape=None,
+            gemm1_clamp_limit=swiglu_limit,
         )
     elif mxfp4_backend in (
         Mxfp4MoeBackend.MARLIN,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py

@@ -360,6 +360,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
             per_act_token_quant=is_per_token,
             per_out_ch_quant=is_per_token,
             block_shape=self.weight_block_size,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
     def apply_monolithic(

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py

@@ -152,6 +152,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
             a1_scale=layer.w13_input_scale,
             a2_scale=layer.w2_input_scale,
             block_shape=self.weight_block_size,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
     def maybe_make_prepare_finalize(

vllm/model_executor/layers/quantization/fp8.py

@@ -851,6 +851,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             a1_scale=a1_scale,
             a2_scale=a2_scale,
             block_shape=self.weight_block_size,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
         # Inject biases into the quant config if the model has them

vllm/model_executor/layers/quantization/modelopt.py

@@ -948,6 +948,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
             w2_scale=w2_scale,
             a1_scale=a1_scale,
             a2_scale=a2_scale,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
     def apply_monolithic(

vllm/model_executor/layers/quantization/online/fp8.py

@@ -371,6 +371,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
             a1_scale=a1_scale,
             a2_scale=a2_scale,
             block_shape=self.weight_block_size,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
         self._maybe_inject_biases(quant_config, layer)

vllm/model_executor/layers/quantization/online/mxfp8.py

@@ -222,6 +222,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
             a1_scale=a1_scale,
             a2_scale=a2_scale,
             block_shape=self.weight_block_size,
+            swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
         self._maybe_inject_biases(quant_config, layer)

vllm/model_executor/layers/quantization/quark/quark_moe.py

@@ -433,6 +433,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
             w2_bias=layer.w2_bias,
             per_act_token_quant=self.input_qscheme == "per_channel",
             per_out_ch_quant=self.weight_qscheme == "per_channel",
+            gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
         )
 
     def apply(
@@ -899,6 +900,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
             w1_scale=layer.w13_weight_scale_2,
             w2_scale=layer.w2_weight_scale_2,
             per_out_ch_quant=True,
+            gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
         )
 
     def apply(

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -302,9 +302,11 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
     y_s_col_stride: tl.int64,
     # Information for float8
     eps,
+    clamp_limit,
     fp8_min: tl.constexpr,
     fp8_max: tl.constexpr,
     use_ue8m0: tl.constexpr,
+    HAS_CLAMP: tl.constexpr,
     # Meta-parameters
     GROUP_SIZE: tl.constexpr,
     BLOCK_M: tl.constexpr,
@@ -336,7 +338,16 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(
     act_in = tl.load(act_in_ptrs)
     mul_in = tl.load(act_in_ptrs + N_2)
 
-    # silu & mul
+    # silu & mul — match C++ silu_and_mul: clamp in fp32 then store back to the
+    # input dtype, run silu in fp32 then narrow, and do the mul at input
+    # precision so HAS_CLAMP True/False share the same multiplication path.
+    if HAS_CLAMP:
+        act_in = tl.minimum(act_in.to(tl.float32), clamp_limit).to(
+            y_ptr.dtype.element_ty
+        )
+        mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to(
+            y_ptr.dtype.element_ty
+        )
     act_in = act_in.to(tl.float32)
     one_f32 = tl.cast(1, tl.float32)
     silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty)
@@ -367,6 +378,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
     output: torch.Tensor | None = None,  # [M, N // 2]
     use_ue8m0: bool | None = None,
     eps: float = 1e-10,
+    clamp_limit: float | None = None,
 ):
     """
     silu+mul + block-fp8 quant with group size 128.
@@ -409,6 +421,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
     assert N_2 % BLOCK_N == 0
     grid = (M // BLOCK_M, N_2 // BLOCK_N)
 
+    has_clamp = clamp_limit is not None
     _silu_mul_per_token_group_quant_fp8_colmajor[grid](
         input,
         output,
@@ -417,9 +430,11 @@ def silu_mul_per_token_group_quant_fp8_colmajor(
         N,
         output_scales.stride(-1),
         eps,
+        clamp_limit if has_clamp else 0.0,
         fp8_min,
         fp8_max,
         use_ue8m0,
+        has_clamp,
         GROUP_SIZE,
         BLOCK_M,
         BLOCK_N,

vllm/model_executor/models/deepseek_v4.py

@@ -1203,10 +1203,10 @@ class DeepseekV4DecoderLayer(nn.Module):
         x: torch.Tensor,
         positions: torch.Tensor,
         input_ids: torch.Tensor | None,
-        post_mix: torch.Tensor | None,
+        post_mix: torch.Tensor | None = None,
-        res_mix: torch.Tensor | None,
+        res_mix: torch.Tensor | None = None,
-        residual: torch.Tensor | None,
+        residual: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         if residual is None:
             # Run standalone hc_pre on first layer
             residual = x

vllm/model_executor/models/deepseek_v4_mtp.py

@@ -141,9 +141,12 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(
             inputs_embeds
         ).unsqueeze(-2)
-        hidden_states = self.mtp_block(
+        hidden_states, residual, post_mix, res_mix = self.mtp_block(
             positions=positions, x=hidden_states, input_ids=None
         )
+        hidden_states = self.mtp_block.hc_post(
+            hidden_states, residual, post_mix, res_mix
+        )
         # Return the flat pre-hc_head residual so it can be re-fed as the
         # next spec step's `previous_hidden_states` when
         # num_speculative_tokens > 1. hc_head is deferred to compute_logits.

vllm/platforms/cuda.py

@@ -110,6 +110,10 @@ def _get_backend_priorities(
 
             return [
                 AttentionBackendEnum.FLASHINFER_MLA,
+                # R1 dims + FP8 KV only; rejected by supports_combination
+                # otherwise. Behind FLASHINFER_MLA: wins past bs≈8, regresses
+                # at bs≤2.
+                AttentionBackendEnum.TOKENSPEED_MLA,
                 AttentionBackendEnum.CUTLASS_MLA,
                 AttentionBackendEnum.FLASH_ATTN_MLA,
                 AttentionBackendEnum.FLASHMLA,

vllm/v1/attention/backends/mla/prefill/registry.py

@@ -43,6 +43,10 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):
         "vllm.v1.attention.backends.mla.prefill.trtllm_ragged."
         "TrtllmRaggedPrefillBackend"
     )
+    TOKENSPEED_MLA = (
+        "vllm.v1.attention.backends.mla.prefill.tokenspeed_mla."
+        "TokenspeedMLAPrefillBackend"
+    )
 
     def get_path(self) -> str:
         """Get the fully qualified class path for this backend."""

vllm/v1/attention/backends/mla/prefill/selector.py

@@ -67,6 +67,7 @@ def _get_mla_prefill_backend_priorities(
             MLAPrefillBackendEnum.FLASH_ATTN,
             MLAPrefillBackendEnum.TRTLLM_RAGGED,
             MLAPrefillBackendEnum.FLASHINFER,
+            MLAPrefillBackendEnum.TOKENSPEED_MLA,
         ]
     else:  # Hopper (SM90) and older
         return [

vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TokenSpeed CuTe DSL backend for MLA prefill."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.model_executor.layers.attention.mla_attention import (
+        MLACommonPrefillMetadata,
+    )
+    from vllm.platforms.interface import DeviceCapability
+
+
+class TokenspeedMLAPrefillBackend(MLAPrefillBackend):
+    """TokenSpeed CuTe DSL backend for MLA prefill."""
+
+    requires_r1_mla_dimensions = True
+
+    @staticmethod
+    def get_name() -> str:
+        return "TOKENSPEED_MLA"
+
+    @classmethod
+    def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:
+        return device_capability.major == 10
+
+    _INSTALL_HINT = (
+        "tokenspeed_mla package is not installed. "
+        "Install it with: `uv pip install tokenspeed-mla`"
+    )
+
+    @classmethod
+    def is_available(cls) -> bool:
+        try:
+            from tokenspeed_mla import (
+                tokenspeed_mla_prefill,  # noqa: F401
+            )
+
+            return True
+        except ImportError:
+            return False
+
+    @classmethod
+    def validate_configuration(
+        cls,
+        device_capability,
+        selector_config,
+    ) -> list[str]:
+        # Replace the generic "required dependencies not available" message
+        # from the base class with a specific install hint so users know
+        # exactly which package to install when they explicitly select this
+        # backend without having tokenspeed_mla installed.
+        reasons = super().validate_configuration(device_capability, selector_config)
+        return [
+            cls._INSTALL_HINT if r == "required dependencies not available" else r
+            for r in reasons
+        ]
+
+    def __init__(
+        self,
+        num_heads: int,
+        scale: float,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        vllm_config: "VllmConfig",
+    ) -> None:
+        super().__init__(
+            num_heads=num_heads,
+            scale=scale,
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            vllm_config=vllm_config,
+        )
+
+        # Pre-JIT BF16 and FP8 prefill kernels. Idempotent — also called from
+        # TokenspeedMLAImpl.__init__; second call is a no-op.
+        from tokenspeed_mla import warmup_compile_prefill
+
+        for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
+            warmup_compile_prefill(
+                q_dtype=q_dtype,
+                d_qk=qk_nope_head_dim + qk_rope_head_dim,
+                d_v=v_head_dim,
+                enable_pdl=False,
+            )
+
+    def prepare_metadata(
+        self,
+        prefill_metadata: "MLACommonPrefillMetadata",
+    ) -> None:
+        super().prepare_metadata(prefill_metadata)
+        # Kernel signature requires `seq_lens` but the implementation never reads
+        # it (per-batch lengths are derived from `cum_seq_lens` diffs); compute
+        # for parity with trtllm_ragged. cuda-graph padding in
+        # `query_start_loc` is saturated to `total_num_tokens`
+        # (gpu_model_runner.py:1905), so trailing diffs are 0 and padded batches
+        # are kernel no-ops — same reason trtllm passes the padded length as
+        # batch_size directly.
+        self._query_seq_lens = (
+            prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]
+        )
+
+    def run_prefill_new_tokens(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        return_softmax_lse: bool,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        from tokenspeed_mla import tokenspeed_mla_prefill
+
+        # `v` arrives as the second half of `kv_nope.split(...)` in
+        # mla_attention.forward_mha — a non-contiguous view of `kv_nope` along
+        # dim=-1. The kernel does `v.reshape(1, total_kv, h_k, 1, d_v)` which
+        # would silently copy on a non-contiguous tensor; force contiguity here
+        # so the copy (if any) happens once outside the kernel call.
+        v = v.contiguous()
+
+        ret = tokenspeed_mla_prefill(
+            query=q,
+            key=k,
+            value=v,
+            seq_lens=self._query_seq_lens,
+            cum_seq_lens=self._prefill_metadata.query_start_loc,
+            max_seq_len=self._prefill_metadata.max_query_len,
+            batch_size=self._query_seq_lens.shape[0],
+            softmax_scale=self.scale,
+            is_causal=True,
+            return_lse=return_softmax_lse,
+            enable_pdl=False,
+        )
+
+        if isinstance(ret, tuple):
+            # Convert from (q_len, num_heads) to (num_heads, q_len)
+            return ret[0], ret[1].transpose(0, 1).contiguous()
+        return ret
+
+    def run_prefill_context_chunk(
+        self,
+        chunk_idx: int,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from tokenspeed_mla import tokenspeed_mla_prefill
+
+        assert self._prefill_metadata.chunked_context is not None
+        chunked = self._prefill_metadata.chunked_context
+
+        # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope`
+        # in `_compute_prefill_context` and arrives non-contiguous.
+        v = v.contiguous()
+
+        attn_out, lse = tokenspeed_mla_prefill(
+            query=q,
+            key=k,
+            value=v,
+            seq_lens=chunked.seq_lens[chunk_idx],
+            cum_seq_lens=chunked.cu_seq_lens[chunk_idx],
+            max_seq_len=chunked.max_seq_lens[chunk_idx],
+            batch_size=chunked.seq_lens[chunk_idx].shape[0],
+            softmax_scale=self.scale,
+            is_causal=False,
+            return_lse=True,
+            cum_seq_lens_q=self._prefill_metadata.query_start_loc,
+            max_seq_len_q=self._prefill_metadata.max_query_len,
+            enable_pdl=False,
+        )
+
+        # Convert from (q_len, num_heads) to (num_heads, q_len)
+        return attn_out, lse.transpose(0, 1).contiguous()

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TokenSpeed CuTe DSL MLA decode backend (Blackwell, FP8 KV cache only)."""
+
+from typing import ClassVar
+
+import torch
+
+from vllm.config.cache import CacheDType
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import (
+    MLACommonBackend,
+    MLACommonImpl,
+    MLACommonMetadata,
+    MLACommonMetadataBuilder,
+    QueryLenSupport,
+)
+from vllm.platforms.interface import DeviceCapability
+from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.v1.attention.backend import (
+    AttentionCGSupport,
+    AttentionLayer,
+    AttentionType,
+    MultipleOf,
+)
+from vllm.v1.attention.backends.utils import KVCacheLayoutType
+
+logger = init_logger(__name__)
+
+# Workspace upper bound for tokenspeed_mla_decode (per-device, lazy):
+#   num_sms * num_heads * MAX_Q_LEN * (kv_lora_rank + 1) * sizeof(float32)
+# Matches the kernel's `get_workspace_size` formula. MAX_Q_LEN=8 covers up to
+# EAGLE3 / MTP-2 spec decoding query lengths; larger q_len fails the kernel's
+# own buffer check.
+_TOKENSPEED_MAX_Q_LEN = 8
+
+_g_workspace: dict[torch.device, torch.Tensor] = {}
+
+
+def _get_workspace(
+    device: torch.device, num_heads: int, kv_lora_rank: int
+) -> torch.Tensor:
+    from tokenspeed_mla import get_num_sm
+
+    needed = (
+        get_num_sm(device) * num_heads * _TOKENSPEED_MAX_Q_LEN * (kv_lora_rank + 1) * 4
+    )
+    existing = _g_workspace.get(device)
+    if existing is None or existing.numel() < needed:
+        _g_workspace[device] = torch.empty(needed, dtype=torch.int8, device=device)
+    return _g_workspace[device]
+
+
+class TokenspeedMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
+
+
+class TokenspeedMLABackend(MLACommonBackend):
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [32, 64]
+
+    @staticmethod
+    def get_name() -> str:
+        return "TOKENSPEED_MLA"
+
+    @staticmethod
+    def get_impl_cls() -> type["TokenspeedMLAImpl"]:
+        return TokenspeedMLAImpl
+
+    @staticmethod
+    def get_builder_cls() -> type["TokenspeedMLAMetadataBuilder"]:
+        return TokenspeedMLAMetadataBuilder
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 10
+
+    @classmethod
+    def supports_combination(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        device_capability: DeviceCapability,
+    ) -> str | None:
+        # Surface a clear install hint up front rather than letting a raw
+        # ModuleNotFoundError fire deep inside `forward_mqa` at first request.
+        try:
+            import tokenspeed_mla  # noqa: F401
+        except ImportError:
+            return (
+                "tokenspeed_mla package is not installed. "
+                "Install it with: `uv pip install tokenspeed-mla`"
+            )
+
+        # tokenspeed_mla CuTe DSL kernel is shape-specialized for DeepSeek R1
+        # MLA dimensions (qk_nope=128, qk_rope=64, v=128). Reject anything else.
+        from vllm.config import get_current_vllm_config
+
+        vllm_config = get_current_vllm_config()
+        if vllm_config.model_config is not None:
+            hf_text_config = vllm_config.model_config.hf_text_config
+            qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 0)
+            qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 0)
+            v_head_dim = getattr(hf_text_config, "v_head_dim", 0)
+            if qk_nope_head_dim != 128 or qk_rope_head_dim != 64 or v_head_dim != 128:
+                return (
+                    "tokenspeed_mla requires DeepSeek R1 MLA dimensions "
+                    "(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), "
+                    f"got ({qk_nope_head_dim}, {qk_rope_head_dim}, {v_head_dim})"
+                )
+        return None
+
+    @classmethod
+    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
+        return "HND"
+
+
+class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA Specific Arguments
+        **mla_args,
+    ) -> None:
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            **mla_args,
+        )
+
+        unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
+        if any(unsupported_features):
+            raise NotImplementedError(
+                "TokenspeedMLAImpl does not support one of the following: "
+                "alibi_slopes, sliding_window, logits_soft_cap"
+            )
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError(
+                "Encoder self-attention and "
+                "encoder/decoder cross-attention "
+                "are not implemented for "
+                "TokenspeedMLAImpl"
+            )
+
+        if not is_quantized_kv_cache(self.kv_cache_dtype):
+            raise NotImplementedError(
+                "TokenspeedMLAImpl requires an FP8 KV cache "
+                "(--kv-cache-dtype fp8 or fp8_e4m3); "
+                f"got kv_cache_dtype={self.kv_cache_dtype!r}."
+            )
+
+        # Allocate (or fetch the cached) workspace lazily on first forward —
+        # __init__ runs before the device is necessarily set on the worker;
+        # we know it for sure at forward time when we see the input tensor.
+        self._workspace_buffer: torch.Tensor | None = None
+        self.softmax_scale: float | None = None
+        self.output_scale: float | None = None
+
+        # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always
+        # runs when tokenspeed is selected, prefill backend may not (user can
+        # pair with flash_attn / trtllm). Idempotent.
+        from tokenspeed_mla import warmup_compile_prefill
+
+        for q_dtype in (torch.bfloat16, torch.float8_e4m3fn):
+            warmup_compile_prefill(
+                q_dtype=q_dtype,
+                d_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
+                d_v=self.v_head_dim,
+                enable_pdl=False,
+            )
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        from tokenspeed_mla import tokenspeed_mla_decode
+
+        assert kv_c_and_k_pe_cache.numel() > 0
+        assert attn_metadata.decode is not None
+
+        if isinstance(q, tuple):
+            q_nope, q_pe = q
+            q = torch.cat([q_nope, q_pe], dim=-1)
+
+        # supports_quant_query_input=True (set in MLACommonImpl) tells the
+        # pipeline to concat+FP8-quantize Q upstream via _decode_concat_quant_fp8_op.
+        # The kernel is shape-specialized for FP8 Q + FP8 KV, so anything else
+        # here means the upstream quant didn't run and the kernel will produce
+        # garbage.
+        assert q.dtype == torch.float8_e4m3fn, (
+            f"TokenspeedMLAImpl expected FP8 query (supports_quant_query_input=True), "
+            f"got {q.dtype}. Pipeline isinstance(q, tuple)={isinstance(q, tuple)}, "
+            f"q_scale={layer._q_scale_float}, k_scale={layer._k_scale_float}."
+        )
+
+        # tokenspeed_mla_decode expects query shape
+        # (num_decodes, q_len_per_request, num_heads, head_dim).
+        if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
+            logger.warning_once(
+                """TokenspeedMLAImpl got a query of uneven length.
+                This usually indicates an issue in batch reordering
+                or incorrect setup in dummy_run."""
+            )
+            q = q.unsqueeze(1)
+        else:
+            q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
+
+        if self.softmax_scale is None:
+            # FP8 KV cache is mandatory for this backend, so q_scale/k_scale
+            # always apply. softmax_scale is bmm1; output_scale is bmm2 — both
+            # required to recover the correct attention output from the FP8
+            # KV cache (V is stored as V_real/k_scale).
+            self.softmax_scale = (
+                self.scale * layer._q_scale_float * layer._k_scale_float
+            )
+            self.output_scale = layer._k_scale_float
+
+        if self._workspace_buffer is None:
+            self._workspace_buffer = _get_workspace(
+                q.device, self.num_heads, self.kv_lora_rank
+            )
+
+        # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size).
+        # tokenspeed_mla_decode wants 3D — pass as-is (no unsqueeze, unlike trtllm).
+        o = tokenspeed_mla_decode(
+            query=q,
+            kv_cache=kv_c_and_k_pe_cache,
+            workspace_buffer=self._workspace_buffer,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            block_tables=attn_metadata.decode.block_table,
+            seq_lens=attn_metadata.decode.seq_lens,
+            max_seq_len=attn_metadata.max_seq_len,
+            softmax_scale=self.softmax_scale,
+            output_scale=self.output_scale,
+            enable_pdl=False,
+        )
+
+        # Flatten the output for consistent shape
+        o = o.view(-1, o.shape[-2], o.shape[-1])
+
+        # tokenspeed_mla_decode does not return LSE.
+        return o, None

vllm/v1/attention/backends/registry.py

@@ -63,6 +63,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     FLASHINFER_MLA = (
         "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"
     )
+    TOKENSPEED_MLA = (
+        "vllm.v1.attention.backends.mla.tokenspeed_mla.TokenspeedMLABackend"
+    )
     FLASHINFER_MLA_SPARSE = (
         "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
         "FlashInferMLASparseBackend"

vllm/v1/kv_cache_interface.py

@@ -422,6 +422,17 @@ class SlidingWindowSpec(AttentionSpec):
 
     @property
     def real_page_size_bytes(self) -> int:
+        # Mirror ``FullAttentionSpec.real_page_size_bytes`` for NVFP4 KV cache.
+        if self.kv_quant_mode.is_nvfp4:
+            last_dim = nvfp4_kv_cache_full_dim(
+                self.head_size
+            ) + nvfp4_kv_cache_full_dim(self.head_size_v)
+            return (
+                self.block_size
+                * self.num_kv_heads
+                * last_dim
+                * get_dtype_size(self.dtype)
+            )
         return (
             self.block_size
             * self.num_kv_heads