[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
69 changed files with 2390 additions and 497 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/
@@ -307,12 +307,12 @@ set(VLLM_EXT_SRC
  "csrc/quantization/activation_kernels.cu"
  "csrc/cuda_utils_kernels.cu"
  "csrc/custom_all_reduce.cu"
-  "csrc/torch_bindings.cpp")
+  "csrc/torch_bindings.cpp"
+  "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
+    "csrc/minimax_reduce_rms_kernel.cu")

  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

@@ -1047,13 +1047,13 @@ endif()
 set(VLLM_MOE_EXT_SRC
  "csrc/moe/torch_bindings.cpp"
  "csrc/moe/moe_align_sum_kernels.cu"
-  "csrc/moe/topk_softmax_kernels.cu")
+  "csrc/moe/topk_softmax_kernels.cu"
+  "csrc/moe/topk_softplus_sqrt_kernels.cu")

 if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC
    "csrc/moe/moe_wna16.cu"
-    "csrc/moe/grouped_topk_kernels.cu"
-    "csrc/moe/topk_softplus_sqrt_kernels.cu")
+    "csrc/moe/grouped_topk_kernels.cu")
 endif()

 if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -29,7 +29,11 @@
 */

 #include <cmath>
-#include <cuda_fp8.h>
+#ifndef USE_ROCM
+  #include <cuda_fp8.h>
+#else
+  #include <hip/hip_fp8.h>
+#endif
 #include <cuda_runtime.h>
 #include <type_traits>

@@ -42,7 +46,23 @@
 #include "type_convert.cuh"

 #ifndef FINAL_MASK
-  #define FINAL_MASK 0xffffffffu
+  #ifdef USE_ROCM
+    #define FINAL_MASK 0xffffffffffffffffULL
+  #else
+    #define FINAL_MASK 0xffffffffu
+  #endif
+#endif
+
+#ifdef USE_ROCM
+// ROCm-compatible FP8 conversion helpers
+__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
+  #if defined(HIP_FP8_TYPE_OCP)
+  __hip_fp8_e4m3 fp8_val(val);
+  #else
+  __hip_fp8_e4m3_fnuz fp8_val(val);
+  #endif
+  return reinterpret_cast<uint8_t&>(fp8_val);
+}
 #endif

 namespace vllm {
@@ -314,9 +334,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
      for (int i = 0; i < kElemsPerLane; i++) {
        float scaled = elements[i] * inv_scale;
        scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
+#ifndef USE_ROCM
        __nv_fp8_storage_t s =
            __nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
        out_bytes[i] = static_cast<uint8_t>(s);
+#else
+      out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
+#endif
      }
      // One 16-byte STG per lane.
      *reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
@@ -384,6 +408,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
  // PDL: enable programmatic stream serialization whenever the hardware
  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable,
  // so leave numAttrs = 0 and launch as a regular kernel.
+#ifndef USE_ROCM
  static int const sm_version = getSMVersion();
  // Host-side guard: the device kernel body is compiled as a no-op for
  // bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
@@ -410,6 +435,15 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
      q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
      num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
      kv_block_stride);
+#else
+  // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
+  // clang-format off
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
+      <<<grid, kBlockSize, 0, stream>>>(
+          q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+          eps, num_tokens_full, num_tokens_insert, num_heads_q,
+          cache_block_size, kv_block_stride);
+#endif
 }

 }  // namespace deepseek_v4_fused_ops
@@ -60,15 +60,6 @@ __device__ __forceinline__ float toFloat(T value) {
  }
 }

-#define FINAL_MASK 0xffffffff
-template <typename T>
-__inline__ __device__ T warpReduceSum(T val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1)
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  return val;
-}
-
 // ====================== TopK softplus_sqrt things
 // ===============================

@@ -272,8 +263,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
      }
    }
    // Compute per-thread scale (using warp reduction when renormalizing).
+    // THREADS_PER_ROW-parameterized butterfly works for both warp sizes (32
+    // on CUDA, 64 on ROCm CDNA) and any THREADS_PER_ROW the dispatch picks.
    if (renormalize) {
-      selected_sum = warpReduceSum(selected_sum);
+#pragma unroll
+      for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+        selected_sum +=
+            VLLM_SHFL_XOR_SYNC_WIDTH(selected_sum, mask, THREADS_PER_ROW);
+      }
    }
    float scale = static_cast<float>(routed_scaling_factor);
    if (renormalize) {
@@ -544,7 +541,6 @@ void topkGatingSoftplusSqrtKernelLauncher(
    const IndType* tid2eid, cudaStream_t stream) {
  static constexpr int WARPS_PER_TB = 4;
  static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
-#ifndef USE_ROCM
  // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
  // elements can be loaded by a warp
  static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
@@ -552,6 +548,19 @@ void topkGatingSoftplusSqrtKernelLauncher(
       std::is_same_v<InputType, __half>)
          ? 4
          : 8;
+  // Narrower LDG (ELTS_PER_LDG=1) used by 192/320/448/576 on ROCm WARP_SIZE=64
+  // where ELTS_PER_LDG=2 fails the EXPERTS%(ELTS_PER_LDG*WARP_SIZE)==0 check.
+  // On CUDA WARP_SIZE=32 the wider LDG already aligns, so the alias collapses
+  // back to BYTES_PER_LDG_MULTIPLE_64 — no behavioral change for CUDA.
+#ifdef USE_ROCM
+  static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
+      (std::is_same_v<InputType, __nv_bfloat16> ||
+       std::is_same_v<InputType, __half>)
+          ? 2
+          : 4;
+#else
+  static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
+      BYTES_PER_LDG_MULTIPLE_64;
 #endif
  switch (num_experts) {
    case 1:
@@ -584,27 +593,29 @@ void topkGatingSoftplusSqrtKernelLauncher(
    case 512:
      LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
      break;
-      // (CUDA only) support multiples of 64 when num_experts is not power of 2.
-      // ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
-      // num_experts, alternatively we can test 4 bytes loading and enable it in
-      // future.
-#ifndef USE_ROCM
+      // Multiples of 64 that are not powers of 2. The kernel requires
+      // EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0. With ELTS_PER_LDG=2
+      // (BYTES_PER_LDG_MULTIPLE_64), this holds for all five values on CUDA
+      // WARP_SIZE=32 but only for 384 on ROCm WARP_SIZE=64. The other four
+      // use BYTES_PER_LDG_MULTIPLE_64_NARROW (ELTS_PER_LDG=1), which
+      // satisfies the assertion for any multiple of 64 on either backend;
+      // on CUDA the narrow alias collapses back to the wider load, so CUDA
+      // behavior is unchanged.
    case 192:
-      LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
    case 320:
-      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
    case 384:
      LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
      break;
    case 448:
-      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
    case 576:
-      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
+      LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
      break;
-#endif
    default: {
      TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
    }
@@ -16,14 +16,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
      "bias) -> ()");
  m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);

-#ifndef USE_ROCM
  m.def(
      "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
      "token_expert_indices, Tensor gating_output, bool renormalize, float "
      "routed_scaling_factor, Tensor? "
      "bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
  m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
-#endif
+
  // Calculate the result of moe by summing up the partial results
  // from all selected experts.
  m.def("moe_sum(Tensor input, Tensor! output) -> ()");
@@ -183,7 +183,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "int forced_token_heads_per_warp=-1) -> ()");
  ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);

-#ifndef USE_ROCM
  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
  // kernel launch.
@@ -194,7 +193,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "float eps, int cache_block_size) -> ()");
  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
-#endif

  // Apply repetition penalties to logits in-place
  ops.def(
@@ -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
@@ -21,3 +21,6 @@ timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
+# tilelang has to be installed for mhc module to be
+# imported correctly.
+tilelang==0.1.9
@@ -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,
+    )
@@ -70,7 +70,8 @@ def test_sqrtsoftplus_bias_uses_deepseek_v4_routing_method():


@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+    not current_platform.is_cuda_alike(),
+    reason="This test is skipped on non-CUDA platform.",
 )
@pytest.mark.parametrize("num_tokens", [1, 33, 128])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@@ -125,7 +126,8 @@ def test_fused_topk_softplus_sqrt(


@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+    not current_platform.is_cuda_alike(),
+    reason="This test is skipped on non-CUDA platform.",
 )
@pytest.mark.parametrize("num_tokens", [1, 33, 128])
@pytest.mark.parametrize("hidden_size", [1024, 2048])
@@ -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
@@ -119,6 +119,7 @@ MoEBackend = Literal[
    "flashinfer_cutedsl",
    "marlin",
    "humming",
+    "triton_unfused",
    "aiter",
    "emulation",
 ]
@@ -150,6 +151,7 @@ class KernelConfig:
    - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
    - "marlin": Use Marlin kernels (weight-only quantization)
    - "humming": Use Humming Mixed Precision kernels
+    - "triton_unfused": Use Triton unfused MoE kernels
    - "aiter": Use AMD AITer kernels (ROCm only)
    - "emulation": use BF16/FP16 GEMM, dequantizing weights and
                   running QDQ on activations.
@@ -626,6 +626,7 @@ class SpeculativeConfig:
                    revision=self.revision,
                    code_revision=self.code_revision,
                    tokenizer_revision=self.target_model_config.tokenizer_revision,
+                    max_model_len=self.max_model_len,  # type: ignore[arg-type]
                    spec_target_max_model_len=self.target_model_config.max_model_len,
                    quantization=self.quantization,
                    enforce_eager=self.target_model_config.enforce_eager,
@@ -837,10 +838,17 @@ class SpeculativeConfig:

            return speculative_max_model_len

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

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

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

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

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

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

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

@@ -312,6 +312,21 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
        As: torch.Tensor,
        Bs: torch.Tensor,
    ) -> torch.Tensor:
+        if As.dtype != Bs.dtype:
+            from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+                _upcast_e8m0_to_fp32,
+            )
+
+            if As.dtype == torch.float8_e8m0fnu:
+                As = _upcast_e8m0_to_fp32(As).contiguous()
+            else:
+                As = As.to(torch.float32)
+
+            if Bs.dtype == torch.float8_e8m0fnu:
+                Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
+            else:
+                Bs = Bs.to(torch.float32)
+
        out_dtype = self.config.out_dtype
        if self.use_triton:
            gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
@@ -169,7 +169,9 @@ class SiluAndMulWithClamp(CustomOp):
    def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
        super().__init__(compile_native=compile_native)
        self.swiglu_limit = float(swiglu_limit)
-        if current_platform.is_cuda_alike() or current_platform.is_xpu():
+        if current_platform.is_rocm():
+            self._forward_method = self.forward_native
+        elif current_platform.is_cuda_alike() or current_platform.is_xpu():
            self.op = torch.ops._C.silu_and_mul_with_clamp
        elif current_platform.is_cpu():
            self._forward_method = self.forward_native
@@ -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(
@@ -300,6 +300,7 @@ class DeepseekCompressor(nn.Module):
        state_cache = self.state_cache.kv_cache
        # kv_state stored in first half, score_state stored in second half
        state_width = state_cache.shape[-1] // 2
+        pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}

        # Store the KV and score (with fused APE addition) in the state.
        # NOTE: PDL is disabled — both this kernel and _fused_kernel below
@@ -324,7 +325,7 @@ class DeepseekCompressor(nn.Module):
            TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]),
            STATE_WIDTH=state_width,
            COMPRESS_RATIO=self.compress_ratio,
-            launch_pdl=False,
+            **pdl_kwargs,
        )

        # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
@@ -373,7 +374,7 @@ class DeepseekCompressor(nn.Module):
            SCALE_DIM=self._scale_dim,
            KV_BLOCK_STRIDE=kv_cache.stride(0),
            num_warps=self._num_warps,
-            launch_pdl=False,
+            **pdl_kwargs,
        )


@@ -28,6 +28,11 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
    fused_inv_rope_fp8_quant,
    fused_q_kv_rmsnorm,
 )
+from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+    rocm_forward_decode_fallback,
+    rocm_inv_rope_einsum,
+    rocm_sparse_attn_prefill,
+)

 if TYPE_CHECKING:
    from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -53,6 +58,7 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import (
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
    GroupShape,
 )
+from vllm.platforms import current_platform
 from vllm.utils.multi_stream_utils import (
    execute_in_parallel,
    maybe_execute_in_parallel,
@@ -198,8 +204,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        # Pick fp8_einsum recipe based on GPU arch:
        # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128
        # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1
-        from vllm.platforms import current_platform
-
        cap = current_platform.get_device_capability()
        assert cap is not None, "DeepseekV4 attention requires a CUDA device"
        self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128)
@@ -222,6 +226,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
            + 1  # 1B pad
        )

+        # Will be None on ROCm for now.
        self.aux_stream_list = mla_modules.aux_stream_list
        # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
        # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins
@@ -303,6 +308,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        )
        o = o_padded[:, : self.n_local_heads, :]

+        # Keep ROCm on the BF16 reference wo_a path util kernel ready.
+        if current_platform.is_rocm():
+            z = rocm_inv_rope_einsum(
+                self.rotary_emb,
+                o,
+                positions,
+                self.rope_head_dim,
+                self.n_local_groups,
+                self.o_lora_rank,
+                self.wo_a,
+            )
+            return self.wo_b(z.flatten(1))
+
        # O projection: inverse RoPE + FP8 quant + einsum + wo_b
        o_fp8, o_scale = fused_inv_rope_fp8_quant(
            o,
@@ -336,12 +354,15 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        return self.wo_b(z.flatten(1))

    def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
-        assert self.aux_stream_list is not None
-        assert len(self.aux_stream_list) >= 3
+        aux_streams = self.aux_stream_list
+        if aux_streams is not None:
+            assert len(aux_streams) >= 3
+            aux_streams = aux_streams[:3]

        # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs
        # on aux streams 0..2 when their owning module exists. ln_events[0]
        # is the fan-out start event; ln_events[1..3] are per-aux done events.
+        # On ROCm, aux_streams is None and execute_in_parallel runs serially.
        aux_fns: list[Callable[[], Any] | None] = [None, None, None]

        if self.compressor is not None:
@@ -385,7 +406,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
            aux_fns,
            self.ln_events[0],
            self.ln_events[1:4],
-            self.aux_stream_list[:3],
+            aux_streams,
            enable=hidden_states.shape[0]
            <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
        )
@@ -419,8 +440,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        # downstream reads q on default). Indexer/compressor go on aux for
        # overlap with default's GEMM + cache write.
        if self.indexer is not None:
-            assert self.aux_stream_list is not None
-            aux_stream = self.aux_stream_list[0]
+            aux_stream = (
+                self.aux_stream_list[0] if self.aux_stream_list is not None else None
+            )
            indexer = self.indexer
            # Local ref so the closure keeps a non-None type for mypy.
            assert self.compressor is not None
@@ -448,8 +470,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
            )
        elif self.compressor is not None:
            # wq_b + kv_insert on default, compressor on aux.
-            assert self.aux_stream_list is not None
-            aux_stream = self.aux_stream_list[0]
+            aux_stream = (
+                self.aux_stream_list[0] if self.aux_stream_list is not None else None
+            )
            compressor = self.compressor

            def wq_b_kv_insert() -> torch.Tensor:
@@ -668,7 +691,7 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
            vllm_config.scheduler_config.max_num_batched_tokens
        )
        self.max_model_len = vllm_config.model_config.max_model_len
-        # DeepseekV4 only supports fp8 kv-cache format for now
+        # DeepseekV4 only supports fp8 kv-cache format for now.
        kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8"

        assert kv_cache_dtype.startswith("fp8"), (
@@ -816,6 +839,25 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
        swa_indices = swa_metadata.decode_swa_indices
        swa_lens = swa_metadata.decode_swa_lens

+        if current_platform.is_rocm():
+            rocm_forward_decode_fallback(
+                q=q,
+                kv_cache=kv_cache,
+                swa_k_cache=self.swa_cache_layer.kv_cache,
+                swa_only=swa_only,
+                topk_indices=topk_indices,
+                topk_lens=topk_lens,
+                swa_indices=swa_indices,
+                swa_lens=swa_lens,
+                attn_sink=self.attn_sink,
+                scale=self.scale,
+                head_dim=self.head_dim,
+                nope_head_dim=self.nope_head_dim,
+                rope_head_dim=self.rope_head_dim,
+                output=output,
+            )
+            return
+
        # We treat queries in the same seq as different queries
        # and later we only attend by generated indices.
        # q arrives pre-padded to self.padded_heads by the outer wrapper.
@@ -980,15 +1022,27 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
                N,
            )

-            output_chunk, _, _ = flash_mla_sparse_fwd(
-                q=q[query_start:query_end],
-                kv=kv.view(-1, 1, q.shape[-1]),
-                indices=combined_indices.unsqueeze(1),
-                sm_scale=self.scale,
-                attn_sink=self.attn_sink,
-                topk_length=combined_lens,
-                out=output[query_start:query_end],
-            )
+            if current_platform.is_rocm():
+                rocm_sparse_attn_prefill(
+                    q=q[query_start:query_end],
+                    kv=kv.view(-1, 1, q.shape[-1]),
+                    indices=combined_indices.unsqueeze(1),
+                    topk_length=combined_lens,
+                    scale=self.scale,
+                    head_dim=self.head_dim,
+                    attn_sink=self.attn_sink,
+                    output=output[query_start:query_end],
+                )
+            else:
+                output_chunk, _, _ = flash_mla_sparse_fwd(
+                    q=q[query_start:query_end],
+                    kv=kv.view(-1, 1, q.shape[-1]),
+                    indices=combined_indices.unsqueeze(1),
+                    sm_scale=self.scale,
+                    attn_sink=self.attn_sink,
+                    topk_length=combined_lens,
+                    out=output[query_start:query_end],
+                )


 class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
@@ -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
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
 from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
    FusedMoEQuantDesc,
+    RoutingMethodType,
    mxfp4_mxfp8_moe_quant_config,
    mxfp4_w4a8_moe_quant_config,
    mxfp4_w4a16_moe_quant_config,
@@ -64,6 +65,8 @@ class Mxfp4MoeBackend(Enum):
    MARLIN = "MARLIN"
    # ROCm AITER backends
    AITER_MXFP4_BF16 = "AITER_MXFP4_BF16"  # W4A16: CK kernel
+    # Keep the legacy name as an alias while the ROCm split backend rename settles.
+    AITER = "AITER_MXFP4_BF16"
    AITER_MXFP4_FP8 = "AITER_MXFP4_FP8"  # W4A8: triton kernel
    # Triton
    TRITON = "TRITON"
@@ -253,6 +256,8 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
    TRTLLM MXFP8; SM90 falls through to Triton_unfused or Marlin (the
    backend-level ``is_supported_config`` check filters by device capability).
    """
+    if current_platform.is_rocm():
+        return [Mxfp4MoeBackend.AITER_MXFP4_BF16]
    _AVAILABLE_BACKENDS = [
        Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
        Mxfp4MoeBackend.DEEPGEMM_MXFP4,
@@ -543,8 +548,22 @@ def select_deepseek_v4_mxfp4_moe_backend(
            activation_format,
        )

+    # DeepSeek-V4 on ROCm is more accurate with the unfused Triton MXFP4 path
+    # than the default AITER path. Prefer Triton-unfused for this routing mode,
+    # while keeping AITER as a fallback if Triton-unfused rejects the config.
+    if (
+        current_platform.is_rocm()
+        and config.routing_method == RoutingMethodType.DeepseekV4
+    ):
+        priority_backends = [
+            Mxfp4MoeBackend.TRITON_UNFUSED,
+            Mxfp4MoeBackend.AITER_MXFP4_BF16,
+        ]
+    else:
+        priority_backends = _get_priority_backends()
+
    # Iterate priority backends: TRTLLM MXFP8, then Triton.
-    for backend in _get_priority_backends():
+    for backend in priority_backends:
        activation_key = _backend_activation_key(backend)
        for k_cls in backend_to_kernel_cls(backend):
            supported, reason = k_cls.is_supported_config(
@@ -1252,6 +1271,64 @@ def convert_weight_to_mxfp4_moe_kernel_format(
            w2_bias,
        )

+    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        if w13_bias is not None:
+            w13_bias = w13_bias.data.to(torch.float32)
+        if w2_bias is not None:
+            w2_bias = w2_bias.data.to(torch.float32)
+
+        e, n, k = w13_weight.shape
+
+        w13_weight.view(torch.uint8).copy_(
+            w13_weight.data.view(torch.uint8)
+            .view(e, n // 2, 2, k)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+            .view(e, n, k)
+        )
+        w13_weight_scale.data = (
+            w13_weight_scale.data.view(e, n // 2, 2, -1)
+            .permute(0, 2, 1, 3)
+            .contiguous()
+            .view(e, n, -1)
+        )
+
+        w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2)
+        w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2)
+
+        w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True)
+        shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
+            w13_weight_scale.view(-1, w13_weight_scale.shape[-1]),
+            num_experts,
+            True,
+        )
+
+        w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False)
+        shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
+            w2_weight_scale.view(-1, w2_weight_scale.shape[-1]),
+            num_experts,
+            False,
+        )
+
+        if w13_bias is not None:
+            w13_bias = (
+                w13_bias.data.view(-1, n // 2, 2)
+                .permute(0, 2, 1)
+                .contiguous()
+                .view(-1, n)
+            )
+
+        return (
+            w13_weight,
+            w2_weight,
+            shuffled_w13_scale,
+            shuffled_w2_scale,
+            w13_bias,
+            w2_bias,
+        )
+
    elif mxfp4_backend in TRITON_BACKENDS:
        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig

@@ -1307,7 +1384,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
    else:
        raise ValueError(
            f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
-            f"Expected TRTLLM or Triton backend."
+            f"Expected TRTLLM, Triton, or AITER backend."
        )


@@ -268,10 +268,13 @@ class LinearBase(PluggableLayer):
        self.quant_config = quant_config
        self.prefix = prefix
        self.allow_fp8_block_shape_mismatch = False
+        self.quant_method: QuantizeMethodBase
        if quant_config is None:
-            self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod()
+            self.quant_method = UnquantizedLinearMethod()
+        elif quant_method := quant_config.get_quant_method(self, prefix=prefix):
+            self.quant_method = quant_method
        else:
-            self.quant_method = quant_config.get_quant_method(self, prefix=prefix)
+            raise ValueError("All linear layers should support quant method.")
        self.return_bias = return_bias
        self.disable_tp = disable_tp
        self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0
@@ -335,8 +338,6 @@ class ReplicatedLinear(LinearBase):
            disable_tp=disable_tp,
        )

-        # All the linear layer supports quant method.
-        assert self.quant_method is not None
        self.quant_method.create_weights(
            self,
            self.input_size,
@@ -389,7 +390,6 @@ class ReplicatedLinear(LinearBase):
        x: torch.Tensor,
    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
        bias = self.bias if not self.skip_bias_add else None
-        assert self.quant_method is not None

        output = self.quant_method.apply(self, x, bias)

@@ -474,7 +474,6 @@ class ColumnParallelLinear(LinearBase):
        self._maybe_allow_fp8_block_shape_mismatch()
        self.gather_output = gather_output

-        assert self.quant_method is not None
        self.quant_method.create_weights(
            layer=self,
            input_size_per_partition=self.input_size_per_partition,
@@ -583,7 +582,6 @@ class ColumnParallelLinear(LinearBase):
        bias = self.bias if not self.skip_bias_add else None

        # Matrix multiply.
-        assert self.quant_method is not None
        output_parallel = self.quant_method.apply(self, input_, bias)

        if self.gather_output and self.tp_size > 1:
@@ -1463,7 +1461,6 @@ class RowParallelLinear(LinearBase):
        self.input_is_parallel = input_is_parallel
        self.reduce_results = reduce_results

-        assert self.quant_method is not None
        self.quant_method.create_weights(
            layer=self,
            input_size_per_partition=self.input_size_per_partition,
@@ -1553,7 +1550,6 @@ class RowParallelLinear(LinearBase):
            input_parallel = split_input[self.tp_rank].contiguous()

        # Matrix multiply.
-        assert self.quant_method is not None
        # Only fuse bias add into GEMM for rank 0 (this ensures that
        # bias will not get added more than once in TP>1 case)
        bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
@@ -234,6 +234,39 @@ def mhc_pre(
    num_tokens = residual_flat.shape[0]
    fn_flat = fn

+    if current_platform.is_rocm():
+        x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32)
+        mixes = torch.matmul(x, fn_flat.t())
+        sqrsum = x.square().sum(dim=-1, keepdim=True)
+        mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
+
+        pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
+        pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps
+
+        post_logits = (
+            mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1]
+            + hc_base[hc_mult : 2 * hc_mult]
+        )
+        post_mix = torch.sigmoid(post_logits) * hc_post_mult_value
+
+        comb_logits = mixes[:, 2 * hc_mult :].view(
+            num_tokens, hc_mult, hc_mult
+        ) * hc_scale[2] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
+        comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
+        comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
+        for _ in range(sinkhorn_repeat - 1):
+            comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
+            comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
+
+        layer_input = torch.sum(
+            pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
+        ).to(torch.bfloat16)
+        return (
+            post_mix.view(*outer_shape, hc_mult, 1),
+            comb_mix.view(*outer_shape, hc_mult, hc_mult),
+            layer_input.view(*outer_shape, hidden_size),
+        )
+
    # these number are from deepgemm kernel impl
    block_k = 64
    block_m = 64
@@ -414,6 +447,14 @@ def mhc_post(
    post_layer_mix: torch.Tensor,
    comb_res_mix: torch.Tensor,
 ) -> torch.Tensor:
+    if current_platform.is_rocm():
+        mixed_residual = torch.einsum(
+            "...ij,...ih->...jh",
+            comb_res_mix.to(torch.float32),
+            residual.to(torch.float32),
+        )
+        post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32)
+        return (mixed_residual + post_term).to(residual.dtype)
    out = torch.empty_like(residual)
    mhc_post_tilelang(
        comb_res_mix,
@@ -551,6 +592,49 @@ def hc_head_fuse_tilelang(
        T.pdl_trigger()


+def _hc_head_fused_reference(
+    hs_flat: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    out: torch.Tensor,
+    hidden_size: int,
+    rms_eps: float,
+    hc_eps: float,
+    hc_mult: int,
+) -> None:
+    """Pure-PyTorch reference for `hc_head_fuse_tilelang`.
+
+    Used on platforms where the tilelang HIP/CUDA backend is not available
+    (e.g. ROCm builds shipping a tilelang wheel without `target.build.tilelang_hip`).
+    Mirrors the math of the tilelang kernel exactly:
+
+        x      = hs_flat.flatten(-2, -1)                # (T, hc_mult * H), fp32
+        mixes  = x @ fn.T                               # (T, hc_mult)
+        rsqrt  = 1 / sqrt(||x||^2 / (hc_mult * H) + rms_eps)
+        pre[m] = sigmoid(mixes[m] * rsqrt * hc_scale[0] + hc_base[m]) + hc_eps
+        out    = sum_m pre[m] * hs_flat[:, m, :]        # cast back to bf16
+
+    `out` is mutated in place to keep the same op contract
+    (`mutates_args=["out"]`).
+    """
+    num_tokens = hs_flat.shape[0]
+    if num_tokens == 0:
+        return
+    x = hs_flat.reshape(num_tokens, hc_mult * hidden_size).to(torch.float32)
+    # fn: (hc_mult, hc_mult * hidden_size) → mixes: (T, hc_mult)
+    mixes = torch.matmul(x, fn.t())
+    sqrsum = x.square().sum(dim=-1, keepdim=True)
+    rsqrt = torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
+    # hc_scale has shape (1,); hc_base has shape (hc_mult,)
+    pre_mix = torch.sigmoid(mixes * rsqrt * hc_scale[0] + hc_base) + hc_eps
+    # weighted sum over the hc_mult channel dim
+    result = torch.sum(pre_mix.unsqueeze(-1) * hs_flat.to(torch.float32), dim=1).to(
+        out.dtype
+    )
+    out.copy_(result)
+
+
 def _hc_head_fused_kernel(
    hs_flat: torch.Tensor,
    fn: torch.Tensor,
@@ -563,8 +647,15 @@ def _hc_head_fused_kernel(
    hc_mult: int,
 ) -> None:
    """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
-    if hs_flat.shape[0] > 0:
-        hc_head_fuse_tilelang(
+    if hs_flat.shape[0] == 0:
+        return
+    if current_platform.is_rocm():
+        # tilelang ships only the CUDA codegen in upstream wheels, so the HIP
+        # FFI target (`target.build.tilelang_hip`) is missing and the JIT call
+        # would raise `ValueError: Cannot find global function ...`. Use a
+        # numerically equivalent torch fallback instead. `mhc_pre` and
+        # `mhc_post` already follow this same pattern above.
+        _hc_head_fused_reference(
            hs_flat,
            fn,
            hc_scale,
@@ -575,6 +666,18 @@ def _hc_head_fused_kernel(
            hc_eps,
            hc_mult,
        )
+        return
+    hc_head_fuse_tilelang(
+        hs_flat,
+        fn,
+        hc_scale,
+        hc_base,
+        out,
+        hidden_size,
+        rms_eps,
+        hc_eps,
+        hc_mult,
+    )


 direct_register_custom_op(
@@ -843,6 +843,15 @@ def w8a8_triton_block_scaled_mm(
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]

+    # Triton cannot currently bind E8M0 scale tensors directly. On ROCm,
+    # DeepSeek-V4 checkpoints store block scales in exponent-only E8M0 format,
+    # so decode them to fp32 before launching the kernel.
+    if current_platform.is_rocm():
+        if As.dtype == torch.float8_e8m0fnu:
+            As = _upcast_e8m0_to_fp32(As).contiguous()
+        if Bs.dtype == torch.float8_e8m0fnu:
+            Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
+
    assert A.shape[-1] == B.shape[-1]
    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
@@ -499,13 +499,31 @@ class SparseAttnIndexer(CustomOp):
        k: torch.Tensor,
        weights: torch.Tensor,
    ):
-        assert not self.skip_k_cache_insert, (
-            "AMD platform doesn't support skip cache insert yet"
-        )
        assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet"
        assert isinstance(q_quant, torch.Tensor), (
            "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"
        )
+        if self.skip_k_cache_insert or not rocm_aiter_ops.is_enabled():
+            from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+                rocm_aiter_sparse_attn_indexer_native,
+            )
+
+            return rocm_aiter_sparse_attn_indexer_native(
+                hidden_states,
+                _encode_layer_name(self.k_cache.prefix),
+                self.k_cache.kv_cache,
+                q_quant,
+                k,
+                weights,
+                self.quant_block_size,
+                self.scale_fmt,
+                self.topk_tokens,
+                self.head_dim,
+                self.max_model_len,
+                self.max_total_seq_len,
+                self.topk_indices_buffer,
+                skip_k_cache_insert=self.skip_k_cache_insert,
+            )
        if rocm_aiter_ops.is_enabled():
            return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(
                hidden_states,
@@ -522,8 +540,4 @@ class SparseAttnIndexer(CustomOp):
                self.max_total_seq_len,
                self.topk_indices_buffer,
            )
-        else:
-            raise RuntimeError(
-                "Sparse attention indexer ROCm custom op requires ROCm "
-                "Aiter ops to be enabled."
-            )
+        raise RuntimeError("Sparse attention indexer ROCm path could not be selected.")
@@ -1245,7 +1245,12 @@ class DeepseekV4Model(nn.Module):
        # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute
        # (compressor kv_score, indexer.weights_proj, indexer.compressor
        # kv_score). fused_wqa_wkv stays on the default stream.
-        aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
+        # Disable them on ROCm because of hang issues.
+        aux_stream_list = (
+            None
+            if current_platform.is_rocm()
+            else [torch.cuda.Stream() for _ in range(3)]
+        )

        self.device = current_platform.device_type
        # Reserved topk indices buffer for all Indexer layers to reuse.
@@ -167,8 +167,12 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
        )

        # Three aux streams shared across all MTP layers, mirroring
-        # DeepseekV4Model.
-        aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
+        # DeepseekV4Model. ROCm runs the same work serially for now.
+        aux_stream_list = (
+            None
+            if current_platform.is_rocm()
+            else [torch.cuda.Stream() for _ in range(3)]
+        )

        # to map the exact layer index from weights
        self.layers = torch.nn.ModuleDict(
@@ -268,28 +268,6 @@ class InputProcessingContext:
        try:
            output = hf_processor(**data, **allowed_kwargs)
        except Exception as exc:
-            # See https://github.com/huggingface/tokenizers/issues/537
-            if (
-                isinstance(exc, RuntimeError)
-                and exc
-                and exc.args[0] == "Already borrowed"
-                and num_tries < max_tries
-            ):
-                logger.warning(
-                    "Failed to acquire tokenizer in current thread. "
-                    "Retrying (%d/%d)...",
-                    num_tries,
-                    max_tries,
-                )
-                time.sleep(0.5)
-                return self.call_hf_processor(
-                    hf_processor,
-                    data,
-                    kwargs,
-                    num_tries=num_tries + 1,
-                    max_tries=max_tries,
-                )
-
            msg = (
                f"Failed to apply {type(hf_processor).__name__} "
                f"on data={data} with kwargs={allowed_kwargs}"
@@ -409,6 +409,7 @@ class RocmPlatform(Platform):
        "gptq",
        "gptq_marlin",  # will be overwritten with gptq
        "fp8",
+        "deepseek_v4_fp8",
        "compressed-tensors",
        "fbgemm_fp8",
        "gguf",
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
-import copy
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
@@ -108,17 +107,10 @@ class BaseRenderer(ABC, Generic[_T]):
        if mm_registry.supports_multimodal_inputs(config.model_config):
            mm_processor_cache = mm_registry.processor_cache_from_config(config)

-            # Deep-copy the tokenizer so the multimodal processor gets its
-            # own Rust tokenizer backend.  Without this, concurrent access
-            # from AsyncMicrobatchTokenizer and call_hf_processor causes
-            # "RuntimeError: Already borrowed" from the Rust RefCell.
-            # See: https://github.com/huggingface/tokenizers/issues/537
-            mm_tokenizer = copy.deepcopy(tokenizer)
-
            with set_default_torch_num_threads():
                self.mm_processor = mm_registry.create_processor(
                    config.model_config,
-                    tokenizer=mm_tokenizer,
+                    tokenizer=self.tokenizer,
                    cache=mm_processor_cache,
                )

@@ -130,11 +122,10 @@ class BaseRenderer(ABC, Generic[_T]):
            # requests don't pollute the sender cache.
            ro_cache = mm_registry.processor_only_cache_from_config(config)
            if ro_cache is not None:
-                ro_tokenizer = copy.deepcopy(tokenizer)
                with set_default_torch_num_threads():
                    self._readonly_mm_processor = mm_registry.create_processor(
                        config.model_config,
-                        tokenizer=ro_tokenizer,
+                        tokenizer=self.tokenizer,
                        cache=ro_cache,
                    )

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from __future__ import annotations

+import copy
 import inspect
 import itertools
 import weakref
@@ -42,7 +43,7 @@ from vllm.multimodal.processing.processor import (
    apply_token_matches,
    find_mm_placeholders,
 )
-from vllm.tokenizers.hf import HfTokenizer
+from vllm.tokenizers.hf import HfTokenizer, maybe_make_thread_pool
 from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
 from vllm.transformers_utils.processor import cached_get_processor
 from vllm.utils.async_utils import make_async
@@ -785,6 +786,14 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
        config: VllmConfig,
        tokenizer: HfTokenizer | None,
    ) -> None:
+        # Ensure the og tokenizer is never modified by maybe_make_thread_pool
+        tokenizer = copy.copy(tokenizer)
+        if (
+            # Skip for mock configs and tokenizers
+            getattr(config.model_config, "enable_prompt_embeds", False)
+            and isinstance(tokenizer, HfTokenizer)
+        ):
+            _ensure_prompt_embeds_placeholder_token(tokenizer)
        super().__init__(config, tokenizer)

        self.use_unified_vision_chunk = getattr(
@@ -795,6 +804,11 @@ class HfRenderer(BaseRenderer[HfTokenizer]):
            safe_apply_chat_template, executor=self._executor
        )

+        if self.tokenizer is not None:
+            maybe_make_thread_pool(
+                self.tokenizer, config.model_config.renderer_num_workers + 1
+            )
+
    def render_messages(
        self,
        messages: list[ChatCompletionMessageParam],
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from .hf import maybe_make_thread_pool
 from .protocol import TokenizerLike
 from .registry import (
    TokenizerRegistry,
@@ -15,4 +16,5 @@ __all__ = [
    "cached_get_tokenizer",
    "get_tokenizer",
    "cached_tokenizer_from_config",
+    "maybe_make_thread_pool",
 ]
@@ -2,8 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
 import copy
+import queue
 from pathlib import Path
-from typing import TypeAlias
+from typing import TypeAlias, TypeVar

 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

@@ -12,6 +13,92 @@ from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_co
 from .protocol import TokenizerLike

 HfTokenizer: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast
+_T = TypeVar("_T", bound=TokenizerLike)
+
+
+class ThreadSafeHFTokenizerMixin:
+    """Mixin class for thread-safe HF fast tokenizers."""
+
+    pass
+
+
+def maybe_make_thread_pool(tokenizer: _T, copies: int = 1):
+    """
+    If `tokenizer` is a `PreTrainedTokenizerFast`, modify the tokenizer
+    in-place to make the public interface thread-safe by routing calls
+    through a deep-copied tokenizer pool.
+
+    Note that:
+    - Only ``TokenizerLike``'s public interface is thread-safe.
+      This doesn't include ``_tokenizer`` property nor any mutation
+      methods like ``add_special_tokens`` or ``add_tokens``.
+    - Adjacent method calls could happen on different deep copies.
+    """
+    if not isinstance(tokenizer, PreTrainedTokenizerFast) or isinstance(
+        tokenizer, ThreadSafeHFTokenizerMixin
+    ):
+        return tokenizer
+
+    og_tokenizer = copy.copy(tokenizer)
+
+    tokenizer_pool: queue.Queue[PreTrainedTokenizerFast] = queue.Queue()
+    for _ in range(copies):
+        tokenizer_pool.put(copy.deepcopy(og_tokenizer))
+
+    @contextlib.contextmanager
+    def _borrow_from_pool():
+        try:
+            tok = tokenizer_pool.get_nowait()
+            yield tok
+        except queue.Empty:
+            tok = copy.deepcopy(og_tokenizer)
+            yield tok
+        finally:
+            tokenizer_pool.put(tok)
+
+    class TokenizerPool(tokenizer.__class__, ThreadSafeHFTokenizerMixin):  # type: ignore
+        def apply_chat_template(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.apply_chat_template(*args, **kwargs)
+
+        def batch_decode(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.batch_decode(*args, **kwargs)
+
+        def batch_encode(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.batch_encode(*args, **kwargs)
+
+        def convert_tokens_to_ids(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.convert_tokens_to_ids(*args, **kwargs)
+
+        def convert_ids_to_tokens(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.convert_ids_to_tokens(*args, **kwargs)
+
+        def convert_tokens_to_string(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.convert_tokens_to_string(*args, **kwargs)
+
+        def decode(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.decode(*args, **kwargs)
+
+        def encode(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok.encode(*args, **kwargs)
+
+        def __call__(self, *args, **kwargs):
+            with _borrow_from_pool() as tok:
+                return tok(*args, **kwargs)
+
+        def __reduce__(self):
+            return maybe_make_thread_pool, (og_tokenizer, copies)
+
+    TokenizerPool.__name__ = f"TokenizerPool{og_tokenizer.__class__.__name__}"
+
+    tokenizer.__class__ = TokenizerPool


 def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
@@ -103,7 +190,10 @@ class CachedHfTokenizer(TokenizerLike):
                    "is a custom tokenizer not yet available in the "
                    "HuggingFace transformers library, consider "
                    "setting `trust_remote_code=True` in LLM or using "
-                    "the `--trust-remote-code` flag in the CLI."
+                    "the `--trust-remote-code` flag in the CLI. If the "
+                    "model was created with a newer version of "
+                    "transformers, consider upgrading: "
+                    "`uv pip install --upgrade transformers`"
                )
                raise RuntimeError(err_msg) from e
            else:
@@ -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,
@@ -3,8 +3,8 @@

 import json
 import os
-import platform
 import subprocess
+import sys
 from dataclasses import dataclass
 from functools import cache

@@ -78,7 +78,7 @@ def parse_id_list(raw_str: str) -> list[int]:


 def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
-    if platform.system() == "Darwin":
+    if sys.platform == "darwin":
        # MacOS has no memory node
        return MemoryNodeInfo(
            total_memory=psutil.virtual_memory().total,
@@ -122,17 +122,14 @@ def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:

 def get_allowed_cpu_list() -> list[LogicalCPUInfo]:
    cpu_list = _get_cpu_list()
-    if platform.system() == "Darwin":
-        return cpu_list
-
-    global_allowed_cpu_id_list = os.sched_getaffinity(0)  # type: ignore[attr-defined]
-    logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list]
-
-    return logical_cpu_list
+    if sys.platform == "linux":
+        allowed = os.sched_getaffinity(0)
+        return [x for x in cpu_list if x.id in allowed]
+    return cpu_list


 def get_visible_memory_node() -> list[int]:
-    if platform.system() == "Darwin":
+    if sys.platform == "darwin":
        return [0]

    allowed_memory_node_list = get_memory_affinity()
@@ -163,7 +160,7 @@ def _synthesize_cpu_list() -> list[LogicalCPUInfo]:


 def _get_cpu_list() -> list[LogicalCPUInfo]:
-    if platform.system() == "Darwin":
+    if sys.platform == "darwin":
        # For MacOS, no user-level CPU affinity and SMT, return all CPUs
        return _synthesize_cpu_list()

@@ -115,22 +115,29 @@ def get_flash_attn_version(
            )
            fa_version = 2

-        # The FA3 kernel rejects s_aux (sinks) when hdim != hdim_v; upgrade to
-        # FA4 on SM90 when available.
+        # Some FA3 unsupported SM90 cases can use FA4 when available.
        if (
            fa_version == 3
-            and has_sinks
-            and head_size is not None
-            and head_size_v is not None
-            and head_size != head_size_v
            and device_capability.major == 9
            and is_fa_version_supported(4)
        ):
-            logger.info_once(
-                "Diff-KV with sinks: upgrading FlashAttention 3 -> 4",
-                scope="local",
-            )
-            fa_version = 4
+            upgrade_reason = None
+            if head_size is not None and head_size > 256:
+                upgrade_reason = f"FA3 does not support head_size={head_size} on SM90"
+            elif (
+                has_sinks
+                and head_size is not None
+                and head_size_v is not None
+                and head_size != head_size_v
+            ):
+                upgrade_reason = "Diff-KV with sinks"
+            if upgrade_reason:
+                logger.info_once(
+                    "%s: upgrading FlashAttention 3 -> 4",
+                    upgrade_reason,
+                    scope="local",
+                )
+                fa_version = 4

        # FA4 currently uses batch-shape-dependent scheduling
        # heuristics on SM100+, which breaks batch invariance.
@@ -638,14 +638,6 @@ class FlashAttentionImpl(AttentionImpl):
            requires_alibi=alibi_slopes is not None,
            head_size=head_size,
        )
-        # head_size > 256 requires FA4 on SM90+; force upgrade from FA3
-        if (
-            head_size > 256
-            and self.vllm_flash_attn_version == 3
-            and current_platform.is_cuda()
-            and current_platform.is_device_capability_family(90)
-        ):
-            self.vllm_flash_attn_version = 4
        logger.info_once(
            "Using FlashAttention version %s",
            self.vllm_flash_attn_version,
@@ -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:
@@ -7,6 +7,7 @@ import torch

 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backend import (
    AttentionBackend,
@@ -360,7 +361,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
            _LAYER_TYPE_C4A: None,
            _LAYER_TYPE_C128A: None,
        }
-        if num_decode_tokens == 0:
+        if num_decode_tokens == 0 or current_platform.is_rocm():
            return out
        for layer_type in self._layer_types:
            # get_mla_metadata() is the official FlashMLA entry point that
@@ -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"
@@ -9,6 +9,7 @@ INT32-packed UE8M0 on SM100) so fp8_einsum skips transform_sf_into_required_layo

 import torch

+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import direct_register_custom_op

@@ -242,6 +243,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
        (scale_inner * tma_aligned_T, 1, tma_aligned_T),
    )
    grid = (tma_aligned_T, n_groups * heads_per_group)
+    pdl_kwargs = {} if current_platform.is_rocm() else {"launch_pdl": False}
    _fused_inv_rope_fp8_quant_per_head[grid](
        o,
        positions,
@@ -265,7 +267,7 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
        HALF_ROPE=half_rope,
        TMA_ALIGNED_SCALES=tma_aligned_scales,
        num_stages=1,
-        launch_pdl=False,
+        **pdl_kwargs,
        num_warps=1,
    )
    return fp8_buf, scale_buf
@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import functools
 import importlib
+import math
 from importlib.util import find_spec

 import torch
+import torch.nn.functional as F

 from vllm.forward_context import get_forward_context
 from vllm.platforms import current_platform
@@ -13,6 +15,11 @@ from vllm.utils.torch_utils import LayerNameType
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
 from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton

+if current_platform.is_rocm():
+    from vllm.platforms.rocm import _ON_GFX942
+else:
+    _ON_GFX942 = False
+

@triton.jit
 def _indexer_k_quant_and_cache_kernel(
@@ -230,6 +237,43 @@ def fp8_paged_mqa_logits_torch(

    fp8_dtype = current_platform.fp8_dtype()
    batch_size, next_n, _, dim = q.size()
+    if next_n == 1:
+        block_size = kv_cache.shape[1]
+        logits = torch.full(
+            [batch_size, max_model_len],
+            float("-inf"),
+            device=q.device,
+            dtype=torch.float32,
+        )
+        if context_lens.dim() > 1:
+            context_lens = context_lens.squeeze(-1)
+        kv_cache_flat = kv_cache.view(-1, block_size * (dim + 4))
+        for i in range(batch_size):
+            q_i = q[i, 0].to(torch.float32)
+            q_scale = weights[i]
+            seq_len = int(context_lens[i].item())
+            assert seq_len <= max_model_len
+            num_pages = cdiv(seq_len, block_size)
+            padded_seq_len = num_pages * block_size
+            pages = block_tables[i, :num_pages]
+            cache = kv_cache_flat[pages]
+            scale_offset = block_size * dim
+            cache_value = (
+                cache[..., :scale_offset].view(dtype=fp8_dtype).to(torch.float32)
+            )
+            cache_scale = (
+                cache[..., scale_offset:].view(dtype=torch.float32).contiguous()
+            )
+            cache_value = cache_value.view(padded_seq_len, dim)
+            cache_scale = cache_scale.view(padded_seq_len)
+            score = F.linear(cache_value, q_i)
+            score = F.relu(score)
+            score *= q_scale[None, :]
+            score = score.sum(dim=1)
+            score *= cache_scale
+            logits[i, :seq_len] = score[:seq_len]
+        return logits
+
    kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
    scale = scale.contiguous().view(torch.float)
    q = q.float()
@@ -241,20 +285,30 @@ def fp8_paged_mqa_logits_torch(
        device=q.device,
        dtype=torch.float32,
    )
-    context_lens = context_lens.tolist()
    for i in range(batch_size):
        context_len = context_lens[i]
-        q_offsets = torch.arange(context_len - next_n, context_len, device="cuda")
+        if context_len.ndim == 0:
+            context_len_i = int(context_len.item())
+            q_offsets = torch.arange(
+                context_len_i - next_n, context_len_i, device=q.device
+            )
+            context_limit = torch.full(
+                (next_n,), context_len_i, dtype=torch.int32, device=q.device
+            )
+        else:
+            context_limit = context_len.to(device=q.device, dtype=torch.int32)
+            q_offsets = context_limit - 1
        weight_slice = (
            weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
        )
-        for block_rk in range(cdiv(context_len, block_size)):
+        max_context_len = int(context_limit.max().item())
+        for block_rk in range(cdiv(max_context_len, block_size)):
            block_idx = block_tables[i][block_rk]
            qx, kx = q[i], kv_cache[block_idx]
            k_offsets = torch.arange(
-                block_rk * block_size, (block_rk + 1) * block_size, device="cuda"
+                block_rk * block_size, (block_rk + 1) * block_size, device=q.device
            )
-            mask = (k_offsets[None, :] < context_len) & (
+            mask = (k_offsets[None, :] < context_limit[:, None]) & (
                k_offsets[None, :] <= q_offsets[:, None]
            )
            s = torch.where(
@@ -331,30 +385,52 @@ def rocm_fp8_paged_mqa_logits(
        aiter_paged_mqa_logits_module = paged_mqa_logits_module()

    if aiter_paged_mqa_logits_module is not None:
-        deepgemm_fp8_paged_mqa_logits = (
-            aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
+        if _ON_GFX942:
+            deepgemm_fp8_paged_mqa_logits = (
+                aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits
+            )
+            batch_size, next_n, heads, _ = q_fp8.shape
+            out_logits = torch.full(
+                [batch_size * next_n, max_model_len],
+                float("-inf"),
+                device="cuda",
+                dtype=torch.float32,
+            )
+            deepgemm_fp8_paged_mqa_logits(
+                q_fp8,
+                kv_cache_fp8,
+                weights,
+                out_logits,
+                context_lens,
+                block_tables,
+                max_model_len,
+                ChunkK=256,
+                Preshuffle=block_size == 64,
+                KVBlockSize=block_size,
+                WavePerEU=2,
+            )
+            return out_logits
+        deepgemm_fp8_paged_mqa_logits_stage1 = (
+            aiter_paged_mqa_logits_module.deepgemm_fp8_paged_mqa_logits_stage1
        )
        batch_size, next_n, heads, _ = q_fp8.shape
-        out_logits = torch.full(
-            [batch_size * next_n, max_model_len],
+        out_qk = torch.full(
+            (heads, batch_size * next_n, max_model_len),
            float("-inf"),
            device="cuda",
            dtype=torch.float32,
        )
-        deepgemm_fp8_paged_mqa_logits(
+        deepgemm_fp8_paged_mqa_logits_stage1(
            q_fp8,
            kv_cache_fp8,
            weights,
-            out_logits,
+            out_qk,
            context_lens,
            block_tables,
            max_model_len,
-            ChunkK=256,
-            Preshuffle=block_size == 64,
-            KVBlockSize=block_size,
-            WavePerEU=2,
+            ChunkQ=heads,
        )
-        return out_logits
+        return out_qk.sum(dim=0)
    else:
        return fp8_paged_mqa_logits_torch(
            q_fp8, kv_cache_fp8, weights, context_lens, block_tables, max_model_len
@@ -464,6 +540,27 @@ def rocm_fp8_mqa_logits(
        return fp8_mqa_logits_torch(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke)


+def _topk_indices_torch(logits: torch.Tensor, topk_tokens: int) -> torch.Tensor:
+    k = min(topk_tokens, logits.shape[-1])
+    values, indices = torch.topk(logits, k=k, dim=-1)
+    indices = indices.to(torch.int32)
+    indices = torch.where(
+        values == float("-inf"),
+        torch.full_like(indices, -1, dtype=torch.int32),
+        indices,
+    )
+    if k == topk_tokens:
+        return indices
+    padded = torch.full(
+        (logits.shape[0], topk_tokens),
+        -1,
+        dtype=torch.int32,
+        device=logits.device,
+    )
+    padded[:, :k] = indices
+    return padded
+
+
 def rocm_aiter_sparse_attn_indexer_fake(
    hidden_states: torch.Tensor,
    k_cache_prefix: LayerNameType,
@@ -482,8 +579,9 @@ def rocm_aiter_sparse_attn_indexer_fake(
    # profile run
    # NOTE(Chen): create the max possible flattened_kv. So that
    # profile_run can get correct memory usage.
+    device = hidden_states.device if k is None else k.device
    _flattened_kv = torch.empty(
-        [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
+        [total_seq_lens, head_dim + 4], device=device, dtype=torch.uint8
    )
    fp8_dtype = current_platform.fp8_dtype()
    _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
@@ -491,7 +589,7 @@ def rocm_aiter_sparse_attn_indexer_fake(
    return topk_indices_buffer


-def rocm_aiter_sparse_attn_indexer(
+def rocm_aiter_sparse_attn_indexer_native(
    hidden_states: torch.Tensor,
    k_cache_prefix: LayerNameType,
    kv_cache: torch.Tensor,
@@ -505,10 +603,12 @@ def rocm_aiter_sparse_attn_indexer(
    max_model_len: int,
    total_seq_lens: int,
    topk_indices_buffer: torch.Tensor | None,
+    skip_k_cache_insert: bool = False,
 ) -> torch.Tensor:
    # careful! this will be None in dummy run
    attn_metadata = get_forward_context().attn_metadata
    fp8_dtype = current_platform.fp8_dtype()
+    from vllm import _custom_ops as ops
    from vllm.utils.torch_utils import _resolve_layer_name

    k_cache_prefix = _resolve_layer_name(k_cache_prefix)
@@ -537,19 +637,33 @@ def rocm_aiter_sparse_attn_indexer(
    has_decode = layer_attn_metadata.num_decodes > 0
    has_prefill = layer_attn_metadata.num_prefills > 0
    num_decode_tokens = layer_attn_metadata.num_decode_tokens
+    device = hidden_states.device if k is None else k.device

    # during speculative decoding, k may be padded to the CUDA graph batch
    # size while slot_mapping only covers actual tokens.
    num_tokens = slot_mapping.shape[0]
-    k = k[:num_tokens]
+    if k is not None:
+        k = k[:num_tokens]
+    elif not skip_k_cache_insert:
+        raise ValueError("k must be provided when skip_k_cache_insert is False")

-    indexer_k_quant_and_cache_triton(
-        k,
-        kv_cache,
-        slot_mapping,
-        quant_block_size,
-        scale_fmt,
-    )
+    if not skip_k_cache_insert:
+        if _ON_GFX942:
+            ops.indexer_k_quant_and_cache(
+                k,
+                kv_cache,
+                slot_mapping,
+                quant_block_size,
+                scale_fmt,
+            )
+        else:
+            indexer_k_quant_and_cache_triton(
+                k,
+                kv_cache,
+                slot_mapping,
+                quant_block_size,
+                scale_fmt,
+            )

    topk_indices_buffer[: hidden_states.shape[0]] = -1
    if has_prefill:
@@ -558,22 +672,31 @@ def rocm_aiter_sparse_attn_indexer(
        for chunk in prefill_metadata.chunks:
            k_fp8 = torch.empty(
                [chunk.total_seq_lens, head_dim],
-                device=k.device,
+                device=device,
                dtype=fp8_dtype,
            )
            k_scale = torch.empty(
                [chunk.total_seq_lens, 4],
-                device=k.device,
+                device=device,
                dtype=torch.uint8,
            )
-            cp_gather_indexer_k_quant_cache_triton(
-                kv_cache,
-                k_fp8,
-                k_scale,
-                chunk.block_table,
-                chunk.cu_seq_lens,
-                chunk.token_to_seq,
-            )
+            if _ON_GFX942:
+                ops.cp_gather_indexer_k_quant_cache(
+                    kv_cache,
+                    k_fp8,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                )
+            else:
+                cp_gather_indexer_k_quant_cache_triton(
+                    kv_cache,
+                    k_fp8,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                    token_to_seq=chunk.token_to_seq,
+                )

            logits = rocm_fp8_mqa_logits(
                q_fp8[chunk.token_start : chunk.token_end],
@@ -582,21 +705,10 @@ def rocm_aiter_sparse_attn_indexer(
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
            )
-            num_rows = logits.shape[0]
-            assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
            topk_indices = topk_indices_buffer[
                chunk.token_start : chunk.token_end, :topk_tokens
            ]
-            torch.ops._C.top_k_per_row_prefill(
-                logits,
-                chunk.cu_seqlen_ks,
-                chunk.cu_seqlen_ke,
-                topk_indices,
-                num_rows,
-                logits.stride(0),
-                logits.stride(1),
-                topk_tokens,
-            )
+            topk_indices.copy_(_topk_indices_torch(logits, topk_tokens))

    if has_decode:
        decode_metadata = layer_attn_metadata.decode
@@ -633,19 +745,8 @@ def rocm_aiter_sparse_attn_indexer(
            max_model_len=max_model_len,
        )

-        num_rows = logits.shape[0]
-        assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
        topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]
-        torch.ops._C.top_k_per_row_decode(
-            logits,
-            next_n,
-            decode_metadata.seq_lens,
-            topk_indices,
-            num_rows,
-            logits.stride(0),
-            logits.stride(1),
-            topk_tokens,
-        )
+        topk_indices.copy_(_topk_indices_torch(logits, topk_tokens)[:num_decode_tokens])

        if decode_metadata.requires_padding:
            # if padded, we need to unpack
@@ -659,3 +760,370 @@ def rocm_aiter_sparse_attn_indexer(
            )

    return topk_indices_buffer
+
+
+def rocm_aiter_sparse_attn_indexer(
+    hidden_states: torch.Tensor,
+    k_cache_prefix: LayerNameType,
+    kv_cache: torch.Tensor,
+    q_fp8: torch.Tensor,
+    k: torch.Tensor,
+    weights: torch.Tensor,
+    quant_block_size: int,
+    scale_fmt: str | None,
+    topk_tokens: int,
+    head_dim: int,
+    max_model_len: int,
+    total_seq_lens: int,
+    topk_indices_buffer: torch.Tensor | None,
+) -> torch.Tensor:
+    return rocm_aiter_sparse_attn_indexer_native(
+        hidden_states,
+        k_cache_prefix,
+        kv_cache,
+        q_fp8,
+        k,
+        weights,
+        quant_block_size,
+        scale_fmt,
+        topk_tokens,
+        head_dim,
+        max_model_len,
+        total_seq_lens,
+        topk_indices_buffer,
+        skip_k_cache_insert=False,
+    )
+
+
+def _decode_e8m0_scales(scale: torch.Tensor) -> torch.Tensor:
+    if scale.dtype == torch.float8_e8m0fnu:
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            _upcast_e8m0_to_fp32,
+        )
+
+        return _upcast_e8m0_to_fp32(scale).contiguous()
+    return scale.to(torch.float32)
+
+
+def _expand_2d_block_scales(
+    scale: torch.Tensor,
+    rows: int,
+    cols: int,
+) -> torch.Tensor:
+    scale = _decode_e8m0_scales(scale)
+    row_blocks, col_blocks = scale.shape[-2:]
+    row_block = math.ceil(rows / row_blocks)
+    col_block = math.ceil(cols / col_blocks)
+    scale = torch.repeat_interleave(scale, row_block, dim=-2)[..., :rows, :]
+    scale = torch.repeat_interleave(scale, col_block, dim=-1)[..., :, :cols]
+    return scale
+
+
+def _apply_gptj_inv_rope_ref(
+    x: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    rope_dim: int,
+) -> torch.Tensor:
+    if rope_dim == 0 or x.numel() == 0:
+        return x
+    half_rot = rope_dim // 2
+    nope_dim = x.shape[-1] - rope_dim
+    dtype = x.dtype
+    x = x.to(torch.float32)
+    cache = cos_sin_cache.index_select(0, positions.to(torch.long))
+    cos = cache[:, :half_rot].to(torch.float32)
+    sin = cache[:, half_rot : 2 * half_rot].to(torch.float32)
+    view_shape = (positions.shape[0],) + (1,) * (x.dim() - 2) + (half_rot,)
+    cos = cos.view(view_shape)
+    sin = sin.view(view_shape)
+    rope = x[..., nope_dim:]
+    y_even = rope[..., 0::2]
+    y_odd = rope[..., 1::2]
+    rope_out = torch.stack(
+        (y_even * cos + y_odd * sin, y_odd * cos - y_even * sin),
+        dim=-1,
+    ).flatten(-2)
+    x = x.clone()
+    x[..., nope_dim:] = rope_out
+    return x.to(dtype)
+
+
+def _apply_inv_rope_ref(
+    rotary_emb: torch.nn.Module,
+    x: torch.Tensor,
+    positions: torch.Tensor,
+    rope_dim: int,
+) -> torch.Tensor:
+    if hasattr(rotary_emb, "forward_native"):
+        try:
+            query, _ = rotary_emb.forward_native(
+                positions,
+                x.clone(),
+                None,
+                inverse=True,
+            )
+            return query
+        except TypeError:
+            pass
+    return _apply_gptj_inv_rope_ref(x, positions, rotary_emb.cos_sin_cache, rope_dim)
+
+
+def rocm_inv_rope_einsum(
+    rotary_emb: torch.nn.Module,
+    o: torch.Tensor,
+    positions: torch.Tensor,
+    rope_head_dim: int,
+    n_local_groups: int,
+    o_lora_rank: int,
+    wo_a: torch.nn.Module,
+) -> torch.Tensor:
+    """Reference inverse-RoPE + WO_A einsum path used on ROCm."""
+    o_ref = _apply_inv_rope_ref(rotary_emb, o, positions, rope_head_dim).to(
+        torch.bfloat16
+    )
+    o_ref = o_ref.view(o.shape[0], n_local_groups, -1)
+
+    hidden_dim = o_ref.shape[-1]
+    if hasattr(wo_a, "weight_scale_inv"):
+        wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
+            torch.float32
+        )
+        wo_a_scale = _expand_2d_block_scales(
+            wo_a.weight_scale_inv.view(
+                n_local_groups, -1, wo_a.weight_scale_inv.shape[-1]
+            ),
+            o_lora_rank,
+            hidden_dim,
+        )
+        wo_a_weight = (wo_a_weight * wo_a_scale).to(torch.bfloat16)
+    else:
+        wo_a_weight = wo_a.weight.view(n_local_groups, o_lora_rank, hidden_dim).to(
+            torch.bfloat16
+        )
+
+    return torch.einsum("tgd,grd->tgr", o_ref, wo_a_weight)
+
+
+def rocm_ref_sparse_attn_prefill(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    topk_length: torch.Tensor | None,
+    scale: float,
+    head_dim: int,
+    attn_sink: torch.Tensor | None,
+) -> torch.Tensor:
+    indices = indices.clone().squeeze(1)
+    s_q, h_q, d_qk = q.shape
+    topk = indices.shape[-1]
+    s_kv = kv.shape[0]
+    if topk_length is not None:
+        mask = torch.arange(topk, device=indices.device).unsqueeze(
+            0
+        ) >= topk_length.unsqueeze(1)
+        indices[mask] = -1
+    invalid_mask = (indices < 0) | (indices >= s_kv)
+    indices[invalid_mask] = 0
+
+    qf = q.float()
+    gathered_kv = kv.index_select(0, indices.flatten()).reshape(s_q, topk, d_qk).float()
+    scores = qf @ gathered_kv.transpose(1, 2)
+    scores *= scale
+    scores[invalid_mask.unsqueeze(1).expand_as(scores)] = float("-inf")
+
+    orig_lse = torch.logsumexp(scores, dim=-1)
+    lse_for_o = orig_lse
+    if attn_sink is not None:
+        lse_for_o = torch.logsumexp(
+            torch.stack(
+                [orig_lse, attn_sink[:h_q].view(1, h_q).expand_as(orig_lse)],
+                dim=0,
+            ),
+            dim=0,
+        )
+    lse_for_o = lse_for_o.clone()
+    lse_for_o[lse_for_o == float("-inf")] = float("+inf")
+    probs = torch.exp(scores - lse_for_o.unsqueeze(-1))
+    out = probs @ gathered_kv[..., :head_dim]
+    lonely_q_mask = orig_lse == float("-inf")
+    out[lonely_q_mask.unsqueeze(-1).expand_as(out)] = 0.0
+    return out.to(torch.bfloat16)
+
+
+def rocm_sparse_attn_prefill(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    topk_length: torch.Tensor | None,
+    scale: float,
+    head_dim: int,
+    attn_sink: torch.Tensor | None,
+    output: torch.Tensor,
+) -> None:
+    output_chunk = rocm_ref_sparse_attn_prefill(
+        q=q,
+        kv=kv,
+        indices=indices,
+        topk_length=topk_length,
+        scale=scale,
+        head_dim=head_dim,
+        attn_sink=attn_sink,
+    )
+    output.copy_(output_chunk.to(output.dtype))
+
+
+def rocm_dequantize_blocked_k_cache(
+    quant_k_cache: torch.Tensor,
+    head_dim: int,
+    nope_head_dim: int,
+    rope_head_dim: int,
+) -> torch.Tensor:
+    fp8_dtype = current_platform.fp8_dtype()
+    tile_size = 64
+    num_tiles = nope_head_dim // tile_size
+
+    num_blocks, block_size, _ = quant_k_cache.shape
+    quant_k_cache = quant_k_cache.view(num_blocks, -1)
+    input_nope_rope = quant_k_cache[
+        :, : block_size * (nope_head_dim + 2 * rope_head_dim)
+    ].view(num_blocks, block_size, nope_head_dim + 2 * rope_head_dim)
+    input_nope = input_nope_rope[:, :, :nope_head_dim].view(fp8_dtype)
+    input_rope = input_nope_rope[:, :, nope_head_dim:].view(torch.bfloat16)
+    input_scale = (
+        quant_k_cache[:, block_size * (nope_head_dim + 2 * rope_head_dim) :]
+        .view(num_blocks, block_size, 8)[:, :, :num_tiles]
+        .view(torch.float8_e8m0fnu)
+    )
+
+    result = torch.empty(
+        (num_blocks, block_size, 1, head_dim),
+        dtype=torch.bfloat16,
+        device=quant_k_cache.device,
+    )
+    result[..., nope_head_dim:] = input_rope.unsqueeze(2)
+    for tile_idx in range(num_tiles):
+        cur_nope = input_nope[
+            ..., tile_idx * tile_size : (tile_idx + 1) * tile_size
+        ].to(torch.bfloat16)
+        cur_scales = input_scale[:, :, tile_idx].to(torch.bfloat16).unsqueeze(-1)
+        result[..., tile_idx * tile_size : (tile_idx + 1) * tile_size] = (
+            cur_nope * cur_scales
+        ).unsqueeze(2)
+    return result
+
+
+def rocm_ref_sparse_attn_decode(
+    q: torch.Tensor,
+    blocked_k: torch.Tensor,
+    indices_in_kvcache: torch.Tensor,
+    topk_length: torch.Tensor | None,
+    scale: float,
+    head_dim: int,
+    attn_sink: torch.Tensor | None,
+    extra_blocked_k: torch.Tensor | None = None,
+    extra_indices_in_kvcache: torch.Tensor | None = None,
+    extra_topk_length: torch.Tensor | None = None,
+) -> torch.Tensor:
+    b, s_q, h_q, d_qk = q.shape
+
+    def process_scope(
+        cur_blocked_k: torch.Tensor,
+        cur_indices: torch.Tensor,
+        cur_topk_length: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cur_indices = cur_indices.reshape(b, s_q, -1)
+        topk = cur_indices.size(-1)
+        fixed_indices = torch.clamp_min(cur_indices, 0)
+        gathered_kv = (
+            cur_blocked_k.view(-1, d_qk)
+            .index_select(0, fixed_indices.view(-1))
+            .view(b, s_q, topk, d_qk)
+        )
+        invalid_mask = cur_indices == -1
+        if cur_topk_length is not None:
+            cur_topk_length = cur_topk_length.reshape(b)
+            invalid_mask |= torch.arange(0, topk, device=invalid_mask.device).view(
+                1, 1, topk
+            ) >= cur_topk_length.view(b, 1, 1)
+        return gathered_kv, invalid_mask
+
+    gathered_kv, invalid_mask = process_scope(
+        blocked_k, indices_in_kvcache, topk_length
+    )
+    if extra_blocked_k is not None:
+        assert extra_indices_in_kvcache is not None
+        gathered_kv1, invalid_mask1 = process_scope(
+            extra_blocked_k, extra_indices_in_kvcache, extra_topk_length
+        )
+        gathered_kv = torch.cat([gathered_kv, gathered_kv1], dim=2)
+        invalid_mask = torch.cat([invalid_mask, invalid_mask1], dim=2)
+
+    gathered_kv = gathered_kv.view(b * s_q, -1, d_qk).float()
+    gathered_kv[gathered_kv != gathered_kv] = 0.0
+    qf = q.float().view(b * s_q, h_q, d_qk)
+    attn_weight = qf @ gathered_kv.transpose(-1, -2)
+    attn_weight *= scale
+    attn_weight[
+        invalid_mask.view(b * s_q, 1, -1).expand(b * s_q, h_q, invalid_mask.size(-1))
+    ] = float("-inf")
+    lse = attn_weight.logsumexp(dim=-1)
+    attn_weight = torch.exp(attn_weight - lse.unsqueeze(-1))
+    output = attn_weight @ gathered_kv[..., :head_dim]
+    output = output.view(b, s_q, h_q, head_dim)
+    lse = lse.view(b, s_q, h_q)
+
+    if attn_sink is not None:
+        output *= (1.0 / (1.0 + torch.exp(attn_sink.view(1, 1, h_q) - lse))).unsqueeze(
+            -1
+        )
+
+    lonely_q_mask = lse == float("-inf")
+    output[lonely_q_mask.unsqueeze(-1).expand_as(output)] = 0.0
+    return output.squeeze(1).to(torch.bfloat16)
+
+
+def rocm_forward_decode_fallback(
+    q: torch.Tensor,
+    kv_cache: torch.Tensor | None,
+    swa_k_cache: torch.Tensor,
+    swa_only: bool,
+    topk_indices: torch.Tensor | None,
+    topk_lens: torch.Tensor | None,
+    swa_indices: torch.Tensor,
+    swa_lens: torch.Tensor,
+    attn_sink: torch.Tensor | None,
+    scale: float,
+    head_dim: int,
+    nope_head_dim: int,
+    rope_head_dim: int,
+    output: torch.Tensor,
+) -> None:
+    blocked_swa = rocm_dequantize_blocked_k_cache(
+        swa_k_cache,
+        head_dim=head_dim,
+        nope_head_dim=nope_head_dim,
+        rope_head_dim=rope_head_dim,
+    )
+    blocked_extra = None
+    if not swa_only:
+        assert kv_cache is not None
+        blocked_extra = rocm_dequantize_blocked_k_cache(
+            kv_cache,
+            head_dim=head_dim,
+            nope_head_dim=nope_head_dim,
+            rope_head_dim=rope_head_dim,
+        )
+    attn_out = rocm_ref_sparse_attn_decode(
+        q=q.unsqueeze(1),
+        blocked_k=blocked_swa,
+        indices_in_kvcache=swa_indices.unsqueeze(1),
+        topk_length=swa_lens,
+        scale=scale,
+        head_dim=head_dim,
+        attn_sink=attn_sink[: q.shape[1]] if attn_sink is not None else None,
+        extra_blocked_k=blocked_extra,
+        extra_indices_in_kvcache=topk_indices,
+        extra_topk_length=topk_lens,
+    )
+    output.copy_(attn_out.to(output.dtype))
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
Hexiang Wang GitHub ganyi tjtanaa tjtanaavllm	628c436301	[New Model][ROCm] Add AMD support for DeepSeek V4 (#40871 ) Signed-off-by: ganyi <ygan@amd.com> Signed-off-by: whx-sjtu <xiaowang990929@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: tjtanaavllm <tunjian.tan@amd.com> Co-authored-by: ganyi <ygan@amd.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaavllm <tunjian.tan@amd.com>	2026-05-05 08:55:37 -07:00
Canlin GuoandGitHub	2228fe6868	[Attention] Move FA3→FA4 upgrade into get_flash_attn_version() (#40815 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2026-05-05 15:43:03 +00:00
Harry MellorandGitHub	84bd8a3c1e	Remove unnecessary runtime asserts from linear layers (#41729 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-05-05 14:42:56 +00:00
Lidang JiangandGitHub	b786ec8e74	[Bugfix] Suggest upgrading Transformers for tokenizer class errors (#38099 ) Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>	2026-05-05 14:10:45 +00:00
yzong-rh GitHub wang.yuqi	20dcd984f9	[Bugfix] Fix `RuntimeError: Already borrowed` by adding thread-safe Hugging Face fast-tokenizer wrappers (#41181 ) Signed-off-by: Yifan Zong <yzong@redhat.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>	2026-05-05 14:04:01 +00:00
Martin HickeyandGitHub	6fca518157	[BugFix][MyPy]: Module has no attribute "sched_getaffinity" [attr-defined] (#41465 ) Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>	2026-05-05 13:20:37 +00:00