[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
44 changed files with 1368 additions and 1365 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
@@ -50,7 +50,6 @@ MTPModelTypes = Literal[
    "pangu_ultra_moe_mtp",
    "step3p5_mtp",
    "hy_v3_mtp",
-    "gemma4_mtp",
 ]
 NgramGPUTypes = Literal["ngram_gpu"]
 DFlashModelTypes = Literal["dflash"]
@@ -492,17 +491,6 @@ class SpeculativeConfig:
                {"n_predict": n_predict, "architectures": ["HYV3MTPModel"]}
            )

-        if hf_config.model_type == "gemma4_assistant":
-            hf_config.model_type = "gemma4_mtp"
-            text_config = getattr(hf_config, "text_config", hf_config)
-            # The assistant runs all decoder layers in a single forward
-            # call to produce one draft token, so n_predict=1.
-            # num_kv_shared_layers must be 0: cross-model KV sharing is
-            # set up by the proposer after model construction.
-            if hasattr(text_config, "num_kv_shared_layers"):
-                text_config.num_kv_shared_layers = 0
-            hf_config.update({"n_predict": 1, "architectures": ["Gemma4MTPModel"]})
-
        return hf_config

    def __post_init__(self):
@@ -638,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,
@@ -849,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(
@@ -1044,14 +1040,6 @@ class SpeculativeConfig:
            slots_per_req += 1
        return slots_per_req

-    def use_gemma4_mtp(self) -> bool:
-        return (
-            self.method == "mtp"
-            and self.draft_model_config is not None
-            and getattr(self.draft_model_config.hf_config, "model_type", None)
-            == "gemma4_mtp"
-        )
-
    def use_eagle(self) -> bool:
        return self.method in ("eagle", "eagle3", "mtp", "dflash")

@@ -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
@@ -1,602 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Inference-only Gemma4 MTP (Multi-Token Prediction) model.
-
-The Gemma4 assistant model is a lightweight decoder that shares KV cache
-with the target (backbone) model.  All assistant decoder layers are
-KV-shared: they only have Q projections (no K/V projections or norms),
-and read K/V from the target model's cache at runtime.
-
-Checkpoint layout (``gemma4_assistant``)::
-
-    model.embed_tokens.*          -- token embeddings
-    model.layers.{i}.*            -- decoder layers (Q-only attention + MLP)
-    model.norm.*                  -- final RMSNorm
-    pre_projection.*              -- Linear(2 * backbone_hidden_size, hidden_size)
-    post_projection.*             -- Linear(hidden_size, backbone_hidden_size)
-    lm_head.*                     -- language model head (tied to embed_tokens)
-    masked_embedding.centroids.*  -- centroid projection (when use_ordered_embeddings)
-    masked_embedding.token_ordering -- token-to-centroid mapping buffer
-"""
-
-from collections.abc import Iterable
-
-import torch
-from torch import nn
-
-from vllm.compilation.decorators import support_torch_compile
-from vllm.config import CacheConfig, VllmConfig
-from vllm.distributed import (
-    get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_gather,
-)
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-from vllm.sequence import IntermediateTensors
-
-from .gemma4 import Gemma4MLP, _get_text_config
-from .utils import (
-    AutoWeightsLoader,
-    WeightsMapper,
-    extract_layer_index,
-    maybe_prefix,
-)
-
-logger = init_logger(__name__)
-
-
-class Gemma4MTPMaskedEmbedder(nn.Module):
-    """Sparse logit computation via centroid-based vocabulary masking.
-
-    Instead of computing logits against the full vocabulary, projects
-    hidden states to centroid scores, selects top-K centroids, and
-    computes logits only for the ~top_k * (vocab_size / num_centroids)
-    tokens belonging to those centroids.
-    """
-
-    token_ordering: torch.Tensor
-
-    def __init__(
-        self,
-        hidden_size: int,
-        vocab_size: int,
-        num_centroids: int,
-        centroid_intermediate_top_k: int,
-    ) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.vocab_size = vocab_size
-        self.num_centroids = num_centroids
-        self.centroid_intermediate_top_k = centroid_intermediate_top_k
-        self.vocab_size_per_centroid = vocab_size // num_centroids
-        self.num_selected = centroid_intermediate_top_k * self.vocab_size_per_centroid
-
-        self.centroids = nn.Linear(hidden_size, num_centroids, bias=False)
-        self.register_buffer(
-            "token_ordering",
-            torch.empty(vocab_size, dtype=torch.long),
-        )
-
-    def _select_and_score(
-        self,
-        hidden_states: torch.Tensor,
-        lm_head_weight: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Centroid selection + sparse dot product.
-
-        Returns:
-            logits: (num_tokens, num_selected) sparse logits.
-            indices: (num_tokens, num_selected) corresponding vocab indices.
-        """
-        num_tokens = hidden_states.shape[0]
-        _, top_k_indices = torch.topk(
-            self.centroids(hidden_states),
-            k=self.centroid_intermediate_top_k,
-            dim=-1,
-        )
-        clusters = self.token_ordering.view(
-            self.num_centroids,
-            self.vocab_size_per_centroid,
-        )
-        selected = clusters[top_k_indices]
-        embeddings = lm_head_weight[selected.reshape(-1)].view(
-            num_tokens,
-            self.num_selected,
-            self.hidden_size,
-        )
-        logits = torch.einsum("td,tsd->ts", hidden_states, embeddings)
-        return logits, selected.view(num_tokens, -1)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        lm_head_weight: torch.Tensor,
-    ) -> torch.Tensor:
-        """Full-vocab logits with non-selected positions masked to -inf."""
-        logits, indices = self._select_and_score(hidden_states, lm_head_weight)
-        output = torch.full(
-            (hidden_states.shape[0], self.vocab_size),
-            fill_value=torch.finfo(hidden_states.dtype).min,
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-        return output.scatter_(-1, indices, logits)
-
-    def get_top_tokens(
-        self,
-        hidden_states: torch.Tensor,
-        lm_head_weight: torch.Tensor,
-    ) -> torch.Tensor:
-        """Sparse argmax — returns vocab token IDs without full-vocab tensor."""
-        logits, indices = self._select_and_score(hidden_states, lm_head_weight)
-        return indices.gather(-1, logits.argmax(-1, keepdim=True)).squeeze(-1)
-
-
-class Gemma4MTPAttention(nn.Module):
-    """Q-only attention for Gemma4 MTP layers.
-
-    K/V come from the target model's KV cache via
-    ``kv_sharing_target_layer_name`` (set by the proposer after
-    model construction).
-    """
-
-    def __init__(
-        self,
-        config,
-        hidden_size: int,
-        num_heads: int,
-        num_kv_heads: int,
-        head_dim: int,
-        max_position_embeddings: int,
-        cache_config: CacheConfig | None = None,
-        quant_config: QuantizationConfig | None = None,
-        attn_logits_soft_cap: float | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.config = config
-        self.hidden_size = hidden_size
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.total_num_heads = num_heads
-        self.num_heads = self.total_num_heads // tp_size
-        self.total_num_kv_heads = num_kv_heads
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
-        self.head_dim = head_dim
-        self.q_size = self.num_heads * self.head_dim
-        self.scaling = 1.0
-
-        self.q_proj = ColumnParallelLinear(
-            hidden_size,
-            self.total_num_heads * self.head_dim,
-            bias=config.attention_bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.q_proj",
-        )
-        self.o_proj = RowParallelLinear(
-            self.total_num_heads * self.head_dim,
-            hidden_size,
-            bias=config.attention_bias,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-        )
-        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-
-        layer_idx = extract_layer_index(prefix)
-        layer_type = config.layer_types[layer_idx]
-        self.is_sliding = layer_type == "sliding_attention"
-        sliding_window = config.sliding_window if self.is_sliding else None
-
-        if layer_type in config.rope_parameters:
-            rope_parameters = dict(config.rope_parameters[layer_type])
-        else:
-            rope_parameters = dict(config.rope_parameters.copy())
-            if self.is_sliding:
-                rope_parameters["rope_theta"] = getattr(
-                    config, "rope_local_base_freq", 10000.0
-                )
-
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            max_position=max_position_embeddings,
-            rope_parameters=rope_parameters,
-            is_neox_style=True,
-        )
-
-        # kv_sharing_target_layer_name is set after model construction
-        # by Gemma4Proposer._setup_gemma4_kv_sharing().
-        self.is_kv_shared_layer = True
-        self.attn = Attention(
-            self.num_heads,
-            self.head_dim,
-            self.scaling,
-            num_kv_heads=self.num_kv_heads,
-            cache_config=cache_config,
-            quant_config=quant_config,
-            logits_soft_cap=attn_logits_soft_cap,
-            per_layer_sliding_window=sliding_window,
-            prefix=f"{prefix}.attn",
-        )
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        **kwargs,
-    ) -> torch.Tensor:
-        q, _ = self.q_proj(hidden_states)
-
-        q = q.unflatten(-1, (self.num_heads, self.head_dim))
-        q = self.q_norm(q)
-        q = q.flatten(-2, -1)
-
-        q, _ = self.rotary_emb(positions, q, None)
-
-        # Attention reads K/V from the target's cache via KV sharing;
-        # these dummy tensors are never consumed but required by the API.
-        num_tokens = q.shape[0]
-        kv_dummy = torch.empty(
-            num_tokens,
-            self.num_kv_heads * self.head_dim,
-            dtype=q.dtype,
-            device=q.device,
-        )
-        attn_output = self.attn(q, kv_dummy, kv_dummy)
-        output, _ = self.o_proj(attn_output)
-        return output
-
-
-class Gemma4MTPDecoderLayer(nn.Module):
-    def __init__(
-        self,
-        config,
-        cache_config: CacheConfig | None = None,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        layer_idx = extract_layer_index(prefix)
-        layer_type = config.layer_types[layer_idx]
-        is_full_attention = layer_type == "full_attention"
-        head_dim = (
-            getattr(config, "global_head_dim", config.head_dim)
-            if is_full_attention
-            else config.head_dim
-        )
-
-        self.self_attn = Gemma4MTPAttention(
-            config=config,
-            hidden_size=self.hidden_size,
-            num_heads=config.num_attention_heads,
-            num_kv_heads=config.num_key_value_heads,
-            head_dim=head_dim,
-            max_position_embeddings=config.max_position_embeddings,
-            cache_config=cache_config,
-            quant_config=quant_config,
-            attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None),
-            prefix=f"{prefix}.self_attn",
-        )
-
-        self.mlp = Gemma4MLP(
-            hidden_size=self.hidden_size,
-            intermediate_size=config.intermediate_size,
-            hidden_activation=config.hidden_activation,
-            quant_config=quant_config,
-            prefix=f"{prefix}.mlp",
-        )
-
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-        self.pre_feedforward_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-        self.post_feedforward_layernorm = RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-
-        self.register_buffer("layer_scalar", torch.ones(1))
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        residual: torch.Tensor | None,
-        **kwargs,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(residual)
-
-        hidden_states = self.self_attn(
-            positions=positions,
-            hidden_states=hidden_states,
-            **kwargs,
-        )
-
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = hidden_states + residual
-        residual = hidden_states
-
-        hidden_states = self.pre_feedforward_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-
-        hidden_states = self.post_feedforward_layernorm(hidden_states)
-        hidden_states = hidden_states + residual
-
-        hidden_states = hidden_states * self.layer_scalar
-        return hidden_states, None
-
-
-class Gemma4MultiTokenPredictor(nn.Module):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-
-        config = vllm_config.speculative_config.draft_model_config.hf_config
-        text_config = _get_text_config(config)
-        self.config = text_config
-
-        self.hidden_size = text_config.hidden_size
-        self.backbone_hidden_size = getattr(
-            config, "backbone_hidden_size", self.hidden_size
-        )
-        self.vocab_size = text_config.vocab_size
-        self.num_mtp_layers = text_config.num_hidden_layers
-
-        self.embed_tokens = VocabParallelEmbedding(
-            self.vocab_size,
-            self.hidden_size,
-        )
-
-        self.pre_projection = ColumnParallelLinear(
-            2 * self.backbone_hidden_size,
-            self.hidden_size,
-            bias=False,
-            gather_output=True,
-            prefix=f"{prefix}.pre_projection",
-        )
-
-        self.post_projection = RowParallelLinear(
-            self.hidden_size,
-            self.backbone_hidden_size,
-            bias=False,
-            input_is_parallel=False,
-            prefix=f"{prefix}.post_projection",
-        )
-
-        self.layers = nn.ModuleList(
-            Gemma4MTPDecoderLayer(
-                text_config,
-                cache_config=vllm_config.cache_config,
-                quant_config=vllm_config.quant_config,
-                prefix=f"{prefix}.layers.{idx}",
-            )
-            for idx in range(self.num_mtp_layers)
-        )
-
-        self.norm = RMSNorm(self.hidden_size, eps=text_config.rms_norm_eps)
-
-        # After embedding sharing, embed_tokens is replaced with the
-        # target model's backbone-dim embedding.  Scale by
-        # sqrt(backbone_hidden_size) to match the target's convention.
-        self.register_buffer(
-            "normalizer",
-            torch.tensor(self.backbone_hidden_size**0.5),
-            persistent=False,
-        )
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids) * self.normalizer
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
-        params_dict = dict(self.named_parameters())
-        params_dict.update(dict(self.named_buffers()))
-        loaded_params: set[str] = set()
-
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-
-        return loaded_params
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        spec_step_idx: int = 0,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Returns (draft_hidden_states, backbone_hidden_states).
-
-        draft_hidden_states: draft-dim, used by compute_logits via lm_head.
-        backbone_hidden_states: backbone-dim, stored in the proposer's
-            hidden-state buffer and fed back as input to the next step.
-        """
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_input_ids(input_ids)
-
-        combined = torch.cat([inputs_embeds, hidden_states], dim=-1)
-        hidden_states, _ = self.pre_projection(combined)
-
-        residual = None
-        for layer in self.layers:
-            hidden_states, residual = layer(
-                positions=positions,
-                hidden_states=hidden_states,
-                residual=residual,
-            )
-
-        draft_hidden_states = self.norm(hidden_states)
-
-        backbone_hidden_states, _ = self.post_projection(draft_hidden_states)
-        return draft_hidden_states, backbone_hidden_states
-
-
-@support_torch_compile
-class Gemma4MTP(nn.Module):
-    """Gemma4 Multi-Token Prediction model for speculative decoding.
-
-    forward() returns (draft_hidden_states, backbone_hidden_states).
-    The proposer uses draft_hidden_states for compute_logits (via
-    the draft-dim lm_head) and backbone_hidden_states for the
-    hidden-state feedback buffer.
-    """
-
-    has_own_lm_head = True
-
-    hf_to_vllm_mapper = WeightsMapper(
-        orig_to_new_prefix={
-            "pre_projection.": "model.pre_projection.",
-            "post_projection.": "model.post_projection.",
-        },
-    )
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        config = vllm_config.speculative_config.draft_model_config.hf_config
-        text_config = _get_text_config(config)
-        self.config = config
-
-        self.model = Gemma4MultiTokenPredictor(
-            vllm_config=vllm_config,
-            prefix=maybe_prefix(prefix, "model"),
-        )
-
-        # lm_head operates in draft-dim.  Tied to embed_tokens at init
-        # so load_weights populates both from a single checkpoint entry.
-        # After embedding sharing, lm_head.weight still references the
-        # original draft-dim tensor.
-        self.lm_head = ParallelLMHead(
-            text_config.vocab_size,
-            text_config.hidden_size,
-            prefix=maybe_prefix(prefix, "lm_head"),
-        )
-        if getattr(config, "tie_word_embeddings", True):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-        self.logits_processor = LogitsProcessor(
-            text_config.vocab_size,
-            soft_cap=getattr(text_config, "final_logit_softcapping", None),
-        )
-
-        if getattr(config, "use_ordered_embeddings", False):
-            num_centroids = getattr(config, "num_centroids", 2048)
-            top_k = getattr(config, "centroid_intermediate_top_k", 32)
-            self.masked_embedding = Gemma4MTPMaskedEmbedder(
-                hidden_size=text_config.hidden_size,
-                vocab_size=text_config.vocab_size,
-                num_centroids=num_centroids,
-                centroid_intermediate_top_k=top_k,
-            )
-            logger.info(
-                "Gemma4 MTP: centroids masking enabled "
-                "(num_centroids=%d, top_k=%d, active_tokens=%d/%d).",
-                num_centroids,
-                top_k,
-                top_k * (text_config.vocab_size // num_centroids),
-                text_config.vocab_size,
-            )
-        else:
-            self.masked_embedding = None
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.embed_input_ids(input_ids)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        spec_step_idx: int = 0,
-        **kwargs: object,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        return self.model(
-            input_ids,
-            positions,
-            hidden_states,
-            intermediate_tensors,
-            inputs_embeds,
-            spec_step_idx,
-        )
-
-    def _get_full_lm_head_weight(self) -> torch.Tensor:
-        lm_head_weight = self.lm_head.weight
-        tp_size = get_tensor_model_parallel_world_size()
-        if tp_size > 1:
-            lm_head_weight = tensor_model_parallel_all_gather(
-                lm_head_weight,
-                dim=0,
-            )
-        return lm_head_weight[: self.masked_embedding.vocab_size]
-
-    def compute_logits(
-        self,
-        hidden_states: torch.Tensor,
-        spec_step_idx: int = 0,
-    ) -> torch.Tensor | None:
-        if self.masked_embedding is not None:
-            return self.masked_embedding(
-                hidden_states,
-                self._get_full_lm_head_weight(),
-            )
-        return self.logits_processor(self.lm_head, hidden_states)
-
-    def get_top_tokens(
-        self,
-        hidden_states: torch.Tensor,
-    ) -> torch.Tensor:
-        """Sparse argmax via centroids masking. Returns token IDs directly."""
-        return self.masked_embedding.get_top_tokens(
-            hidden_states,
-            self._get_full_lm_head_weight(),
-        )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
@@ -601,7 +601,6 @@ _SPECULATIVE_DECODING_MODELS = {
    "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
    "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
    "DeepSeekV4MTPModel": ("deepseek_v4_mtp", "DeepSeekV4MTP"),
-    "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"),
    "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
    "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),
    "Exaone4_5_MTP": ("exaone4_5_mtp", "Exaone4_5_MTP"),
@@ -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,
@@ -512,17 +512,6 @@ class LongCatFlashMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
        return getattr(self.hf_text_config, "num_nextn_predict_layers", 1)


-class Gemma4MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
-    def get_hidden_size(self) -> int:
-        # The speculator buffer must match the backbone (target) model's
-        # hidden dimension, not the draft model's smaller dimension.
-        return getattr(self.hf_config, "backbone_hidden_size",
-                       super().get_hidden_size())
-
-    def get_num_hidden_layers(self) -> int:
-        return getattr(self.hf_text_config, "num_hidden_layers", 0)
-
-
 class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
    def is_mm_prefix_lm(self) -> bool:
        return (
@@ -552,7 +541,6 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
    "falcon": FalconModelArchConfigConvertor,
    "gemma4": Gemma4ModelArchConfigConvertor,
    "gemma4_text": Gemma4ModelArchConfigConvertor,
-    "gemma4_mtp": Gemma4MTPModelArchConfigConvertor,
    "RefinedWeb": FalconModelArchConfigConvertor,
    "RefinedWebModel": FalconModelArchConfigConvertor,
    "nemotron-nas": NemotronNasModelArchConfigConvertor,
@@ -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"
@@ -1,335 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Gemma4 MTP (Multi-Token Prediction) proposer for speculative decoding.
-
-The Gemma4 assistant model runs all decoder layers per draft step
-(producing one token), and all its attention layers share KV cache
-with the target model via cross-model KV sharing.
-"""
-
-from collections import defaultdict
-from copy import copy
-
-import torch
-import torch.nn as nn
-
-from vllm.config import VllmConfig, get_layers_from_vllm_config, replace
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
-from vllm.v1.attention.backend import CommonAttentionMetadata
-from vllm.v1.kv_cache_interface import (
-    KVCacheConfig,
-    KVCacheSpec,
-    UniformTypeKVCacheSpecs,
-)
-from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
-from vllm.v1.worker.utils import AttentionGroup
-
-logger = init_logger(__name__)
-
-
-class Gemma4Proposer(SpecDecodeBaseProposer):
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        device: torch.device,
-        runner=None,
-    ):
-        super().__init__(
-            vllm_config,
-            device,
-            pass_hidden_states_to_model=True,
-            runner=runner,
-        )
-        # All draft steps predict from the same position (the last
-        # target-model position), so positions and seq_lens must not
-        # advance between steps.
-        self.constant_draft_positions = True
-
-        # Per-group block tables for multi-group KV cache models.
-        # Populated by gpu_model_runner during _prepare_inputs.
-        self._per_group_block_tables: dict[int, torch.Tensor] = {}
-
-        # Centroids CUDA graphs — populated in load_model if centroids
-        # masking is active. _centroids_sizes is pre-sorted for fast
-        # lookup in _greedy_sample.
-        self._centroids_sizes: list[int] = []
-        self._centroids_graphs: dict[int, torch.cuda.CUDAGraph] = {}
-        self._centroids_inputs: dict[int, torch.Tensor] = {}
-        self._centroids_outputs: dict[int, torch.Tensor] = {}
-
-    def set_per_group_block_table(self, gid: int, block_table: torch.Tensor) -> None:
-        self._per_group_block_tables[gid] = block_table
-
-    def model_returns_tuple(self) -> bool:
-        # forward() returns (draft_hidden_states, backbone_hidden_states).
-        # The proposer uses draft_hidden_states for compute_logits and
-        # backbone_hidden_states for the hidden-state feedback buffer.
-        return True
-
-    def build_per_group_and_layer_attn_metadata(
-        self,
-        common_attn_metadata: CommonAttentionMetadata,
-        draft_index: int = 0,
-    ) -> tuple[list[object], dict[str, object]]:
-        """Build attention metadata using the correct block table per group.
-
-        Gemma4 has multiple KV cache groups (sliding vs full attention)
-        with different block tables.  The base class receives a single
-        common_attn_metadata whose block_table belongs to one group.
-        We swap in the correct block table for each draft attention group.
-        """
-        per_group_attn_metadata: list[object] = []
-        per_layer_attn_metadata: dict[str, object] = {}
-        for attn_group in self.draft_attn_groups:
-            gid = attn_group.kv_cache_group_id
-            if gid in self._per_group_block_tables:
-                cm = copy(common_attn_metadata)
-                cm.block_table_tensor = self._per_group_block_tables[gid]
-            else:
-                cm = common_attn_metadata
-            attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
-                common_attn_metadata=cm, draft_index=draft_index
-            )
-            per_group_attn_metadata.append(attn_metadata)
-            for layer_name in attn_group.layer_names:
-                per_layer_attn_metadata[layer_name] = attn_metadata
-        return per_group_attn_metadata, per_layer_attn_metadata
-
-    def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self._centroids_sizes:
-            T = hidden_states.shape[0]
-            for size in self._centroids_sizes:
-                if size >= T:
-                    self._centroids_inputs[size][:T].copy_(hidden_states)
-                    self._centroids_graphs[size].replay()
-                    return self._centroids_outputs[size][:T].clone()
-            return self.model.get_top_tokens(hidden_states)
-        return super()._greedy_sample(hidden_states)
-
-    def _setup_centroids_cuda_graphs(self) -> None:
-        """Capture CUDA graphs for centroids get_top_tokens at key sizes."""
-        masked_emb = self.model.masked_embedding
-        lm_head_weight = self.model._get_full_lm_head_weight()
-
-        for size in [1, 2, 4, 8, 16, 32, 64]:
-            static_input = torch.zeros(
-                size,
-                masked_emb.hidden_size,
-                dtype=self.dtype,
-                device=self.device,
-            )
-            for _ in range(3):
-                masked_emb.get_top_tokens(static_input, lm_head_weight)
-            torch.cuda.synchronize()
-
-            g = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(g):
-                static_output = masked_emb.get_top_tokens(
-                    static_input,
-                    lm_head_weight,
-                )
-            self._centroids_graphs[size] = g
-            self._centroids_inputs[size] = static_input
-            self._centroids_outputs[size] = static_output
-
-        self._centroids_sizes = sorted(self._centroids_graphs)
-        logger.info(
-            "Gemma4 MTP: captured centroids CUDA graphs for sizes %s.",
-            self._centroids_sizes,
-        )
-
-    def _create_draft_vllm_config(self) -> VllmConfig:
-        """Preserve the target's forced TRITON_ATTN backend for draft layers.
-
-        Gemma4 forces TRITON_ATTN due to heterogeneous head dimensions
-        (head_dim=256 sliding, global_head_dim=512 full). The base class
-        resets attention_config.backend to None for draft models, causing
-        sliding layers to fall back to FLASH_ATTN which cannot handle
-        KV-shared cache. Override to carry the target's backend through.
-        """
-        base = super()._create_draft_vllm_config()
-        target_backend = self.vllm_config.attention_config.backend
-        if target_backend is not None:
-            base = replace(
-                base,
-                attention_config=replace(
-                    base.attention_config,
-                    backend=target_backend,
-                ),
-            )
-        return base
-
-    def _maybe_share_lm_head(self, target_language_model: nn.Module) -> None:
-        """Gemma4 MTP always keeps its own draft-dim lm_head.
-
-        The draft model's lm_head operates in draft hidden_size (e.g. 256),
-        which differs from the target's backbone hidden_size (e.g. 1536).
-        Sharing would break compute_logits (and centroids masking when
-        use_ordered_embeddings is enabled).
-        """
-        logger.info(
-            "Gemma4 MTP: keeping draft model's own lm_head (draft_dim != backbone_dim)."
-        )
-
-    def load_model(self, target_model: nn.Module) -> None:
-        target_attn_layer_names = set(
-            get_layers_from_vllm_config(
-                self.vllm_config,
-                AttentionLayerBase,
-            ).keys()
-        )
-
-        super().load_model(target_model)
-
-        self._setup_gemma4_kv_sharing(target_attn_layer_names)
-
-        if getattr(self.model, "masked_embedding", None) is not None:
-            self._setup_centroids_cuda_graphs()
-
-    def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
-        """Draft layers span multiple KV cache groups (sliding + full
-        attention with different head dimensions), so skip the base
-        class single-group assertion."""
-
-    def initialize_attn_backend(
-        self,
-        kv_cache_config: KVCacheConfig,
-        kernel_block_sizes: list[int] | None = None,
-    ) -> None:
-        """Create separate AttentionGroup objects per KV cache spec
-        so that each head-dim variant gets its own metadata builder."""
-        all_attn_layers = get_layers_from_vllm_config(
-            self.vllm_config,
-            AttentionLayerBase,
-        )
-
-        layer_to_gid: dict[str, int] = {}
-        layer_to_spec: dict[str, KVCacheSpec] = {}
-        for gid, group in enumerate(kv_cache_config.kv_cache_groups):
-            group_spec = group.kv_cache_spec
-            for ln in group.layer_names:
-                layer_to_gid[ln] = gid
-                if isinstance(group_spec, UniformTypeKVCacheSpecs):
-                    if ln in group_spec.kv_cache_specs:
-                        layer_to_spec[ln] = group_spec.kv_cache_specs[ln]
-                    else:
-                        tgt = getattr(
-                            all_attn_layers.get(ln),
-                            "kv_sharing_target_layer_name",
-                            None,
-                        )
-                        if tgt and tgt in group_spec.kv_cache_specs:
-                            layer_to_spec[ln] = group_spec.kv_cache_specs[tgt]
-                        else:
-                            layer_to_spec[ln] = group_spec
-                else:
-                    layer_to_spec[ln] = group_spec
-
-        attention_groups: dict[tuple[str, KVCacheSpec], AttentionGroup] = {}
-        for layer_name in self._draft_attn_layer_names:
-            if layer_name not in layer_to_spec:
-                continue
-            attn_layer = all_attn_layers[layer_name]
-            attn_backend = attn_layer.get_attn_backend()
-            spec = layer_to_spec[layer_name]
-            gid = layer_to_gid[layer_name]
-            group_key = (attn_backend.full_cls_name(), spec)
-
-            if group_key not in attention_groups:
-                kernel_block_size = (
-                    kernel_block_sizes[gid]
-                    if kernel_block_sizes is not None and gid < len(kernel_block_sizes)
-                    else None
-                )
-                attn_group = AttentionGroup(
-                    backend=attn_backend,
-                    layer_names=[layer_name],
-                    kv_cache_spec=spec,
-                    kv_cache_group_id=gid,
-                )
-                attn_group.create_metadata_builders(
-                    self.vllm_config,
-                    self.device,
-                    kernel_block_size=kernel_block_size,
-                )
-                attention_groups[group_key] = attn_group
-            else:
-                attention_groups[group_key].layer_names.append(layer_name)
-
-        self.draft_attn_groups = list(attention_groups.values())
-        if self.draft_attn_groups:
-            self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id
-            self.block_size = (
-                self.draft_attn_groups[0]
-                .get_metadata_builder()
-                .kv_cache_spec.block_size
-            )
-        else:
-            self.kv_cache_gid = 0
-            self.block_size = kv_cache_config.kv_cache_groups[
-                0
-            ].kv_cache_spec.block_size
-        logger.debug("Using block size %d for drafting layers", self.block_size)
-
-    def _setup_gemma4_kv_sharing(
-        self,
-        target_attn_layer_names: set[str],
-    ) -> None:
-        """Wire draft layers to share KV with the target model.
-
-        Each draft decoder layer is mapped to the last non-KV-shared
-        target layer of the same attention type (sliding or full).
-        """
-        draft_config = self.speculative_config.draft_model_config.hf_config
-        draft_text_config = draft_config.get_text_config()
-        target_config = self.vllm_config.model_config.hf_config
-        target_text_config = target_config.get_text_config()
-        target_layer_types = getattr(target_text_config, "layer_types", [])
-
-        if not (hasattr(self.model, "model") and hasattr(self.model.model, "layers")):
-            return
-
-        target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
-        num_non_shared = len(target_layer_types) - target_num_kv_shared
-        type_to_target_indices: dict[str, list[int]] = defaultdict(list)
-        for idx, lt in enumerate(target_layer_types[:num_non_shared]):
-            type_to_target_indices[lt].append(idx)
-
-        target_prefix = "model.layers"
-        for name in target_attn_layer_names:
-            if ".layers." in name:
-                target_prefix = name.split(".layers.")[0] + ".layers"
-                break
-
-        draft_layer_types = getattr(draft_text_config, "layer_types", [])
-        for draft_idx, layer in enumerate(self.model.model.layers):
-            if not hasattr(layer, "self_attn"):
-                continue
-            attn = getattr(layer.self_attn, "attn", None)
-            if attn is None:
-                continue
-
-            draft_layer_type = (
-                draft_layer_types[draft_idx]
-                if draft_idx < len(draft_layer_types)
-                else "full_attention"
-            )
-            candidates = type_to_target_indices.get(draft_layer_type, [])
-            if not candidates:
-                logger.warning(
-                    "No target layer of type '%s' for draft layer %d",
-                    draft_layer_type,
-                    draft_idx,
-                )
-                continue
-
-            target_idx = candidates[-1]
-            target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
-            attn.kv_sharing_target_layer_name = target_layer_name
-            logger.info(
-                "Gemma4 MTP: draft layer %d (%s) -> %s",
-                draft_idx,
-                draft_layer_type,
-                target_layer_name,
-            )
@@ -105,12 +105,6 @@ class SpecDecodeBaseProposer:
        )
        self.needs_extra_input_slots = self.net_num_new_slots_per_request > 0

-        # When True, all draft steps reuse the same position as the
-        # first step instead of advancing by one each iteration.
-        # Used by draft models with Q-only attention that share KV
-        # with the target and always predict from the same position.
-        self.constant_draft_positions: bool = False
-
        self.parallel_drafting_token_id: int = 0
        self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None
        if self.parallel_drafting:
@@ -394,9 +388,9 @@ class SpecDecodeBaseProposer:
        return {name: view for name in self._draft_attn_layer_names}

    def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
-        """Initialize cudagraph dispatcher keys for the drafter.
+        """Initialize cudagraph dispatcher keys for eagle.

-        Only supports PIECEWISE cudagraphs (via mixed_mode).
+        Eagle only supports PIECEWISE cudagraphs (via mixed_mode).
        This should be called after adjust_cudagraph_sizes_for_spec_decode.
        """
        if (
@@ -505,12 +499,6 @@ class SpecDecodeBaseProposer:
            positions = self.positions[token_indices_to_sample]
        hidden_states = hidden_states[token_indices_to_sample]

-        if self.constant_draft_positions:
-            # Write the sampling positions into the front of the
-            # positions buffer so that subsequent loop iterations
-            # (which read via _get_positions) use the correct values.
-            self.positions[:batch_size] = positions
-
        if any(isinstance(md, TreeAttentionMetadata) for md in per_group_attn_metadata):
            # Draft using tree attention - requires full logits for top-k
            logits = self.model.compute_logits(sample_hidden_states)
@@ -568,25 +556,59 @@ class SpecDecodeBaseProposer:
            # cast to int32 is crucial when eagle model is compiled.
            # tensor.argmax() returns int64 by default.
            input_ids = draft_token_ids_list[-1].int()
+            # Use fused kernel for slot mapping and metadata updates.
+            # Write clamped positions directly into the positions buffer to
+            # avoid an extra D2D copy for the common (non-mrope) case.
+            positions_1d = positions[0] if self.uses_mrope else positions
+            if self.uses_mrope:
+                out_pos = self.mrope_positions[0, :batch_size]
+            elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+                out_pos = self.xdrope_positions[0, :batch_size]
+            else:
+                out_pos = self.positions[:batch_size]
+            eagle_step_update_slot_mapping_and_metadata(
+                positions_1d=positions_1d,
+                block_table_tensor=common_attn_metadata.block_table_tensor,
+                seq_lens=common_attn_metadata.seq_lens,
+                block_size=block_size,
+                max_model_len=self.max_model_len,
+                out_clamped_positions=out_pos,
+                out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],
+                input_batch_size=input_batch_size,
+            )
+            common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
+            if self.uses_mrope:
+                self.mrope_positions[1:, :batch_size] = self.mrope_positions[
+                    0, :batch_size
+                ]
+                positions = self.mrope_positions[:, :batch_size]
+            elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+                self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[
+                    0, :batch_size
+                ]
+                positions = self.xdrope_positions[0, :batch_size]
+            else:
+                positions = self.positions[:batch_size]
+            # Increment the maximum sequence length. We increment max_seq_len
+            # unconditionally even though some seq_lens may have been capped above,
+            # as max_seq_len serves as an upper bound for sequence lengths.
+            common_attn_metadata.max_seq_len = min(
+                common_attn_metadata.max_seq_len + 1, self.max_model_len
+            )

-            if not self.constant_draft_positions:
-                positions = self._update_positions_dependent_metadata(
-                    positions,
-                    common_attn_metadata,
-                    batch_size,
-                    input_batch_size,
-                    block_size,
-                )
+            # Also update the CPU-side shadow; NOTE: this is hacky and should be
+            # removed in when common_attn_metadata.seq_lens_cpu is deprecated.
+            if common_attn_metadata._seq_lens_cpu is not None:
+                common_attn_metadata._seq_lens_cpu += 1
+            if common_attn_metadata._num_computed_tokens_cpu is not None:
+                common_attn_metadata._num_computed_tokens_cpu += 1
+            if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
+                common_attn_metadata.seq_lens_cpu_upper_bound += 1

-            # Rebuild attention metadata. When draft positions are constant
-            # (e.g. Gemma4 MTP), common_attn_metadata is invariant across
-            # loop iterations so we build once and reuse.
-            if not self.constant_draft_positions or token_index == 0:
-                _, per_layer_attn_metadata = (
-                    self.build_per_group_and_layer_attn_metadata(
-                        common_attn_metadata, draft_index=token_index + 1
-                    )
-                )
+            # Rebuild attention metadata
+            _, per_layer_attn_metadata = self.build_per_group_and_layer_attn_metadata(
+                common_attn_metadata, draft_index=token_index + 1
+            )

            # copy inputs to buffer for cudagraph
            self.input_ids[:batch_size] = input_ids
@@ -632,58 +654,6 @@ class SpecDecodeBaseProposer:
        draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
        return draft_token_ids

-    def _update_positions_dependent_metadata(
-        self,
-        positions: torch.Tensor,
-        common_attn_metadata,
-        batch_size: int,
-        input_batch_size: int,
-        block_size: int,
-    ) -> torch.Tensor:
-        """Update positions, slot mappings, and sequence metadata for the
-        next draft step. Returns the updated positions tensor."""
-        positions_1d = positions[0] if self.uses_mrope else positions
-        if self.uses_mrope:
-            out_pos = self.mrope_positions[0, :batch_size]
-        elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
-            out_pos = self.xdrope_positions[0, :batch_size]
-        else:
-            out_pos = self.positions[:batch_size]
-        eagle_step_update_slot_mapping_and_metadata(
-            positions_1d=positions_1d,
-            block_table_tensor=common_attn_metadata.block_table_tensor,
-            seq_lens=common_attn_metadata.seq_lens,
-            block_size=block_size,
-            max_model_len=self.max_model_len,
-            out_clamped_positions=out_pos,
-            out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],
-            input_batch_size=input_batch_size,
-        )
-        common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
-        if self.uses_mrope:
-            self.mrope_positions[1:, :batch_size] = self.mrope_positions[0, :batch_size]
-            positions = self.mrope_positions[:, :batch_size]
-        elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
-            self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[
-                0, :batch_size
-            ]
-            positions = self.xdrope_positions[0, :batch_size]
-        else:
-            positions = self.positions[:batch_size]
-        common_attn_metadata.max_seq_len = min(
-            common_attn_metadata.max_seq_len + 1,
-            self.max_model_len,
-        )
-
-        if common_attn_metadata._seq_lens_cpu is not None:
-            common_attn_metadata._seq_lens_cpu += 1
-        if common_attn_metadata._num_computed_tokens_cpu is not None:
-            common_attn_metadata._num_computed_tokens_cpu += 1
-        if common_attn_metadata.seq_lens_cpu_upper_bound is not None:
-            common_attn_metadata.seq_lens_cpu_upper_bound += 1
-
-        return positions
-
    def set_inputs_first_pass(
        self,
        target_token_ids: torch.Tensor,
@@ -169,7 +169,6 @@ from vllm.v1.spec_decode.dflash import DFlashProposer
 from vllm.v1.spec_decode.draft_model import DraftModelProposer
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
-from vllm.v1.spec_decode.gemma4 import Gemma4Proposer
 from vllm.v1.spec_decode.medusa import MedusaProposer
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.spec_decode.ngram_proposer_gpu import (
@@ -525,7 +524,6 @@ class GPUModelRunner(
                | DraftModelProposer
                | MedusaProposer
                | ExtractHiddenStatesProposer
-                | Gemma4Proposer
            )
            if self.speculative_config.method == "ngram":
                from vllm.v1.spec_decode.ngram_proposer import NgramProposer
@@ -554,8 +552,6 @@ class GPUModelRunner(
                self._ngram_pinned_val_buf = torch.zeros(
                    self.max_num_reqs, dtype=torch.int32, pin_memory=True
                )
-            elif self.speculative_config.use_gemma4_mtp():
-                self.drafter = Gemma4Proposer(self.vllm_config, self.device, self)
            elif self.speculative_config.use_dflash():
                self.drafter = DFlashProposer(self.vllm_config, self.device, self)
                self.use_aux_hidden_state_outputs = True
@@ -2314,18 +2310,11 @@ class GPUModelRunner(
                cm.slot_mapping = slot_mappings[kv_cache_gid]

            if self.speculative_config and spec_decode_common_attn_metadata is None:
-                if isinstance(
-                    self.drafter, (EagleProposer, DFlashProposer, Gemma4Proposer)
-                ):
+                if isinstance(self.drafter, (EagleProposer, DFlashProposer)):
                    if self.drafter.kv_cache_gid == kv_cache_gid:
                        spec_decode_common_attn_metadata = cm
                else:
                    spec_decode_common_attn_metadata = cm
-            # Capture per-group block tables for multi-group proposers.
-            if self.speculative_config and isinstance(self.drafter, Gemma4Proposer):
-                self.drafter.set_per_group_block_table(
-                    kv_cache_gid, cm.block_table_tensor
-                )

            for attn_gid in range(len(self.attn_groups[kv_cache_gid])):
                if ubatch_slices is not None:
@@ -4287,8 +4276,7 @@ class GPUModelRunner(
                    EagleProposer
                    | DFlashProposer
                    | DraftModelProposer
-                    | ExtractHiddenStatesProposer
-                    | Gemma4Proposer,
+                    | ExtractHiddenStatesProposer,
                )
                sampled_token_ids = sampler_output.sampled_token_ids
                if input_fits_in_drafter:
@@ -4684,8 +4672,7 @@ class GPUModelRunner(
            or spec_config.uses_draft_model()
        ):
            assert isinstance(
-                self.drafter,
-                EagleProposer | DFlashProposer | DraftModelProposer | Gemma4Proposer,
+                self.drafter, EagleProposer | DFlashProposer | DraftModelProposer
            )

            if spec_config.disable_padded_drafter_batch:
@@ -5607,8 +5594,7 @@ class GPUModelRunner(
                    EagleProposer
                    | DFlashProposer
                    | DraftModelProposer
-                    | ExtractHiddenStatesProposer
-                    | Gemma4Proposer,
+                    | ExtractHiddenStatesProposer,
                )
                assert self.speculative_config is not None
                # Eagle currently only supports PIECEWISE cudagraphs.
@@ -6409,8 +6395,7 @@ class GPUModelRunner(
            or self.speculative_config.uses_draft_model()
        ):
            assert isinstance(
-                self.drafter,
-                EagleProposer | DFlashProposer | DraftModelProposer | Gemma4Proposer,
+                self.drafter, EagleProposer | DFlashProposer | DraftModelProposer
            )
            self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes)

@@ -6463,10 +6448,7 @@ class GPUModelRunner(
        ):
            assert isinstance(
                self.drafter,
-                EagleProposer
-                | DFlashProposer
-                | ExtractHiddenStatesProposer
-                | Gemma4Proposer,
+                EagleProposer | DFlashProposer | ExtractHiddenStatesProposer,
            )
            self.drafter.initialize_cudagraph_keys(cudagraph_mode)
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