[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl

The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
update version
2026-05-06 07:45:16 +00:00 · 2026-05-06 07:28:20 +00:00 · 2026-05-06 07:20:48 +00:00 · 2026-05-06 03:14:49 +00:00 · 2026-05-06 02:53:46 +00:00 · 2026-05-06 02:49:09 +00:00
38 changed files with 1309 additions and 288 deletions
@@ -309,6 +309,7 @@ steps:
        depends_on: ~

      - label: "Build release image - x86_64 - CPU"
+        key: build-cpu-release-image-x86
        depends_on:
          - block-cpu-release-image-build
          - input-release-version
@@ -327,7 +328,8 @@ steps:
        depends_on: ~

      - label: "Build release image - arm64 - CPU"
-        depends_on: 
+        key: build-cpu-release-image-arm64
+        depends_on:
          - block-arm64-cpu-release-image-build
          - input-release-version
        agents:
@@ -436,6 +438,41 @@ 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:
@@ -8,8 +8,6 @@ if [ -z "${RELEASE_VERSION}" ]; then
  RELEASE_VERSION="1.0.0.dev"
 fi

-ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
-
 buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
 To download the wheel (by commit):
 \`\`\`
@@ -25,95 +23,5 @@ aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-
 aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
 \`\`\`

-
-To download and upload the image:
-
-\`\`\`
-# Download images:
-
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
-docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
-docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
-docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
-
-# Tag and push images:
-
-## CUDA
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
-docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
-docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai:latest-x86_64
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
-docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
-docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
-docker push vllm/vllm-openai:latest-x86_64-cu129
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
-docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
-docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-docker push vllm/vllm-openai:latest-aarch64
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
-docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
-docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-docker push vllm/vllm-openai:latest-aarch64-cu129
-docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-
-## ROCm
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
-docker push vllm/vllm-openai-rocm:latest
-docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
-docker push vllm/vllm-openai-rocm:latest-base
-docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
-
-## CPU
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
-docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
-docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
-docker push vllm/vllm-openai-cpu:latest-x86_64
-docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
-
-docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
-docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
-docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-docker push vllm/vllm-openai-cpu:latest-arm64
-docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-
-# Create multi-arch manifest:
-
-docker manifest rm vllm/vllm-openai:latest
-docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-docker manifest push vllm/vllm-openai:latest
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
-
-docker manifest rm vllm/vllm-openai:latest-cu129
-docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
-docker manifest push vllm/vllm-openai:latest-cu129
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
-
-docker manifest rm vllm/vllm-openai-cpu:latest || true
-docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
-docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
-docker manifest push vllm/vllm-openai-cpu:latest
-docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
-\`\`\`
+Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
 EOF
@@ -0,0 +1,180 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Publish release Docker images from ECR to DockerHub.
+# Pulls per-arch images, tags with latest and versioned tags, pushes them,
+# then creates and pushes multi-arch manifests.
+
+set -euo pipefail
+
+RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
+if [ -z "${RELEASE_VERSION}" ]; then
+  echo "ERROR: release-version metadata not set"
+  exit 1
+fi
+
+COMMIT="$BUILDKITE_COMMIT"
+ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
+
+echo "========================================"
+echo "Publishing release images v${RELEASE_VERSION}"
+echo "  Commit: ${COMMIT}"
+echo "  ROCm base cache key: ${ROCM_BASE_CACHE_KEY}"
+echo "========================================"
+
+# Login to ECR to pull staging images
+aws ecr-public get-login-password --region us-east-1 | \
+  docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
+
+# ---- CUDA (default: 13.0) ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
+docker push vllm/vllm-openai:latest-x86_64
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
+docker push vllm/vllm-openai:latest-aarch64
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
+
+docker manifest rm vllm/vllm-openai:latest || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true
+docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
+docker manifest push vllm/vllm-openai:latest
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
+
+# ---- CUDA 12.9 ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
+docker push vllm/vllm-openai:latest-x86_64-cu129
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker push vllm/vllm-openai:latest-aarch64-cu129
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+
+docker manifest rm vllm/vllm-openai:latest-cu129 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true
+docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
+docker manifest push vllm/vllm-openai:latest-cu129
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
+
+# ---- Ubuntu 24.04 (CUDA 13.0) ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
+docker push vllm/vllm-openai:latest-x86_64-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
+docker push vllm/vllm-openai:latest-aarch64-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
+
+docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true
+docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
+docker manifest push vllm/vllm-openai:latest-ubuntu2404
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404
+
+# ---- Ubuntu 24.04 (CUDA 12.9) ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
+docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
+docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
+docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
+
+docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true
+docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true
+docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
+docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404
+
+# ---- ROCm ----
+
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm
+docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}
+docker push vllm/vllm-openai-rocm:latest
+docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
+docker push vllm/vllm-openai-rocm:latest-base
+docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
+
+# ---- CPU ----
+# CPU images are behind separate block steps and may not have been built.
+# All-or-nothing: inspect both arches first, then either publish everything
+# (per-arch + multi-arch manifest) or skip everything. Publishing only one
+# arch would leave `:latest-x86_64` pointing at the new release while the
+# `:latest` multi-arch manifest still resolves to the previous release.
+
+CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
+CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
+
+CPU_X86_AVAILABLE=false
+CPU_ARM_AVAILABLE=false
+docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true
+docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true
+
+if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then
+  docker pull "${CPU_X86_TAG}"
+  docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64
+  docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
+  docker push vllm/vllm-openai-cpu:latest-x86_64
+  docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
+
+  docker pull "${CPU_ARM_TAG}"
+  docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64
+  docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
+  docker push vllm/vllm-openai-cpu:latest-arm64
+  docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
+
+  docker manifest rm vllm/vllm-openai-cpu:latest || true
+  docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true
+  docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
+  docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
+  docker manifest push vllm/vllm-openai-cpu:latest
+  docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
+elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then
+  echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)"
+else
+  # Partial state: one arch built, the other did not. Fail loudly rather than
+  # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch)
+  # disagree on which release they point at.
+  echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})."
+  echo "       Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags."
+  echo "       Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended."
+  exit 1
+fi
+
+echo ""
+echo "Successfully published release images for v${RELEASE_VERSION}"
@@ -51,6 +51,7 @@ vllm serve "$MODEL" \
  --offload-num-in-group 2 \
  --offload-prefetch-step 1 \
  --offload-params w13_weight w2_weight \
+  --generation-config vllm \
  --port "$PORT" \
  ${EXTRA_ARGS+"${EXTRA_ARGS[@]}"} &
 SERVER_PID=$!
@@ -39,10 +39,11 @@ fi

 set -x # avoid printing secrets above

-# install twine from pypi
+# install twine and sdist build prerequisites from pypi
 python3 -m venv /tmp/vllm-release-env
 source /tmp/vllm-release-env/bin/activate
 pip install twine
+pip install -r requirements/build/cuda.txt
 python3 -m twine --version

 # copy release wheels to local directory
@@ -17,7 +17,7 @@ steps:
 - label: V1 attention (B200)
  key: v1-attention-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
    - vllm/config/attention.py
    - vllm/model_executor/layers/attention
@@ -14,7 +14,7 @@ steps:

 - label: Attention Benchmarks Smoke Test (B200)
  key: attention-benchmarks-smoke-test-b200
-  device: b200
+  device: b200-k8s
  num_gpus: 2
  optional: true
  working_dir: "/vllm-workspace/"
@@ -43,7 +43,7 @@ steps:
  key: asynctp-correctness-tests-b200
  timeout_in_minutes: 50
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  commands:
@@ -68,7 +68,7 @@ steps:
  key: fusion-and-compile-unit-tests-2xb200
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
  - csrc/quantization/fp4/
  - vllm/model_executor/layers/quantization/
@@ -137,7 +137,7 @@ steps:
  key: fusion-e2e-config-sweep-b200
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  num_devices: 1
  optional: true
  commands:
@@ -209,7 +209,7 @@ steps:
  key: fusion-e2e-tp2-b200
  timeout_in_minutes: 20
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  num_devices: 2
  source_file_dependencies:
    - csrc/quantization/
@@ -212,7 +212,7 @@ steps:

 - label: Distributed Tests (2 GPUs)(B200)
  key: distributed-tests-2-gpus-b200
-  device: b200
+  device: b200-k8s
  optional: true
  working_dir: "/vllm-workspace/"
  num_devices: 2
@@ -25,7 +25,7 @@ steps:
 - label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
  key: qwen3-30b-a3b-fp8-block-accuracy-b200
  timeout_in_minutes: 60
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  working_dir: "/vllm-workspace"
@@ -125,7 +125,7 @@ steps:
  key: kernels-b200
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  # optional: true
  source_file_dependencies:
  - csrc/quantization/fp4/
@@ -212,7 +212,7 @@ steps:
 - label: Kernels Fp4 MoE Test (B200)
  key: kernels-fp4-moe-test-b200
  timeout_in_minutes: 60
-  device: b200
+  device: b200-k8s
  num_devices: 1
  optional: true
  commands:
@@ -51,7 +51,7 @@ steps:
 - label: LM Eval Qwen3.5 Models (B200)
  key: lm-eval-qwen3-5-models-b200
  timeout_in_minutes: 120
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  source_file_dependencies:
@@ -84,7 +84,7 @@ steps:
  
 - label: MoE Refactor Integration Test (B200 - TEMPORARY)
  key: moe-refactor-integration-test-b200-temporary
-  device: b200
+  device: b200-k8s
  optional: true
  num_devices: 2
  commands:
@@ -224,7 +224,7 @@ steps:
 - label: Batch Invariance (B200)
  key: batch-invariance-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
    - vllm/v1/attention
    - vllm/model_executor/layers
@@ -25,7 +25,7 @@ steps:
  key: quantized-moe-test-b200
  timeout_in_minutes: 60
  working_dir: "/vllm-workspace/"
-  device: b200
+  device: b200-k8s
  source_file_dependencies:
  - tests/quantization/test_blackwell_moe.py
  - vllm/model_executor/models/deepseek_v2.py
@@ -75,7 +75,7 @@ steps:
 - label: Spec Decode Draft Model Nightly B200
  key: spec-decode-draft-model-nightly-b200
  timeout_in_minutes: 30
-  device: b200
+  device: b200-k8s
  optional: true
  source_file_dependencies:
    - vllm/v1/spec_decode/
@@ -203,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.
@@ -223,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 |
@@ -49,7 +49,7 @@ ijson # Required for mistral streaming tool parser
 setproctitle # Used to set process names for better debugging and monitoring
 openai-harmony >= 0.0.3  # Required for gpt-oss
 anthropic >= 0.71.0
-model-hosting-container-standards >= 0.1.13, < 1.0.0
+model-hosting-container-standards >= 0.1.14, < 1.0.0
 mcp
 opentelemetry-sdk >= 1.27.0
 opentelemetry-api >= 1.27.0
@@ -23,3 +23,6 @@ fastsafetensors >= 0.2.2
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
 nvidia-cutlass-dsl>=4.4.2
 quack-kernels>=0.3.3
+
+# Tokenspeed_MLA for faster mla with spec decode
+tokenspeed-mla==0.1.1
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Parity: tokenspeed_mla_decode vs flashinfer trtllm_batch_decode_with_kv_cache_mla."""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.has_device_capability(100):
+    pytest.skip(
+        reason="tokenspeed_mla / TRT-LLM MLA decode require Blackwell (SM100+).",
+        allow_module_level=True,
+    )
+
+try:
+    from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+except ImportError:
+    pytest.skip(reason="flashinfer not installed", allow_module_level=True)
+
+try:
+    from tokenspeed_mla import get_num_sm, tokenspeed_mla_decode
+except ImportError:
+    pytest.skip(reason="tokenspeed_mla not installed", allow_module_level=True)
+
+
+FLASHINFER_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
+_TS_MAX_Q_LEN = 8
+
+
+def _ts_workspace(device, num_heads, kv_lora_rank):
+    needed = get_num_sm(device) * num_heads * _TS_MAX_Q_LEN * (kv_lora_rank + 1) * 4
+    return torch.empty(needed, dtype=torch.int8, device=device)
+
+
+@pytest.mark.parametrize("bs", [1, 2, 4, 16])
+@pytest.mark.parametrize("block_size", [32, 64])
+@pytest.mark.parametrize("q_len_per_request", [1, 2, 4])
+def test_tokenspeed_vs_trtllm_decode(bs: int, block_size: int, q_len_per_request: int):
+    """Match tokenspeed_mla_decode against TRT-LLM batch decode MLA.
+
+    Both kernels consume the same FP8 KV cache, paged block table, and
+    seq_lens. The only structural difference is rank: TRT-LLM expects 4D
+    (`unsqueeze(1)` for the kv-head dim) while tokenspeed expects 3D. We
+    pass each kernel its preferred shape from the same underlying tensor.
+    """
+    torch.set_default_device("cuda")
+    torch.manual_seed(42)
+
+    # Deepseek R1 dims — both kernels are R1-shape-specialized.
+    num_heads = 128
+    kv_lora_rank = 512
+    qk_nope_head_dim = 128
+    qk_rope_head_dim = 64
+    qk_head_dim = kv_lora_rank + qk_rope_head_dim
+    scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
+
+    MAX_SEQ_LEN = 1024
+
+    seq_lens = [torch.randint(2, MAX_SEQ_LEN, (1,)).item() for _ in range(bs)]
+    seq_lens[-1] = MAX_SEQ_LEN
+    max_seq_len = max(seq_lens)
+    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
+
+    blocks_per_seq = (seq_lens_tensor + block_size - 1) // block_size
+    max_num_blocks_per_seq = max(blocks_per_seq.max().item(), 4)
+    total_blocks_needed = sum(blocks_per_seq).item()
+    all_block_ids = torch.randperm(total_blocks_needed, dtype=torch.int32)
+
+    block_tables = torch.zeros((bs, max_num_blocks_per_seq), dtype=torch.int32)
+    block_id = 0
+    for i in range(bs):
+        n = blocks_per_seq[i].item()
+        block_tables[i, :n] = all_block_ids[block_id : block_id + n]
+        block_id += n
+
+    # KV cache: build in BF16 then cast once to FP8 so both kernels see the
+    # exact same quantized values. Shape (num_blocks, block_size, qk_head_dim).
+    kv_cache_bf16 = torch.randn(
+        block_tables.numel(), block_size, qk_head_dim, dtype=torch.bfloat16
+    )
+    kv_cache = kv_cache_bf16.to(torch.float8_e4m3fn)
+
+    # Query: (bs, q_len_per_request, num_heads, qk_head_dim) — same layout as
+    # FlashInferMLAImpl.forward_mqa. Cast to FP8 to match KV.
+    q = torch.randn(
+        bs, q_len_per_request, num_heads, qk_head_dim, dtype=torch.bfloat16
+    ).to(torch.float8_e4m3fn)
+
+    # --- TRT-LLM reference ---
+    fi_workspace = torch.zeros(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8)
+    out_ref = trtllm_batch_decode_with_kv_cache_mla(
+        query=q,
+        kv_cache=kv_cache.unsqueeze(1),
+        workspace_buffer=fi_workspace,
+        qk_nope_head_dim=qk_nope_head_dim,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        block_tables=block_tables,
+        seq_lens=seq_lens_tensor,
+        max_seq_len=max_seq_len,
+        bmm1_scale=scale,
+    )
+
+    # --- TokenSpeed candidate ---
+    ts_workspace = _ts_workspace(q.device, num_heads, kv_lora_rank)
+    out_ts = tokenspeed_mla_decode(
+        query=q,
+        kv_cache=kv_cache,
+        workspace_buffer=ts_workspace,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        block_tables=block_tables,
+        seq_lens=seq_lens_tensor,
+        max_seq_len=max_seq_len,
+        softmax_scale=scale,
+    )
+
+    # Both kernels output v_head_dim=kv_lora_rank=512 per head.
+    # Output dtypes can differ; compare in float32.
+    out_ref_f = out_ref.to(torch.float32)
+    out_ts_f = out_ts.to(torch.float32)
+    assert out_ref_f.shape == out_ts_f.shape, (
+        f"shape mismatch: trtllm={tuple(out_ref_f.shape)} "
+        f"tokenspeed={tuple(out_ts_f.shape)}"
+    )
+
+    torch.testing.assert_close(out_ts_f, out_ref_f, atol=2e-2, rtol=2e-2)
@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Numeric accuracy parity: tokenspeed_mla_prefill vs trtllm_ragged_attention_deepseek.
+
+Two cases mirror what the vLLM MLA prefill backend does in production:
+- `test_prefill_no_context`: causal Q==KV ragged batch (run_prefill_new_tokens).
+- `test_prefill_with_context`: non-causal Q ragged + KV ragged with
+  per-request kv_len > q_len (run_prefill_context_chunk).
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.has_device_capability(100):
+    pytest.skip(
+        reason="tokenspeed_mla / TRT-LLM ragged require Blackwell (SM100+).",
+        allow_module_level=True,
+    )
+
+try:
+    from flashinfer.prefill import trtllm_ragged_attention_deepseek
+except ImportError:
+    pytest.skip(reason="flashinfer not installed", allow_module_level=True)
+
+try:
+    from tokenspeed_mla import tokenspeed_mla_prefill, warmup_compile_prefill
+except ImportError:
+    pytest.skip(reason="tokenspeed_mla not installed", allow_module_level=True)
+
+
+FLASHINFER_WORKSPACE_BUFFER_SIZE = 384 * 1024 * 1024
+
+
+# Deepseek R1 dimensions — both kernels are shape-specialized for these.
+NUM_HEADS = 128
+KV_LORA_RANK = 512
+QK_NOPE_HEAD_DIM = 128
+QK_ROPE_HEAD_DIM = 64
+V_HEAD_DIM = 128
+QK_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM  # 192
+SCALE = QK_HEAD_DIM**-0.5
+
+
+def _make_q_kv(
+    seq_lens: list[int],
+    kv_lens: list[int],
+    dtype: torch.dtype,
+):
+    """Build ragged Q (qk_head_dim) and K (qk_head_dim) / V (v_head_dim)."""
+    total_q = sum(seq_lens)
+    total_kv = sum(kv_lens)
+
+    q = torch.randn(total_q, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16).to(dtype)
+    k = torch.randn(total_kv, NUM_HEADS, QK_HEAD_DIM, dtype=torch.bfloat16).to(dtype)
+    v = torch.randn(total_kv, NUM_HEADS, V_HEAD_DIM, dtype=torch.bfloat16).to(dtype)
+    return q, k, v
+
+
+def _cumsum_int32(lens: list[int]) -> torch.Tensor:
+    out = torch.zeros(len(lens) + 1, dtype=torch.int32)
+    out[1:] = torch.tensor(lens, dtype=torch.int32).cumsum(0)
+    return out
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("bs", [1, 4, 16])
+@pytest.mark.parametrize("max_q_len", [64, 256, 1024])
+def test_prefill_no_context(dtype: torch.dtype, bs: int, max_q_len: int):
+    """Causal Q==KV ragged: matches the run_prefill_new_tokens code path."""
+    torch.set_default_device("cuda")
+    torch.manual_seed(0)
+
+    if dtype == torch.float8_e4m3fn:
+        warmup_compile_prefill(
+            q_dtype=torch.float8_e4m3fn,
+            d_qk=QK_HEAD_DIM,
+            d_v=V_HEAD_DIM,
+            enable_pdl=False,
+        )
+
+    seq_lens = [int(torch.randint(2, max_q_len + 1, (1,)).item()) for _ in range(bs)]
+    seq_lens[-1] = max_q_len  # pin the last so max_q_len is hit
+
+    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32)
+    cum_seq_lens = _cumsum_int32(seq_lens)
+
+    q, k, v = _make_q_kv(seq_lens, seq_lens, dtype)
+
+    # --- TRT-LLM reference ---
+    workspace = torch.zeros(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8)
+    out_ref = torch.empty(q.shape[0], q.shape[1], v.shape[2], dtype=torch.bfloat16)
+    ref_ret = trtllm_ragged_attention_deepseek(
+        query=q,
+        key=k,
+        value=v,
+        workspace_buffer=workspace,
+        seq_lens=seq_lens_tensor,
+        max_q_len=max_q_len,
+        max_kv_len=max_q_len,
+        bmm1_scale=SCALE,
+        bmm2_scale=1.0,
+        o_sf_scale=1.0,
+        batch_size=bs,
+        window_left=-1,
+        cum_seq_lens_q=cum_seq_lens,
+        cum_seq_lens_kv=cum_seq_lens,
+        enable_pdl=False,
+        is_causal=True,
+        return_lse=False,
+        out=out_ref,
+    )
+    out_ref = ref_ret if not isinstance(ref_ret, tuple) else ref_ret[0]
+
+    # --- TokenSpeed candidate ---
+    out_ts = tokenspeed_mla_prefill(
+        query=q,
+        key=k,
+        value=v,
+        seq_lens=seq_lens_tensor,
+        cum_seq_lens=cum_seq_lens,
+        max_seq_len=max_q_len,
+        batch_size=bs,
+        softmax_scale=SCALE,
+        is_causal=True,
+        return_lse=False,
+        enable_pdl=False,
+    )
+    if isinstance(out_ts, tuple):
+        out_ts = out_ts[0]
+
+    out_ref_f = out_ref.to(torch.float32)
+    out_ts_f = out_ts.to(torch.float32)
+    assert out_ref_f.shape == out_ts_f.shape, (
+        f"shape mismatch: trtllm={tuple(out_ref_f.shape)} "
+        f"tokenspeed={tuple(out_ts_f.shape)}"
+    )
+
+    if dtype == torch.float8_e4m3fn:
+        atol, rtol = 5e-2, 5e-2
+    else:
+        atol, rtol = 1e-2, 1e-2
+    torch.testing.assert_close(out_ts_f, out_ref_f, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("bs", [1, 4, 16])
+def test_prefill_with_context(dtype: torch.dtype, bs: int):
+    """Non-causal Q ragged + KV ragged: run_prefill_context_chunk path.
+
+    Per-request KV length is independent of (and >=) Q length, mimicking the
+    chunked-context call site where KV is the cache chunk and Q is the new tokens.
+    """
+    torch.set_default_device("cuda")
+    torch.manual_seed(1)
+
+    if dtype == torch.float8_e4m3fn:
+        warmup_compile_prefill(
+            q_dtype=torch.float8_e4m3fn,
+            d_qk=QK_HEAD_DIM,
+            d_v=V_HEAD_DIM,
+            enable_pdl=False,
+        )
+
+    q_lens = [int(torch.randint(16, 257, (1,)).item()) for _ in range(bs)]
+    kv_lens = [q_lens[i] + int(torch.randint(0, 1025, (1,)).item()) for i in range(bs)]
+
+    kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32)
+    cum_q = _cumsum_int32(q_lens)
+    cum_kv = _cumsum_int32(kv_lens)
+    max_q_len = max(q_lens)
+    max_kv_len = max(kv_lens)
+
+    q, k, v = _make_q_kv(q_lens, kv_lens, dtype)
+
+    # --- TRT-LLM reference ---
+    workspace = torch.zeros(FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8)
+    out_ref = torch.empty(q.shape[0], q.shape[1], v.shape[2], dtype=torch.bfloat16)
+    ref_ret = trtllm_ragged_attention_deepseek(
+        query=q,
+        key=k,
+        value=v,
+        workspace_buffer=workspace,
+        seq_lens=kv_lens_t,
+        max_q_len=max_q_len,
+        max_kv_len=max_kv_len,
+        bmm1_scale=SCALE,
+        bmm2_scale=1.0,
+        o_sf_scale=1.0,
+        batch_size=bs,
+        window_left=-1,
+        cum_seq_lens_q=cum_q,
+        cum_seq_lens_kv=cum_kv,
+        enable_pdl=False,
+        is_causal=False,
+        return_lse=True,
+        out=out_ref,
+    )
+    out_ref, lse_ref = ref_ret[0], ref_ret[1]
+
+    # --- TokenSpeed candidate ---
+    ts_ret = tokenspeed_mla_prefill(
+        query=q,
+        key=k,
+        value=v,
+        seq_lens=kv_lens_t,
+        cum_seq_lens=cum_kv,
+        max_seq_len=max_kv_len,
+        batch_size=bs,
+        softmax_scale=SCALE,
+        is_causal=False,
+        return_lse=True,
+        cum_seq_lens_q=cum_q,
+        max_seq_len_q=max_q_len,
+        enable_pdl=False,
+    )
+    out_ts, lse_ts = ts_ret[0], ts_ret[1]
+
+    if dtype == torch.float8_e4m3fn:
+        atol, rtol = 5e-2, 5e-2
+    else:
+        atol, rtol = 1e-2, 1e-2
+    torch.testing.assert_close(
+        out_ts.to(torch.float32),
+        out_ref.to(torch.float32),
+        atol=atol,
+        rtol=rtol,
+    )
+
+    # LSE: trtllm returns (q_len, num_heads). Tokenspeed convention should
+    # match shape-by-shape — if it doesn't, the LSE transpose contract that
+    # merge_attn_states relies on is broken and this assert surfaces it.
+    assert lse_ref.shape == lse_ts.shape, (
+        f"LSE shape mismatch: trtllm={tuple(lse_ref.shape)} "
+        f"tokenspeed={tuple(lse_ts.shape)}"
+    )
+    # Log-base normalization: trtllm returns LSE in log2, tokenspeed and
+    # vLLM's merge_attn_states (triton_merge_attn_states.py:138) both use
+    # natural-log. Convert trtllm's log2 LSE to natural log before
+    # comparison, otherwise we'd be comparing different bases (factor ln 2).
+    import math
+
+    torch.testing.assert_close(
+        lse_ts.to(torch.float32),
+        lse_ref.to(torch.float32) * math.log(2),
+        atol=5e-3,
+        rtol=5e-3,
+    )
@@ -590,6 +590,33 @@ def _test_extract_tool_calls_streaming(
    ]
    assert_tool_calls(actual_tool_calls, expected_tool_calls)

+    if expected_tool_calls:
+        assert len(tool_parser.streamed_args_for_tool) == len(expected_tool_calls)
+        assert len(tool_parser.prev_tool_call_arr) == len(expected_tool_calls)
+        for i in range(len(expected_tool_calls)):
+            assert (
+                tool_parser.prev_tool_call_arr[i]["arguments"]
+                == tool_parser.streamed_args_for_tool[i]
+            )
+            assert tool_parser.streamed_args_for_tool[i] == function_args_strs[i]
+            assert (
+                tool_parser.prev_tool_call_arr[i]["name"]
+                == expected_tool_calls[i].function.name
+            )
+
+        # Simulate the serving layer's unstreamed-args check
+        index = len(tool_parser.prev_tool_call_arr) - 1
+        args = tool_parser.prev_tool_call_arr[index].get("arguments", {})
+        expected_call = (
+            args if isinstance(args, str) else json.dumps(args, ensure_ascii=False)
+        )
+        actual_call = tool_parser.streamed_args_for_tool[index]
+        remaining_call = expected_call.replace(actual_call, "", 1)
+        assert remaining_call == ""
+    else:
+        assert len(tool_parser.streamed_args_for_tool) == 0
+        assert len(tool_parser.prev_tool_call_arr) == 0
+

@pytest.mark.parametrize(
    ids=[
@@ -855,6 +882,8 @@ def test_extract_tool_calls_streaming_v11_no_tools(
        previous_text = current_text

    assert collected_content == model_output
+    assert len(mistral_tool_parser.streamed_args_for_tool) == 0
+    assert len(mistral_tool_parser.prev_tool_call_arr) == 0


@pytest.mark.parametrize(
@@ -22,7 +22,6 @@ from vllm.config.vllm import set_current_vllm_config
 from vllm.model_executor.layers.attention.mla_attention import (
    QueryLenSupport,
    _DecodeConcatQuantFP8,
-    get_mla_prefill_scale,
 )
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -31,6 +30,7 @@ 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.registry import AttentionBackendEnum
 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
@@ -622,6 +622,19 @@ def run_attention_backend(
            k_scale=k_scale,
        )

+        # Attach prefill backend (normally created by MLAAttention.__init__)
+        prefill_scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
+        prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+        mock_layer.prefill_backend = prefill_backend_cls(
+            num_heads=num_heads,
+            scale=prefill_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,
+        )
+
        # Populate static_forward_context with mock attention layers
        for layer_name in layer_names:
            vllm_config.compilation_config.static_forward_context[layer_name] = (
@@ -787,7 +800,8 @@ def test_backend_correctness(
        f"MLA dimensions don't match: {total_head_size} != {head_size}"
    )
    decode_scale = 1.0 / (total_head_size**0.5)
-    prefill_scale = get_mla_prefill_scale(vllm_config.model_config)
+    qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+    prefill_scale = qk_head_dim**-0.5

    # 2. Generate data and compute SDPA reference output for MLA
    all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
@@ -2,17 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for MLA prefill backend selector."""

-from types import SimpleNamespace
 from unittest.mock import MagicMock, patch

 import pytest
 import torch

 from vllm.config import AttentionConfig, ModelConfig, VllmConfig
-from vllm.model_executor.layers.attention.mla_attention import get_mla_prefill_scale
-from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
-    yarn_get_mscale,
-)
 from vllm.platforms.interface import DeviceCapability
 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
 from vllm.v1.attention.backends.mla.prefill.selector import (
@@ -58,62 +53,6 @@ def _make_vllm_config(
    return mock_vllm_config


-class TestMLAPrefillScale:
-    """Tests for the MLA prefill softmax scale."""
-
-    def test_uses_qk_head_dim_for_deepseek_v2_style_mla(self):
-        model_config = SimpleNamespace(
-            hf_text_config=SimpleNamespace(
-                q_lora_rank=None,
-                kv_lora_rank=512,
-                qk_nope_head_dim=128,
-                qk_rope_head_dim=64,
-                v_head_dim=128,
-                rope_parameters={"rope_type": "default"},
-            )
-        )
-
-        assert get_mla_prefill_scale(model_config) == pytest.approx(192**-0.5)
-
-    def test_applies_deepseek_yarn_mscale(self):
-        model_config = SimpleNamespace(
-            hf_text_config=SimpleNamespace(
-                q_lora_rank=None,
-                kv_lora_rank=512,
-                qk_nope_head_dim=128,
-                qk_rope_head_dim=64,
-                v_head_dim=128,
-                rope_parameters={
-                    "rope_type": "yarn",
-                    "factor": 40,
-                    "mscale_all_dim": 0.707,
-                },
-            )
-        )
-
-        mscale = yarn_get_mscale(40, 0.707)
-        assert get_mla_prefill_scale(model_config) == pytest.approx(
-            192**-0.5 * mscale * mscale
-        )
-
-    def test_deepseek_v4_style_mla_does_not_apply_yarn_mscale(self):
-        model_config = SimpleNamespace(
-            hf_text_config=SimpleNamespace(
-                compress_ratios=[4],
-                q_lora_rank=1536,
-                head_dim=128,
-                qk_rope_head_dim=64,
-                rope_parameters={
-                    "rope_type": "yarn",
-                    "factor": 40,
-                    "mscale_all_dim": 0.707,
-                },
-            )
-        )
-
-        assert get_mla_prefill_scale(model_config) == pytest.approx(128**-0.5)
-
-
 class TestGetMLAPrefillBackend:
    """Tests for get_mla_prefill_backend (public API)."""

@@ -6,6 +6,7 @@ import pytest

 from tests.utils import get_attn_backend_list_based_on_platform
 from vllm import LLM, SamplingParams
+from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
 from vllm.platforms import current_platform
 from vllm.sampling_params import StructuredOutputsParams

@@ -77,3 +78,23 @@ def test_eagle_max_len(
            "is longer than the eagle max length"
        )
        assert o.outputs[0].text == "a b c d e " * 15
+
+
+@pytest.mark.parametrize("spec_max_model_len", [80, 150])
+def test_mtp_speculative_config_max_model_len(spec_max_model_len: int):
+    """Regression test for #41456: max_model_len in speculative config
+    should be respected for the draft model."""
+    model_config = ModelConfig(
+        model="XiaomiMiMo/MiMo-7B-Base",
+        runner="generate",
+        max_model_len=200,
+        trust_remote_code=True,
+    )
+    spec_config = SpeculativeConfig(
+        target_model_config=model_config,
+        target_parallel_config=ParallelConfig(),
+        method="mtp",
+        num_speculative_tokens=1,
+        max_model_len=spec_max_model_len,
+    )
+    assert spec_config.draft_model_config.max_model_len == spec_max_model_len
@@ -626,6 +626,7 @@ class SpeculativeConfig:
                    revision=self.revision,
                    code_revision=self.code_revision,
                    tokenizer_revision=self.target_model_config.tokenizer_revision,
+                    max_model_len=self.max_model_len,  # type: ignore[arg-type]
                    spec_target_max_model_len=self.target_model_config.max_model_len,
                    quantization=self.quantization,
                    enforce_eager=self.target_model_config.enforce_eager,
@@ -837,10 +838,17 @@ class SpeculativeConfig:

            return speculative_max_model_len

-        return min(
+        result = min(
            draft_max_model_len,
            target_max_model_len,
        )
+        if result != draft_max_model_len:
+            logger.info(
+                "Overriding draft model max model len from %d to %d",
+                draft_max_model_len,
+                result,
+            )
+        return result

    @staticmethod
    def _verify_and_get_draft_tp(
@@ -23,7 +23,9 @@ from openai.types.responses import (
    ResponseOutputItem,
    ResponseOutputItemAddedEvent,
    ResponseOutputItemDoneEvent,
+    ResponseOutputMessage,
    ResponsePrompt,
+    ResponseReasoningItem,
    ResponseReasoningTextDeltaEvent,
    ResponseReasoningTextDoneEvent,
    ResponseStatus,
@@ -451,18 +453,21 @@ class ResponsesRequest(OpenAIBaseModel):

    @model_validator(mode="before")
    @classmethod
-    def function_call_parsing(cls, data):
-        """Parse function_call dictionaries into ResponseFunctionToolCall objects.
-        This ensures Pydantic can properly resolve union types in the input field.
-        Function calls provided as dicts are converted to ResponseFunctionToolCall
-        objects before validation, while invalid structures are left for Pydantic
-        to reject with appropriate error messages.
-        """
+    def input_item_parsing(cls, data):
+        """Parse input items that are missing required fields or that Pydantic
+        cannot disambiguate in a Union of TypedDict / BaseModel types.

+        Specifically handles:
+        - function_call -> ResponseFunctionToolCall
+        - reasoning     -> ResponseReasoningItem (auto-generates id)
+        - message(role=assistant) -> ResponseOutputMessage (auto-generates
+          id/status and annotations)
+
+        Invalid structures are left for Pydantic to reject.
+        """
        input_data = data.get("input")

        # Early return for None, strings, or bytes
-        # (strings are iterable but shouldn't be processed)
        if input_data is None or isinstance(input_data, (str, bytes)):
            return data

@@ -476,16 +481,61 @@ class ResponsesRequest(OpenAIBaseModel):

        processed_input = []
        for item in input_data:
-            if isinstance(item, dict) and item.get("type") == "function_call":
+            if not isinstance(item, dict):
+                processed_input.append(item)
+                continue
+
+            item_type = item.get("type")
+
+            if item_type == "function_call":
                try:
                    processed_input.append(ResponseFunctionToolCall(**item))
                except ValidationError:
-                    # Let Pydantic handle validation for malformed function calls
                    logger.debug(
                        "Failed to parse function_call to ResponseFunctionToolCall, "
                        "leaving for Pydantic validation"
                    )
                    processed_input.append(item)
+
+            elif item_type == "reasoning":
+                if "id" not in item:
+                    item = {**item, "id": f"rs_{random_uuid()}"}
+                try:
+                    processed_input.append(ResponseReasoningItem(**item))
+                except ValidationError:
+                    logger.debug(
+                        "Failed to parse reasoning to ResponseReasoningItem, "
+                        "leaving for Pydantic validation"
+                    )
+                    processed_input.append(item)
+
+            elif item_type == "message" and item.get("role") == "assistant":
+                item = dict(item)
+                if "id" not in item:
+                    item["id"] = f"msg_{random_uuid()}"
+                if "status" not in item:
+                    item["status"] = "completed"
+                # ResponseOutputText requires annotations
+                if isinstance(item.get("content"), list):
+                    new_content = []
+                    for c in item["content"]:
+                        if (
+                            isinstance(c, dict)
+                            and c.get("type") == "output_text"
+                            and "annotations" not in c
+                        ):
+                            c = {**c, "annotations": []}
+                        new_content.append(c)
+                    item["content"] = new_content
+                try:
+                    processed_input.append(ResponseOutputMessage(**item))
+                except ValidationError:
+                    logger.debug(
+                        "Failed to parse assistant message to ResponseOutputMessage, "
+                        "leaving for Pydantic validation"
+                    )
+                    processed_input.append(item)
+
            else:
                processed_input.append(item)

@@ -238,9 +238,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
    kFp8StaticTensorSym,
    kNvfp4Dynamic,
 )
-from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
-    yarn_get_mscale,
-)
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.utils.math_utils import cdiv, round_down
@@ -262,7 +259,10 @@ from vllm.v1.attention.backend import (
    MLAAttentionImpl,
    SparseMLAAttentionImpl,
 )
-from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackend
+from vllm.v1.attention.backends.mla.prefill import (
+    MLAPrefillBackend,
+    get_mla_prefill_backend,
+)
 from vllm.v1.attention.backends.utils import (
    get_dcp_local_seq_lens,
    split_decodes_and_prefills,
@@ -454,20 +454,32 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
        self.use_direct_call = not current_platform.opaque_attention_op()

-        compilation_config = get_current_vllm_config().compilation_config
+        vllm_config = get_current_vllm_config()
+        compilation_config = vllm_config.compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

+        prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+        self.prefill_backend = prefill_backend_cls(
+            num_heads=self.num_heads,
+            scale=self.scale,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_nope_head_dim=self.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            v_head_dim=self.v_head_dim,
+            vllm_config=vllm_config,
+        )
+
        self.kv_cache = torch.tensor([])

        self.use_sparse = use_sparse

-        vllm_config = get_current_vllm_config_or_none()
+        _vllm_config = get_current_vllm_config_or_none()
        self.dcp_a2a = (
-            vllm_config is not None
-            and vllm_config.parallel_config.decode_context_parallel_size > 1
-            and vllm_config.parallel_config.dcp_comm_backend == "a2a"
+            _vllm_config is not None
+            and _vllm_config.parallel_config.decode_context_parallel_size > 1
+            and _vllm_config.parallel_config.dcp_comm_backend == "a2a"
        )

        # Initialize q/k/v range constants.
@@ -1330,35 +1342,6 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:
    )


-def get_mla_prefill_scale(model_config: ModelConfig) -> float:
-    hf_text_config = model_config.hf_text_config
-    mla_dims = get_mla_dims(model_config)
-    qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
-    scale = qk_head_dim**-0.5
-
-    # Deepseek V4 disables YaRN mscale for attention; Deepseek V2/V3 applies
-    # the same mscale correction when constructing the MLA attention module.
-    if hasattr(hf_text_config, "compress_ratios"):
-        return scale
-
-    rope_parameters = getattr(hf_text_config, "rope_parameters", None)
-    if rope_parameters is None:
-        rope_parameters = getattr(hf_text_config, "rope_scaling", None)
-
-    if rope_parameters is None:
-        return scale
-
-    rope_type = rope_parameters.get("rope_type", rope_parameters.get("type"))
-    apply_yarn_scaling = rope_parameters.get("apply_yarn_scaling", True)
-    if rope_type != "default" and apply_yarn_scaling:
-        mscale_all_dim = rope_parameters.get("mscale_all_dim", False)
-        scaling_factor = rope_parameters["factor"]
-        mscale = yarn_get_mscale(float(scaling_factor), float(mscale_all_dim))
-        scale *= mscale * mscale
-
-    return scale
-
-
@functools.cache
 def backend_supports_prefill_query_quantization() -> bool:
    """Check if the selected MLA prefill backend supports query quantization.
@@ -1384,6 +1367,7 @@ def backend_supports_prefill_query_quantization() -> bool:
    return backend_cls.get_name() in (
        "FLASHINFER",
        "TRTLLM_RAGGED",
+        "TOKENSPEED_MLA",
    )


@@ -1554,20 +1538,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                device=device,
            )

-        from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
-
-        prefill_backend_cls = get_mla_prefill_backend(vllm_config)
-        self._prefill_backend = prefill_backend_cls(
-            num_heads=self.num_heads,
-            scale=get_mla_prefill_scale(self.model_config),
-            kv_lora_rank=self.mla_dims.kv_lora_rank,
-            qk_nope_head_dim=self.mla_dims.qk_nope_head_dim,
-            qk_rope_head_dim=self.mla_dims.qk_rope_head_dim,
-            v_head_dim=self.mla_dims.v_head_dim,
-            vllm_config=vllm_config,
-            device=device,
-            layer_names=layer_names,
-        )
+        self._prefill_backend = self.compilation_config.static_forward_context[
+            layer_names[0]
+        ].prefill_backend

        supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY
        self._init_reorder_batch_threshold(
@@ -140,6 +140,8 @@ def _fwd_kernel_ep_scatter_2(
    offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD)
    mask_s = offset_in_s < SCALE_HIDDEN_SIZE

+    output_tensor_stride0 = output_tensor_stride0.to(tl.int64)
+
    for token_id in range(start_token_id, total_token_num, grid_num):
        to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask)
        to_copy_s = tl.load(
@@ -154,12 +156,13 @@ def _fwd_kernel_ep_scatter_2(

            if expert_id >= 0:
                dest_token_index = tl.atomic_add(expert_start_loc + expert_id, 1)
+                dest_token_index_i64 = dest_token_index.to(tl.int64)
                tl.store(
                    output_index + token_id * output_index_stride0 + topk_index,
                    dest_token_index,
                )
                output_tensor_ptr = (
-                    output_tensor + dest_token_index * output_tensor_stride0
+                    output_tensor + dest_token_index_i64 * output_tensor_stride0
                )
                output_tensor_scale_ptr = (
                    output_tensor_scale + dest_token_index * output_tensor_scale_stride0
@@ -623,13 +623,6 @@ class MistralToolParser(ToolParser):
        if len(delta_tool_calls) > 0:
            delta.tool_calls = delta_tool_calls

-        # HACK: serving_chat.py inspects the internal state of tool parsers
-        # when determining its final streaming delta, automatically
-        # adding autocompleted JSON.
-        # These two lines avoid that nonsense while ensuring finish_reason
-        # is set to tool_calls when at least one tool is called.
-        if delta_tool_calls and not self.prev_tool_call_arr:
-            self.prev_tool_call_arr = [{"arguments": {}}]
        return delta

    def _generate_delta_tool_call(self, delta_text: str) -> list[DeltaToolCall]:
@@ -642,6 +635,8 @@ class MistralToolParser(ToolParser):
            StreamingState.PARSING_ARGUMENTS,
        ] and delta_text.startswith(self.bot_token):
            self.current_tool_id += 1
+            self.streamed_args_for_tool.append("")
+            self.prev_tool_call_arr.append({})
            self.streaming_state = StreamingState.PARSING_NAME
            delta_text = delta_text.replace(self.bot_token, "", 1)
        if self.streaming_state == StreamingState.PARSING_NAME:
@@ -655,6 +650,9 @@ class MistralToolParser(ToolParser):
                self.current_tool_name += delta_function_name
                # HF tokenizers may include [ARGS] in the text
                self.current_tool_name = self.current_tool_name.replace("[ARGS]", "")
+                self.prev_tool_call_arr[self.current_tool_id]["name"] = (
+                    self.current_tool_name
+                )
                delta_text = delta_text[len(delta_function_name) :]
                self.streaming_state = StreamingState.PARSING_ARGUMENTS
            else:
@@ -671,6 +669,10 @@ class MistralToolParser(ToolParser):
                self.streaming_state = StreamingState.TOOL_COMPLETE
            else:
                delta_arguments = delta_text
+            self.streamed_args_for_tool[self.current_tool_id] += delta_arguments
+            self.prev_tool_call_arr[self.current_tool_id]["arguments"] = (
+                self.streamed_args_for_tool[self.current_tool_id]
+            )
            ret = []
            if self.current_tool_name or delta_arguments:
                ret += [
@@ -820,9 +822,12 @@ class MistralToolParser(ToolParser):
                    if self.current_tool_mistral_id is not None:
                        current_tool_call.id = self.current_tool_mistral_id
                        self.current_tool_mistral_id = None
+                    self._track_streamed_args_pre_v11(current_tool_call)
                    delta_tool_calls.append(current_tool_call)
                current_tool_call_modified = False
                self.current_tool_id += 1
+                self.streamed_args_for_tool.append("")
+                self.prev_tool_call_arr.append({})
                self.current_tool_mistral_id = MistralToolCall.generate_random_id()
                current_tool_call = DeltaToolCall(
                    index=self.current_tool_id,
@@ -835,6 +840,9 @@ class MistralToolParser(ToolParser):
                # we have the complete tool name
                current_tool_call_modified = True
                current_tool_call.function.name = self.current_tool_name
+                self.prev_tool_call_arr[self.current_tool_id]["name"] = (
+                    self.current_tool_name
+                )
                self.current_tool_name = None
            if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED:
                self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
@@ -860,16 +868,9 @@ class MistralToolParser(ToolParser):
            if self.current_tool_mistral_id is not None:
                current_tool_call.id = self.current_tool_mistral_id
                self.current_tool_mistral_id = None
+            self._track_streamed_args_pre_v11(current_tool_call)
            delta_tool_calls.append(current_tool_call)

-        # HACK: serving_chat.py inspects the internal state of tool parsers
-        # when determining it's final streaming delta, automatically
-        # adding autocompleted JSON.
-        # These two lines avoid that nonsense while ensuring finish_reason
-        # is set to tool_calls when at least one tool is called.
-        if delta_tool_calls and not self.prev_tool_call_arr:
-            self.prev_tool_call_arr = [{"arguments": {}}]
-
        if content or len(delta_tool_calls) > 0:
            delta_message = DeltaMessage()
            if content:
@@ -883,6 +884,16 @@ class MistralToolParser(ToolParser):
            else:
                return None

+    def _track_streamed_args_pre_v11(self, tool_call: DeltaToolCall) -> None:
+        r"""Accumulate `tool_call` arguments into the streaming state."""
+        if tool_call.function is not None and tool_call.function.arguments is not None:
+            self.streamed_args_for_tool[self.current_tool_id] += (
+                tool_call.function.arguments
+            )
+            self.prev_tool_call_arr[self.current_tool_id]["arguments"] = (
+                self.streamed_args_for_tool[self.current_tool_id]
+            )
+
    def _split_delta(
        self,
        delta_text: str,
@@ -81,8 +81,6 @@ class MLAPrefillBackend(ABC):
        qk_rope_head_dim: int,
        v_head_dim: int,
        vllm_config: "VllmConfig",
-        device: torch.device,
-        layer_names: list[str] | None = None,
    ) -> None:
        self.num_heads = num_heads
        self.scale = scale
@@ -91,8 +89,6 @@ class MLAPrefillBackend(ABC):
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.vllm_config = vllm_config
-        self.device = device
-        self.layer_names = layer_names

    def prepare_metadata(  # noqa: B027
        self,
@@ -44,8 +44,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
        qk_rope_head_dim: int,
        v_head_dim: int,
        vllm_config: "VllmConfig",
-        device: torch.device,
-        layer_names: list[str] | None = None,
    ) -> None:
        super().__init__(
            num_heads=num_heads,
@@ -55,8 +53,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
            qk_rope_head_dim=qk_rope_head_dim,
            v_head_dim=v_head_dim,
            vllm_config=vllm_config,
-            device=device,
-            layer_names=layer_names,
        )

        # Handle the differences between the flash_attn_varlen from
@@ -9,6 +9,7 @@ import torch
 import vllm.envs as envs
 from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
 from vllm.v1.attention.backends.utils import (
+    PerLayerParameters,
    get_per_layer_parameters,
    infer_global_hyperparameters,
 )
@@ -62,8 +63,6 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
        qk_rope_head_dim: int,
        v_head_dim: int,
        vllm_config: "VllmConfig",
-        device: torch.device,
-        layer_names: list[str] | None = None,
    ) -> None:
        super().__init__(
            num_heads=num_heads,
@@ -73,25 +72,11 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
            qk_rope_head_dim=qk_rope_head_dim,
            v_head_dim=v_head_dim,
            vllm_config=vllm_config,
-            device=device,
-            layer_names=layer_names,
        )

        self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
        self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []
-        if layer_names is None:
-            raise ValueError(
-                "FlashInferPrefillBackend requires layer_names to "
-                "initialize global hyperparameters."
-            )
-
-        from vllm.model_executor.layers.attention.mla_attention import (
-            MLACommonImpl,
-        )
-
-        self._global_hyperparameters = infer_global_hyperparameters(
-            get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl)  # type: ignore[type-abstract]
-        )
+        self._global_hyperparameters: PerLayerParameters | None = None

    def _ensure_chunks(
        self,
@@ -106,10 +91,36 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
                    )
                )

+    def _resolve_global_hyperparameters(self) -> PerLayerParameters:
+        if self._global_hyperparameters is not None:
+            return self._global_hyperparameters
+
+        from vllm.model_executor.layers.attention.mla_attention import (
+            MLAAttention,
+            MLACommonImpl,
+        )
+
+        forward_context = self.vllm_config.compilation_config.static_forward_context
+        layer_names = [
+            name
+            for name, layer in forward_context.items()
+            if isinstance(layer, MLAAttention)
+        ]
+
+        self._global_hyperparameters = infer_global_hyperparameters(
+            get_per_layer_parameters(
+                self.vllm_config,
+                layer_names,
+                MLACommonImpl,  # type: ignore[type-abstract]
+            )
+        )
+        return self._global_hyperparameters
+
    def prepare_metadata(
        self,
        prefill_metadata: "MLACommonPrefillMetadata",
    ) -> None:
+        global_hyperparameters = self._resolve_global_hyperparameters()
        qo_indptr = prefill_metadata.query_start_loc
        has_context = prefill_metadata.chunked_context is not None
        (workspace_buffer,) = current_workspace_manager().get_simultaneous(
@@ -144,9 +155,9 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
            head_dim_qk=head_dim_qk,
            head_dim_vo=head_dim_vo,
            causal=True,
-            sm_scale=self._global_hyperparameters.sm_scale,
-            window_left=self._global_hyperparameters.window_left,
-            logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
+            sm_scale=global_hyperparameters.sm_scale,
+            window_left=global_hyperparameters.window_left,
+            logits_soft_cap=global_hyperparameters.logits_soft_cap,
            q_data_type=prefill_metadata.q_data_type,
            o_data_type=prefill_metadata.output_dtype,
        )
@@ -165,9 +176,9 @@ class FlashInferPrefillBackend(MLAPrefillBackend):
                    head_dim_qk=head_dim_qk,
                    head_dim_vo=head_dim_vo,
                    causal=False,
-                    sm_scale=self._global_hyperparameters.sm_scale,
-                    window_left=self._global_hyperparameters.window_left,
-                    logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
+                    sm_scale=global_hyperparameters.sm_scale,
+                    window_left=global_hyperparameters.window_left,
+                    logits_soft_cap=global_hyperparameters.logits_soft_cap,
                    q_data_type=prefill_metadata.q_data_type,
                    o_data_type=prefill_metadata.output_dtype,
                )
@@ -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."""
@@ -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 [
@@ -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()
@@ -51,8 +51,6 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
        qk_rope_head_dim: int,
        v_head_dim: int,
        vllm_config: "VllmConfig",
-        device: torch.device,
-        layer_names: list[str] | None = None,
    ) -> None:
        super().__init__(
            num_heads=num_heads,
@@ -62,8 +60,6 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend):
            qk_rope_head_dim=qk_rope_head_dim,
            v_head_dim=v_head_dim,
            vllm_config=vllm_config,
-            device=device,
-            layer_names=layer_names,
        )

    def _get_workspace_buffer(self) -> torch.Tensor:
@@ -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
@@ -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"
Author	SHA1	Message	Date
Yongye ZhuandClaude Opus 4.7	58c8a5eaa5	[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:45:16 +00:00
Yongye Zhu	c4547482ca	update version Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:28:20 +00:00
Yongye Zhu	91ef0afcb2	adding to cuda.txt dependency Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 07:20:48 +00:00
Yongye Zhu	97cd2c41ad	fix precommit Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 03:14:49 +00:00
Yongye ZhuandClaude Opus 4.7	c001535038	[Attention][TokenSpeed MLA] Force prefill V tensor contiguous before kernel call `v` arrives at both `run_prefill_new_tokens` and `run_prefill_context_chunk` as the second half of `kv_nope.split([qk_nope_head_dim, v_head_dim], dim=-1)` in mla_attention.py — a non-contiguous view along the last dim. The kernel internally does `v.reshape(1, total_kv, h_k, 1, d_v)` which silently copies when the input is non-contiguous; pull that copy out so the layout is predictable at the kernel boundary. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:53:46 +00:00
Yongye ZhuandClaude Opus 4.7	de6bc297df	[Attention][TokenSpeed MLA] Surface install hint when package missing Previously a user explicitly selecting TOKENSPEED_MLA without `tokenspeed_mla` installed got either a generic "required dependencies not available" message (prefill backend) or a raw ModuleNotFoundError deep inside forward_mqa at the first request (decode backend). Now both backends fail at startup with the exact install command: `uv pip install tokenspeed-mla`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	0012818287	[Attention][TokenSpeed MLA] Fix trtllm LSE parity test: log2 → natural log trtllm_ragged_attention_deepseek returns LSE in log2; tokenspeed and merge_attn_states use natural log. Multiply the trtllm reference by ln 2 before comparison. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	73cd7e25ae	[Attention][TokenSpeed MLA] Warm up BF16 prefill compile, drop seq_lens computation Pre-JIT both BF16 and FP8 prefill kernels at backend init since the dtype isn't visible from `__init__` — depends on `use_prefill_query_quantization`. Move the per-forward `seq_lens` computation into `prepare_metadata` and document the cuda-graph padding interaction with `query_start_loc`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	964c6eb485	[Attention][TokenSpeed MLA] Fix decode FP8 numerics: pass output_scale and assert FP8 Q The decode kernel needs both bmm scales to recover correct outputs from an FP8 KV cache: bmm1 (softmax_scale = scale * q_scale * k_scale) and bmm2 (output_scale = k_scale, since V is stored as V_real / k_scale). We were only passing bmm1, which left bmm2 = 1.0 and produced silently wrong output. Also assert query dtype is float8_e4m3fn on entry to forward_mqa. supports_quant_query_input=True (inherited from MLACommonImpl) tells the upstream pipeline to FP8-quantize Q via _decode_concat_quant_fp8_op; the kernel is shape-specialized for FP8 Q + FP8 KV, so any other dtype here means the upstream quant path didn't run and the kernel will produce garbage. Failing loud beats failing silent. Verified: gsm8k matches reference with TOKENSPEED_MLA decode + FLASH_ATTN prefill on Kimi-K2.5-NVFP4 / TP=4 / B200. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
Yongye ZhuandClaude Opus 4.7	d0e6514bf8	[Attention] Add TOKENSPEED_MLA backend for DeepSeek R1 prefill + decode on Blackwell Wires the tokenspeed_mla CuTe DSL kernels into vLLM as a new MLA backend, covering both prefill (tokenspeed_mla_prefill) and decode (tokenspeed_mla_decode). Targets Blackwell (SM100) with FP8 KV cache and DeepSeek R1 MLA dimensions; users opt in via -ac '{"backend":"TOKENSPEED_MLA","mla_prefill_backend":"TOKENSPEED_MLA"}'. Includes numeric parity tests against the trtllm reference kernels. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>	2026-05-06 02:49:09 +00:00
ChaunceyandGitHub	c7aa186d67	[Frontend] Supports resubmitting output items with missing fields in Responses API (#41355 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2026-05-05 22:21:33 -04:00
Kevin H. Luu GitHub OpenAI Codex	f653761252	[CI] Route part of B200 jobs to b200-k8s (#41453 ) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>	2026-05-05 19:00:30 -07:00
Andreas KaratzasandGitHub	4a8ae26e53	[ROCm][CI] Use vLLM generation defaults for DeepSeek prefetch-offload eval (#41575 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-05-06 01:08:12 +00:00
Kevin H. LuuandGitHub	1333864408	[CI] Automate Docker Hub release image publishing (#40415 ) Signed-off-by: khluu <khluu000@gmail.com>	2026-05-06 00:15:23 +00:00
Matthew BonanniandGitHub	01b9b5af67	[Attention] Minor refactor: layer takes ownership of the MLA prefill backend (#41744 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2026-05-05 23:22:41 +00:00
Dhruvil Bhatt GitHub EC2 Default User Robert Shaw	8c57b6e7bc	Bump model-hosting-container-standards to >= 0.1.14 (#39755 ) Signed-off-by: EC2 Default User <ec2-user@ip-172-31-20-13.us-west-2.compute.internal> Co-authored-by: EC2 Default User <ec2-user@ip-172-31-20-13.us-west-2.compute.internal> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-05-05 19:09:57 -04:00
Lanze LiuandGitHub	79246b5ea6	[Spec Decode] Fix max_model_len logging in speculative config for draft model (#41571 ) Signed-off-by: Lanze Liu <lanzetech@gmail.com>	2026-05-05 21:56:06 +00:00
Matej Sirovatka GitHub Tyler Michael Smith	48954de237	Fix DeepGEMM ep_scatter output address overflow (#39213 ) Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>	2026-05-05 18:56:56 +00:00
Julien DenizeandGitHub	c6235ed180	[BUGFIX] Support streamed_args_for_tool in MistralToolParser (#41730 ) Signed-off-by: juliendenize <julien.denize@mistral.ai>	2026-05-05 17:48:53 +00:00