Merge branch 'main' into wentao-optimize-async-scheduling-copy

Signed-off-by: whyiug <whyiug@hotmail.com>

.buildkite/ci_config_intel.yaml

@@ -0,0 +1,23 @@
+name: vllm_intel_ci
+job_dirs:
+  - ".buildkite/intel_jobs"
+run_all_patterns:
+  - "docker/Dockerfile"
+  - "CMakeLists.txt"
+  - "requirements/common.txt"
+  - "requirements/xpu.txt"
+  - "requirements/build.txt"
+  - "requirements/test.txt"
+  - "setup.py"
+  - "csrc/"
+  - "cmake/"
+run_all_exclude_patterns:
+  - "docker/Dockerfile."
+  - "csrc/cpu/"
+  - "csrc/rocm/"
+  - "cmake/hipify.py"
+  - "cmake/cpu_extension.cmake"
+registries: public.ecr.aws/q9t5s3a7
+repositories:
+  main: "vllm-ci-test-repo"
+  premerge: "vllm-ci-test-repo"

.buildkite/image_build/image_build_xpu.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+set -e
+
+if [[ $# -lt 3 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+
+# authenticate with AWS ECR
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
+aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
+
+# skip build if image already exists
+if ! docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu &> /dev/null; then
+  echo "Image not found, proceeding with build..."
+else
+  echo "Image found"
+  exit 0
+fi
+
+# build
+docker build \
+  --file docker/Dockerfile.xpu \
+  --build-arg max_jobs=16 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu \
+  --progress plain .
+
+# push
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-xpu

.buildkite/intel_jobs/test-intel.yaml

@@ -0,0 +1,64 @@
+group: Intel
+steps:
+  - label: ":docker: Build XPU image"
+    soft_fail: true
+    depends_on: []
+    key: image-build-xpu
+    commands:
+      - bash -lc '.buildkite/image_build/image_build_xpu.sh "public.ecr.aws/q9t5s3a7" "vllm-ci-test-repo" "$BUILDKITE_COMMIT"'
+    env:
+      DOCKER_BUILDKIT: "1"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+  - label: "XPU example Test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml 
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'pip install tblib==3.1.0 &&
+        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager &&
+        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE &&
+        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp &&
+        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN &&
+        python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8 &&
+        python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192 &&
+        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 &&
+        python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel'
+  - label: "XPU V1 test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml 
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'cd tests &&
+        pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
+        pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
+        pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py &&
+        pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
+        pytest -v -s v1/structured_output &&
+        pytest -v -s v1/test_serial_utils.py &&
+        pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_tree_attention.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py &&
+        pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py'

.buildkite/release-pipeline.yaml

@@ -90,6 +90,14 @@ steps:
         env:
           DOCKER_BUILDKIT: "1"
 
+  - label: "Generate and upload wheel indices"
+    depends_on: "build-wheels"
+    allow_dependency_failure: true
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+
   - group: "Build release Docker images"
     key: "build-release-images"
     steps:
@@ -603,7 +611,7 @@ steps:
       - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
     env:
       S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm700"
+      VARIANT: "rocm721"
 
   # ROCm Job 6: Build ROCm Release Docker Image
   - label: ":docker: Build release image - x86_64 - ROCm"
@@ -673,6 +681,7 @@ steps:
   - label: "Publish nightly ROCm image to DockerHub"
     depends_on:
       - build-rocm-release-image
+    if: build.env("NIGHTLY") == "1"
     agents:
       queue: small_cpu_queue_release
     commands:

.buildkite/scripts/generate-and-upload-nightly-index.sh

@@ -0,0 +1,84 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Generate and upload wheel indices for all wheels in the commit directory.
+# This script should run once after all wheels have been built and uploaded.
+
+# ======== setup ========
+
+BUCKET="vllm-wheels"
+INDICES_OUTPUT_DIR="indices"
+DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
+PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
+SUBPATH=$BUILDKITE_COMMIT
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# detect if python3.12+ is available
+has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
+if [[ "$has_new_python" -eq 0 ]]; then
+    # use new python from docker
+    docker pull python:3-slim
+    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
+fi
+
+echo "Using python interpreter: $PYTHON"
+echo "Python version: $($PYTHON --version)"
+
+# ======== generate and upload indices ========
+
+# list all wheels in the commit directory
+echo "Existing wheels on S3:"
+aws s3 ls "$S3_COMMIT_PREFIX"
+obj_json="objects.json"
+aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
+mkdir -p "$INDICES_OUTPUT_DIR"
+
+# call script to generate indices for all existing wheels
+# these indices have relative paths that work as long as they are next to the wheel directory in s3
+# i.e., the wheels are always in s3://vllm-wheels/<commit>/
+# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
+alias_args=()
+if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
+    alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
+fi
+
+# HACK: we do not need regex module here, but it is required by pre-commit hook
+# To avoid any external dependency, we simply replace it back to the stdlib re module
+sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
+$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
+
+# copy indices to /<commit>/ unconditionally
+echo "Uploading indices to $S3_COMMIT_PREFIX"
+aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
+
+# copy to /nightly/ only if it is on the main branch and not a PR
+if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
+    echo "Uploading indices to overwrite /nightly/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
+fi
+
+# detect version from any wheel in the commit directory
+# download the first wheel we find to extract version metadata
+first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
+if [[ -z "$first_wheel_key" ]]; then
+    echo "Error: No wheels found in $S3_COMMIT_PREFIX"
+    exit 1
+fi
+first_wheel=$(basename "$first_wheel_key")
+aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
+version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+rm -f "/tmp/${first_wheel}"
+echo "Version in wheel: $version"
+pure_version="${version%%+*}"
+echo "Pure version (without variant): $pure_version"
+
+# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
+if [[ "$version" != *"dev"* ]]; then
+    echo "Re-generating indices for /$pure_version/"
+    rm -rf "${INDICES_OUTPUT_DIR:?}"
+    mkdir -p "$INDICES_OUTPUT_DIR"
+    # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
+    $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
+fi

.buildkite/scripts/hardware_ci/run-intel-test.sh

@@ -0,0 +1,276 @@
+#!/bin/bash
+
+# This script runs tests inside the Intel XPU docker container.
+# It mirrors the structure of run-amd-test.sh while keeping Intel-specific
+# container setup and allowing commands to be sourced from YAML or env.
+#
+# Command sources (in priority order):
+#   1) VLLM_TEST_COMMANDS env var (preferred, preserves quoting)
+#   2) Positional args (legacy)
+#   3) One or more YAML files with a commands list (test-area style)
+###############################################################################
+set -o pipefail
+
+DRY_RUN=${DRY_RUN:-0}
+if [[ "${1:-}" == "--dry-run" ]]; then
+  DRY_RUN=1
+  shift
+fi
+
+# Export Python path
+export PYTHONPATH=".."
+
+###############################################################################
+# Helper Functions
+###############################################################################
+
+cleanup_docker() {
+  docker_root=$(docker info -f '{{.DockerRootDir}}')
+  if [ -z "$docker_root" ]; then
+    echo "Failed to determine Docker root directory." >&2
+    exit 1
+  fi
+  echo "Docker root directory: $docker_root"
+
+  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
+  threshold=70
+  if [ "$disk_usage" -gt "$threshold" ]; then
+    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    docker image prune -f
+    docker volume prune -f && docker system prune --force --filter "until=72h" --all
+    echo "Docker images and volumes cleanup completed."
+  else
+    echo "Disk usage is below $threshold%. No cleanup needed."
+  fi
+}
+
+re_quote_pytest_markers() {
+  local input="$1"
+  local output=""
+  local collecting=false
+  local marker_buf=""
+
+  local flat="${input//$'\n'/ }"
+  local restore_glob
+  restore_glob="$(shopt -p -o noglob 2>/dev/null || true)"
+  set -o noglob
+  local -a words
+  read -ra words <<< "$flat"
+  eval "$restore_glob"
+
+  for word in "${words[@]}"; do
+    if $collecting; then
+      if [[ "$word" == *"'"* ]]; then
+        if [[ -n "$marker_buf" ]]; then
+          output+="${marker_buf} "
+          marker_buf=""
+        fi
+        output+="${word} "
+        collecting=false
+        continue
+      fi
+
+      local is_boundary=false
+      case "$word" in
+        "&&"|"||"|";"|"|")
+          is_boundary=true ;;
+        --*)
+          is_boundary=true ;;
+        -[a-zA-Z])
+          is_boundary=true ;;
+        */*)
+          is_boundary=true ;;
+        *.py|*.py::*)
+          is_boundary=true ;;
+        *=*)
+          if [[ "$word" =~ ^[A-Z_][A-Z0-9_]*= ]]; then
+            is_boundary=true
+          fi
+          ;;
+      esac
+
+      if $is_boundary; then
+        if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
+          output+="'${marker_buf}' "
+        else
+          output+="${marker_buf} "
+        fi
+        collecting=false
+        marker_buf=""
+        if [[ "$word" == "-m" || "$word" == "-k" ]]; then
+          output+="${word} "
+          collecting=true
+        else
+          output+="${word} "
+        fi
+      else
+        if [[ -n "$marker_buf" ]]; then
+          marker_buf+=" ${word}"
+        else
+          marker_buf="${word}"
+        fi
+      fi
+    elif [[ "$word" == "-m" || "$word" == "-k" ]]; then
+      output+="${word} "
+      collecting=true
+      marker_buf=""
+    else
+      output+="${word} "
+    fi
+  done
+
+  if $collecting && [[ -n "$marker_buf" ]]; then
+    if [[ "$marker_buf" == *" "* || "$marker_buf" == *"("* ]]; then
+      output+="'${marker_buf}'"
+    else
+      output+="${marker_buf}"
+    fi
+  fi
+
+  echo "${output% }"
+}
+
+apply_intel_test_overrides() {
+  local cmds="$1"
+  # Placeholder for Intel-specific exclusions/overrides.
+  echo "$cmds"
+}
+
+is_yaml_file() {
+  local p="$1"
+  [[ -f "$p" && "$p" == *.yaml ]]
+}
+
+extract_yaml_commands() {
+  local yaml_path="$1"
+  awk '
+    $1 == "commands:" { in_cmds=1; next }
+    in_cmds && $0 ~ /^[[:space:]]*-[[:space:]]/ {
+      sub(/^[[:space:]]*-[[:space:]]/, "");
+      print;
+      next
+    }
+    in_cmds && $0 ~ /^[^[:space:]]/ { exit }
+  ' "$yaml_path"
+}
+
+###############################################################################
+# Main
+###############################################################################
+
+default_image_name="${REGISTRY}/${REPO}:${BUILDKITE_COMMIT}-xpu"
+#default_image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-xpu"
+image_name="${IMAGE_TAG_XPU:-${default_image_name}}"
+container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+# ---- Command source selection ----
+commands=""
+if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
+  commands="${VLLM_TEST_COMMANDS}"
+  echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
+elif [[ $# -gt 0 ]]; then
+  all_yaml=true
+  for arg in "$@"; do
+    if ! is_yaml_file "$arg"; then
+      all_yaml=false
+      break
+    fi
+  done
+
+  if $all_yaml; then
+    for yaml in "$@"; do
+      mapfile -t COMMANDS < <(extract_yaml_commands "$yaml")
+      if [[ ${#COMMANDS[@]} -eq 0 ]]; then
+        echo "Error: No commands found in ${yaml}" >&2
+        exit 1
+      fi
+      for cmd in "${COMMANDS[@]}"; do
+        if [[ -z "$commands" ]]; then
+          commands="${cmd}"
+        else
+          commands+=" && ${cmd}"
+        fi
+      done
+    done
+    echo "Commands sourced from YAML files: $*"
+  else
+    commands="$*"
+    echo "Commands sourced from positional args (legacy mode)"
+  fi
+else
+  SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+  DEFAULT_YAML="${SCRIPT_DIR}/intel-test.yaml"
+  if [[ ! -f "${DEFAULT_YAML}" ]]; then
+    echo "Error: YAML file not found: ${DEFAULT_YAML}" >&2
+    exit 1
+  fi
+  mapfile -t COMMANDS < <(extract_yaml_commands "${DEFAULT_YAML}")
+  if [[ ${#COMMANDS[@]} -eq 0 ]]; then
+    echo "Error: No commands found in ${DEFAULT_YAML}" >&2
+    exit 1
+  fi
+  for cmd in "${COMMANDS[@]}"; do
+    if [[ -z "$commands" ]]; then
+      commands="${cmd}"
+    else
+      commands+=" && ${cmd}"
+    fi
+  done
+  echo "Commands sourced from default YAML: ${DEFAULT_YAML}"
+fi
+
+if [[ -z "$commands" ]]; then
+  echo "Error: No test commands provided." >&2
+  exit 1
+fi
+
+echo "Raw commands: $commands"
+commands=$(re_quote_pytest_markers "$commands")
+echo "After re-quoting: $commands"
+commands=$(apply_intel_test_overrides "$commands")
+echo "Final commands: $commands"
+
+# Dry-run mode prints final commands and exits before Docker.
+if [[ "$DRY_RUN" == "1" ]]; then
+  echo "DRY_RUN=1 set, skipping Docker execution."
+  exit 0
+fi
+
+# --- Docker housekeeping ---
+cleanup_docker
+
+# --- Build or pull test image ---
+if [[ -n "${IMAGE_TAG_XPU:-}" ]]; then
+  echo "Using prebuilt XPU image: ${IMAGE_TAG_XPU}"
+  docker pull "${IMAGE_TAG_XPU}"
+else
+  echo "Using prebuilt XPU image: ${image_name}"
+  docker pull "${image_name}"
+fi
+
+remove_docker_container() {
+  docker rm -f "${container_name}" || true
+  docker image rm -f "${image_name}" || true
+  docker system prune -f || true
+}
+trap remove_docker_container EXIT
+
+# --- Single-node job ---
+
+if [[ -z "${ZE_AFFINITY_MASK:-}" ]]; then
+  echo "Warning: ZE_AFFINITY_MASK is not set. Proceeding without device affinity." >&2
+fi
+
+docker run \
+    --device /dev/dri:/dev/dri \
+    --net=host \
+    --ipc=host \
+    --privileged \
+    -v /dev/dri/by-path:/dev/dri/by-path \
+    --entrypoint="" \
+    -e "HF_TOKEN=${HF_TOKEN:-}" \
+    -e "ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-}" \
+    -e "CMDS=${commands}" \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c 'set -e; echo "ZE_AFFINITY_MASK is ${ZE_AFFINITY_MASK:-}"; eval "$CMDS"'

.buildkite/scripts/upload-nightly-wheels.sh

@@ -2,27 +2,14 @@
 
 set -ex
 
-# ======== part 0: setup ========
+# Upload a single wheel to S3 (rename linux -> manylinux).
+# Index generation is handled separately by generate-and-upload-nightly-index.sh.
 
 BUCKET="vllm-wheels"
-INDICES_OUTPUT_DIR="indices"
-DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
-PYTHON=${PYTHON_PROG:=python3} # try to read from env var, otherwise use python3
 SUBPATH=$BUILDKITE_COMMIT
 S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
 
-# detect if python3.10+ is available
-has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
-if [[ "$has_new_python" -eq 0 ]]; then
-    # use new python from docker
-    docker pull python:3-slim
-    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
-fi
-
-echo "Using python interpreter: $PYTHON"
-echo "Python version: $($PYTHON --version)"
-
-# ========= part 1: collect, rename & upload the wheel ==========
+# ========= collect, rename & upload the wheel ==========
 
 # Assume wheels are in artifacts/dist/*.whl
 wheel_files=(artifacts/dist/*.whl)
@@ -52,56 +39,8 @@ echo "Renamed wheel to: $wheel"
 # Extract the version from the wheel
 version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
 echo "Version in wheel: $version"
-pure_version="${version%%+*}"
-echo "Pure version (without variant): $pure_version"
 
 # copy wheel to its own bucket
 aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
 
-# ========= part 2: generate and upload indices ==========
-# generate indices for all existing wheels in the commit directory
-# this script might be run multiple times if there are multiple variants being built
-# so we need to guarantee there is little chance for "TOCTOU" issues
-# i.e., one process is generating indices while another is uploading a new wheel
-# so we need to ensure no time-consuming operations happen below
-
-# list all wheels in the commit directory
-echo "Existing wheels on S3:"
-aws s3 ls "$S3_COMMIT_PREFIX"
-obj_json="objects.json"
-aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
-mkdir -p "$INDICES_OUTPUT_DIR"
-
-# call script to generate indices for all existing wheels
-# this indices have relative paths that could work as long as it is next to the wheel directory in s3
-# i.e., the wheels are always in s3://vllm-wheels/<commit>/
-# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
-alias_args=()
-if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
-    alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
-fi
-
-# HACK: we do not need regex module here, but it is required by pre-commit hook
-# To avoid any external dependency, we simply replace it back to the stdlib re module
-sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
-$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
-
-# copy indices to /<commit>/ unconditionally
-echo "Uploading indices to $S3_COMMIT_PREFIX"
-aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
-
-# copy to /nightly/ only if it is on the main branch and not a PR 
-if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
-    echo "Uploading indices to overwrite /nightly/"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
-fi
-
-# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
-if [[ "$version" != *"dev"* ]]; then
-    echo "Re-generating indices for /$pure_version/"
-    rm -rf "${INDICES_OUTPUT_DIR:?}/*"
-    mkdir -p "$INDICES_OUTPUT_DIR"
-    # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
-    $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
-fi
+echo "Wheel uploaded. Index generation is handled by a separate step."

.buildkite/test_areas/expert_parallelism.yaml

@@ -8,8 +8,10 @@ steps:
   source_file_dependencies:
   - vllm/distributed/eplb
   - tests/distributed/test_eplb_algo.py
+  - tests/distributed/test_eplb_utils.py
   commands:
   - pytest -v -s distributed/test_eplb_algo.py
+  - pytest -v -s distributed/test_eplb_utils.py
 
 - label: EPLB Execution
   timeout_in_minutes: 20

.github/mergify.yml

@@ -234,6 +234,36 @@ pull_request_rules:
       add:
         - rocm
 
+- name: label-xpu
+  description: Automatically apply intel-gpu label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^docker/Dockerfile.xpu
+      - files~=^\\.buildkite/intel_jobs/
+      - files=\.buildkite/ci_config_intel.yaml
+      - files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
+      - files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
+      - files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+      - files=vllm/distributed/device_communicators/xpu_communicator.py
+      - files=vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+      - files=vllm/v1/attention/ops/xpu_mla_sparse.py
+      - files=vllm/v1/worker/xpu_worker.py
+      - files=vllm/v1/worker/xpu_model_runner.py
+      - files=vllm/_xpu_ops.py
+      - files~=^vllm/lora/ops/xpu_ops
+      - files=vllm/lora/punica_wrapper/punica_xpu.py
+      - files=vllm/platforms/xpu.py
+      - title~=(?i)Intel gpu
+      - title~=(?i)XPU
+      - title~=(?i)Intel
+      - title~=(?i)BMG
+      - title~=(?i)Arc
+  actions:
+    label:
+      add:
+        - intel-gpu
+
 - name: label-cpu
   description: Automatically apply cpu label
   conditions:

CMakeLists.txt

@@ -363,7 +363,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # - sm80 doesn't support fp8 computation
   # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
   # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
   # marlin arches for other files
   cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
 
@@ -523,12 +523,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   endif()
 
 
-  # The cutlass_scaled_mm kernels for Geforce Blackwell SM120 (c3x, i.e. CUTLASS 3.x) require
+  # The cutlass_scaled_mm kernels for Blackwell SM12x (c3x, i.e. CUTLASS 3.x) require
   # CUDA 12.8 or later
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0f" "${CUDA_ARCHS}")
   else()
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
   endif()
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
     set(SRCS
@@ -616,12 +616,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
-  # The nvfp4_scaled_mm_sm120 kernels for Geforce Blackwell SM120 require
+  # The nvfp4_scaled_mm_sm120 kernels for Blackwell SM12x require
   # CUDA 12.8 or later
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(FP4_ARCHS "12.0f" "${CUDA_ARCHS}")
   else()
-    cuda_archs_loose_intersection(FP4_ARCHS "12.0a" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(FP4_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
   endif()
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_ARCHS)
     set(SRCS
@@ -1050,7 +1050,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # - sm80 doesn't support fp8 computation
   # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
   # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(MARLIN_MOE_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
   # moe marlin arches for other files
   cuda_archs_loose_intersection(MARLIN_MOE_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
   if (MARLIN_MOE_OTHER_ARCHS)

benchmarks/attention_benchmarks/benchmark.py

@@ -546,10 +546,7 @@ def main():
         args.prefill_backends = yaml_config.get("prefill_backends", None)
 
         # Check for special modes
-        if "mode" in yaml_config:
-            args.mode = yaml_config["mode"]
-        else:
-            args.mode = None
+        args.mode = yaml_config.get("mode", None)
 
         # Batch specs and sizes
         # Support both explicit batch_specs and generated batch_spec_ranges
@@ -572,10 +569,7 @@ def main():
             elif "batch_specs" in yaml_config:
                 args.batch_specs = yaml_config["batch_specs"]
 
-        if "batch_sizes" in yaml_config:
-            args.batch_sizes = yaml_config["batch_sizes"]
-        else:
-            args.batch_sizes = None
+        args.batch_sizes = yaml_config.get("batch_sizes", None)
 
         # Model config
         if "model" in yaml_config:

benchmarks/kernels/benchmark_moe.py

@@ -627,9 +627,8 @@ class BenchmarkWorker:
                 need_device_guard = True
 
         with (
-            torch.accelerator.device_index(self.device_id)
-            if need_device_guard
-            else nullcontext()
+            # Ray restricts each worker to one GPU; use local index 0
+            torch.accelerator.device_index(0) if need_device_guard else nullcontext()
         ):
             for idx, config in enumerate(tqdm(search_space)):
                 try:

cmake/external_projects/qutlass.cmake

@@ -32,16 +32,16 @@ endif()
 message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-  cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0f" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_ARCHS "10.0f;12.0f" "${CUDA_ARCHS}")
 else()
-  cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0a;10.3a" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;12.1a;10.0a;10.3a" "${CUDA_ARCHS}")
 endif()
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
 
   if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)")
     set(QUTLASS_TARGET_CC 100)
-  elseif(QUTLASS_ARCHS MATCHES "12\\.0a")
+  elseif(QUTLASS_ARCHS MATCHES "12\\.[01][af]?")
     set(QUTLASS_TARGET_CC 120)
   else()
     message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.")
@@ -96,7 +96,7 @@ else()
       "[QUTLASS] Skipping build: CUDA 12.8 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION}).")
   else()
     message(STATUS
-      "[QUTLASS] Skipping build: no supported arch (12.0a / 10.0a) found in "
+      "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
       "CUDA_ARCHS='${CUDA_ARCHS}'.")
   endif()
 endif()

cmake/utils.cmake

@@ -355,8 +355,11 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
   list(REMOVE_DUPLICATES _PTX_ARCHS)
   list(REMOVE_DUPLICATES _SRC_CUDA_ARCHS)
 
-  # If x.0a or x.0f is in SRC_CUDA_ARCHS and x.0 is in CUDA_ARCHS then we should
-  # remove x.0a or x.0f from SRC_CUDA_ARCHS and add x.0a or x.0f to _CUDA_ARCHS
+  # Handle architecture-specific suffixes (a/f) for SRC entries.
+  # First try exact base match (x.y), then cross-suffix match (x.ya / x.yf).
+  # For 'f' (family) suffix: if no exact/cross match, fall back to major-version
+  # match — e.g. SRC="12.0f" matches TGT="12.1a" since SM121 is in the SM12x
+  # family. The output uses TGT's value to preserve the user's compilation flags.
   set(_CUDA_ARCHS)
   foreach(_arch ${_SRC_CUDA_ARCHS})
     if(_arch MATCHES "[af]$")
@@ -365,6 +368,38 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
       if ("${_base}" IN_LIST TGT_CUDA_ARCHS)
         list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}")
         list(APPEND _CUDA_ARCHS "${_arch}")
+      elseif("${_base}a" IN_LIST _TGT_CUDA_ARCHS)
+        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}a")
+        list(APPEND _CUDA_ARCHS "${_base}a")
+      elseif("${_base}f" IN_LIST _TGT_CUDA_ARCHS)
+        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_base}f")
+        list(APPEND _CUDA_ARCHS "${_base}f")
+      elseif(_arch MATCHES "f$")
+        # Family suffix: match any TGT entry in the same major version family.
+        string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" _src_major "${_base}")
+        foreach(_tgt ${_TGT_CUDA_ARCHS})
+          string(REGEX REPLACE "[af]$" "" _tgt_base "${_tgt}")
+          string(REGEX REPLACE "^([0-9]+)\\..*$" "\\1" _tgt_major "${_tgt_base}")
+          if(_tgt_major STREQUAL _src_major)
+            list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_tgt}")
+            list(APPEND _CUDA_ARCHS "${_tgt}")
+            break()
+          endif()
+        endforeach()
+      endif()
+    endif()
+  endforeach()
+
+  # Symmetric handling: if TGT has x.ya/f and SRC has x.y (without suffix),
+  # preserve TGT's suffix in the output.
+  set(_tgt_copy ${_TGT_CUDA_ARCHS})
+  foreach(_arch ${_tgt_copy})
+    if(_arch MATCHES "[af]$")
+      string(REGEX REPLACE "[af]$" "" _base "${_arch}")
+      if ("${_base}" IN_LIST _SRC_CUDA_ARCHS)
+        list(REMOVE_ITEM _TGT_CUDA_ARCHS "${_arch}")
+        list(REMOVE_ITEM _SRC_CUDA_ARCHS "${_base}")
+        list(APPEND _CUDA_ARCHS "${_arch}")
       endif()
     endif()
   endforeach()

csrc/moe/marlin_moe_wna16/kernel.h

@@ -13,7 +13,7 @@
       const int4 *__restrict__ b_bias_ptr,                            \
       const float *__restrict__ a_scales_ptr,                         \
       const int4 *__restrict__ scales_ptr,                            \
-      const uint16_t *__restrict__ global_scale_ptr,                  \
+      const float *__restrict__ global_scale_ptr,                     \
       const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
       const int32_t *__restrict__ sorted_token_ids_ptr,               \
       const int32_t *__restrict__ expert_ids_ptr,                     \

csrc/moe/marlin_moe_wna16/marlin_template.h

@@ -260,7 +260,7 @@ __global__ void Marlin(
     // fp16 quantization scales. shape (k/groupsize, n)
     const int4* __restrict__ scales_ptr,
     // fp16 global scale (for nvfp4// only)
-    const uint16_t* __restrict__ global_scale_ptr,
+    const float* __restrict__ global_scale_ptr,
     // 4bit packed zero-points of shape
     // (k/groupsize, n/pack_factor)
     const int4* __restrict__ zp_ptr,
@@ -308,7 +308,14 @@ __global__ void Marlin(
   constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
 
   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
-  constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
+  static constexpr auto num_bits =
+      vllm::ScalarType::from_id(b_type_id).size_bits();
+  // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
+  // num_bits == 4
+  constexpr bool use_fp16_accum =
+      a_type_id == vllm::kFloat16.id() &&
+      (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
+       !(group_blocks == -1 && num_bits == 4));
   #else
   constexpr bool use_fp16_accum = false;
   #endif
@@ -357,7 +364,7 @@ __global__ void Marlin(
       has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
       has_zp && !is_zp_float && !(b_type == vllm::kU8);
 
-  c_scalar_t2 global_scale;
+  float global_scale_f32 = 1.0f;
 
   constexpr bool has_act_order = group_blocks == 0;
 
@@ -507,11 +514,12 @@ __global__ void Marlin(
 
       if (mul_topk_weights) {
         idx = idx < prob_m_top_k ? idx : 0;
-        c_scalar_t2 topk_weight_val =
-            Cdtype::num2num2(Cdtype::float2num(topk_weights_ptr[idx]));
+        float topk_weight_tmp = topk_weights_ptr[idx];
         if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
-          topk_weight_val = __hmul2(topk_weight_val, global_scale);
+          topk_weight_tmp *= global_scale_f32;
         }
+        c_scalar_t2 topk_weight_val =
+            Cdtype::num2num2(Cdtype::float2num(topk_weight_tmp));
         sh_block_topk_weights[threadIdx.x] = topk_weight_val;
       }
     }
@@ -532,8 +540,7 @@ __global__ void Marlin(
     expert_id = expert_ids_ptr[block_id];
 
     if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
-      uint16_t val = global_scale_ptr[expert_id];
-      global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
+      global_scale_f32 = global_scale_ptr[expert_id];
     }
 
     B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
@@ -1784,6 +1791,13 @@ __global__ void Marlin(
     // We first reorder in shared memory to guarantee the most efficient final
     // global write patterns
     auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
+      if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+        if (!mul_topk_weights) {
+          c0 *= global_scale_f32;
+          c1 *= global_scale_f32;
+        }
+      }
+
       c_scalar_t2 res =
           Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
 
@@ -1800,11 +1814,6 @@ __global__ void Marlin(
         res = __hmul2(res, tmp_scale);
       }
 
-      if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
-        if (!mul_topk_weights) {
-          res = __hmul2(res, global_scale);
-        }
-      }
       if (has_bias && last) {
         c_scalar_t2 tmp_bias = b_bias[0];
         if constexpr (m_block_size_8) {

csrc/moe/marlin_moe_wna16/ops.cu

@@ -382,7 +382,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
   const int4* bias_ptr = (const int4*)b_bias;
   const float* a_s_ptr = (const float*)a_s;
   const int4* b_s_ptr = (const int4*)b_s;
-  const uint16_t* g_s_ptr = (const uint16_t*)g_s;
+  const float* g_s_ptr = (const float*)g_s;
   const int4* zp_ptr = (const int4*)zp;
   const int* g_idx_ptr = (const int*)g_idx;
   const int* perm_ptr = (const int*)perm;
@@ -759,7 +759,7 @@ torch::Tensor moe_wna16_marlin_gemm(
     TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
                 "global_scale can only be used for nvfp4 format.");
   } else {
-    global_scale = torch::empty({0}, options);
+    global_scale = torch::empty({0}, options_fp32);
     TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
                 "the global_scale parameter must be passed for nvfp4 format.");
   }
@@ -842,8 +842,8 @@ torch::Tensor moe_wna16_marlin_gemm(
 
   TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
               "scalar type of a_scales must be float");
-  TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
-              "scalar type of global_scale must be the same with c");
+  TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
+              "scalar type of global_scale must be float");
   if (a_type.size_bits() == 16) {
     TORCH_CHECK(
         a.scalar_type() == c.scalar_type(),

csrc/quantization/activation_kernels.cu

@@ -189,10 +189,7 @@ __device__ __forceinline__ void cp_async_wait<0>() {
 }
 
 __device__ __forceinline__ float clip(float v, float mmin, float mmax) {
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
   return fminf(mmax, fmaxf(v, mmin));
-#else
-#endif
 }
 
 __device__ __forceinline__ __nv_bfloat16 clip(__nv_bfloat16 v,

csrc/quantization/marlin/kernel.h

@@ -13,7 +13,7 @@
       const int4 *__restrict__ b_bias_ptr,                                     \
       const float *__restrict__ a_scales_ptr,                                  \
       const int4 *__restrict__ scales_ptr,                                     \
-      const uint16_t *__restrict__ global_scale_ptr,                           \
+      const float *__restrict__ global_scale_ptr,                              \
       const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx,          \
       int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
       bool has_bias, bool use_atomic_add, bool use_fp32_reduce,                \

csrc/quantization/marlin/marlin.cu

@@ -57,7 +57,7 @@ torch::Tensor marlin_gemm(
     int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
     bool is_zp_float) {
   TORCH_CHECK_NOT_IMPLEMENTED(false,
-                              "marlin_gemm(..) requires CUDA_ARCH >= 8.0");
+                              "marlin_gemm(..) requires CUDA_ARCH >= 7.5");
   return torch::empty({1, 1});
 }
 
@@ -356,7 +356,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
   const int4* bias_ptr = (const int4*)b_bias;
   const float* a_s_ptr = (const float*)a_s;
   const int4* b_s_ptr = (const int4*)b_s;
-  const uint16_t* g_s_ptr = (const uint16_t*)g_s;
+  const float* g_s_ptr = (const float*)g_s;
 
   const int4* zp_ptr = (const int4*)zp;
   const int* g_idx_ptr = (const int*)g_idx;
@@ -751,7 +751,7 @@ torch::Tensor marlin_gemm(
     TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
                 "global_scale can only be used for nvfp4 format.");
   } else {
-    global_scale = torch::empty({0}, options);
+    global_scale = torch::empty({0}, options_fp32);
     TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
                 "the global_scale parameter must be passed for nvfp4 format.");
   }
@@ -832,8 +832,8 @@ torch::Tensor marlin_gemm(
 
   TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
               "scalar type of a_scales must be float");
-  TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
-              "scalar type of global_scale must be the same with c");
+  TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
+              "scalar type of global_scale must be float");
   if (a_type.size_bits() == 16) {
     TORCH_CHECK(
         a.scalar_type() == c.scalar_type(),

csrc/quantization/marlin/marlin_template.h

@@ -251,8 +251,8 @@ __global__ void Marlin(
     const float* __restrict__ a_scales_ptr,
     // fp16 quantization scales. shape (k/groupsize, n)
     const int4* __restrict__ scales_ptr,
-    // fp16 global scale (for nvfp4// only)
-    const uint16_t* __restrict__ global_scale_ptr,
+    // float global scale (for nvfp4// only)
+    const float* __restrict__ global_scale_ptr,
     // 4bit packed zero-points of shape
     // (k/groupsize, n/pack_factor)
     const int4* __restrict__ zp_ptr,
@@ -292,7 +292,13 @@ __global__ void Marlin(
   #endif
 
   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
-  constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
+  constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits();
+  // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
+  // num_bits == 4
+  constexpr bool use_fp16_accum =
+      a_type_id == vllm::kFloat16.id() &&
+      (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
+       !(group_blocks == -1 && num_bits == 4));
   #else
   constexpr bool use_fp16_accum = false;
   #endif
@@ -342,11 +348,10 @@ __global__ void Marlin(
       has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
       has_zp && !is_zp_float && !(b_type == vllm::kU8);
 
-  c_scalar_t2 global_scale;
+  float global_scale_f32 = 1.0f;
 
   if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
-    uint16_t val = global_scale_ptr[0];
-    global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
+    global_scale_f32 = global_scale_ptr[0];
   }
 
   constexpr bool has_act_order = group_blocks == 0;
@@ -1644,6 +1649,10 @@ __global__ void Marlin(
     // We first reorder in shared memory to guarantee the most efficient final
     // global write patterns
     auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
+      if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
+        c0 *= global_scale_f32;
+        c1 *= global_scale_f32;
+      }
       c_scalar_t2 res =
           Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
 
@@ -1659,10 +1668,6 @@ __global__ void Marlin(
         }
         res = __hmul2(res, tmp_scale);
       }
-
-      if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
-        res = __hmul2(res, global_scale);
-      }
       if (has_bias && last) {
         c_scalar_t2 tmp_bias = b_bias[0];
         if constexpr (m_block_size_8) {

docker/Dockerfile

@@ -596,6 +596,25 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
     && flashinfer show-config
 
+# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
+# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
+# which fails without internet. This step caches them at build time.
+RUN python3 <<'PYEOF'
+from flashinfer.jit import env as jit_env
+from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
+from flashinfer.artifacts import ArtifactPath, CheckSumHash
+
+download_trtllm_headers(
+    'bmm',
+    jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
+    f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
+    ArtifactPath.TRTLLM_GEN_BMM,
+    get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
+)
+
+print('FlashInfer TRTLLM BMM headers downloaded successfully')
+PYEOF
+
 # ============================================================
 # OPENAI API SERVER DEPENDENCIES
 # Pre-install these to avoid reinstalling on every vLLM wheel rebuild

docker/Dockerfile.rocm

@@ -386,6 +386,9 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
 # will not be imported by other tests
 RUN mkdir src && mv vllm src/vllm
 
+# This is a workaround to ensure pytest exits with the correct status code in CI tests.
+RUN echo "import os\n\ndef pytest_sessionfinish(session, exitstatus):\n    os._exit(int(exitstatus))" > /vllm-workspace/conftest.py
+
 # -----------------------
 # Final vLLM image
 FROM base AS final

docker/Dockerfile.rocm_base

@@ -1,7 +1,7 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
-ARG TRITON_BRANCH="57c693b6"
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.1-complete
+ARG TRITON_BRANCH="ba5c1517"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
-ARG PYTORCH_BRANCH="89075173"
+ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17
 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
 ARG PYTORCH_VISION_BRANCH="v0.24.1"
 ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
@@ -114,6 +114,8 @@ ARG TRITON_REPO
 RUN git clone ${TRITON_REPO}
 RUN cd triton \
     && git checkout ${TRITON_BRANCH} \
+    && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \
+    && git cherry-pick 555d04f \
     && if [ ! -f setup.py ]; then cd python; fi \
     && python3 setup.py bdist_wheel --dist-dir=dist \
     && mkdir -p /app/install && cp dist/*.whl /app/install
@@ -142,10 +144,14 @@ ARG PYTORCH_VISION_REPO
 ARG PYTORCH_AUDIO_REPO
 ARG USE_SCCACHE
 
+RUN apt-get update && apt-get install -y pkg-config liblzma-dev
 RUN git clone ${PYTORCH_REPO} pytorch
-RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
-    && pip install -r requirements.txt && git submodule update --init --recursive \
-    && python3 tools/amd_build/build_amd.py \
+RUN cd pytorch && git checkout ${PYTORCH_BRANCH}
+RUN cd pytorch \
+    && pip install -r requirements.txt && git submodule update --init --recursive
+RUN cd pytorch/third_party/kineto \
+    && git remote add rocm https://github.com/ROCm/kineto && git fetch rocm && git checkout 2d73be3 
+RUN cd pytorch && python3 tools/amd_build/build_amd.py \
     && if [ "$USE_SCCACHE" = "1" ]; then \
            export HIP_CLANG_PATH=/opt/sccache-wrappers \
            && export CMAKE_C_COMPILER_LAUNCHER=sccache \
@@ -239,7 +245,7 @@ RUN pip install pyyaml && cd aiter \
            export HIP_CLANG_PATH=/opt/sccache-wrappers \
            && sccache --show-stats; \
        fi \
-    && GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
     && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
     && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install

docs/deployment/frameworks/helm.md

@@ -17,6 +17,8 @@ Before you begin, ensure that you have the following:
 
 ## Installing the chart
 
+This guide uses the Helm chart at [examples/online_serving/chart-helm](../../../examples/online_serving/chart-helm).
+
 To install the chart with the release name `test-vllm`:
 
 ```bash

docs/getting_started/installation/gpu.rocm.inc.md

@@ -172,8 +172,11 @@ uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700
 --8<-- [end:build-wheel-from-source]
 --8<-- [start:pre-built-images]
 
-vLLM offers an official Docker image for deployment.
-The image can be used to run OpenAI compatible server and is available on Docker Hub as [vllm/vllm-openai-rocm](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags).
+vLLM offers official Docker images for deployment.
+The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-rocm](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags).
+
+- `vllm/vllm-openai-rocm:latest` — stable release
+- `vllm/vllm-openai-rocm:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
 
 ```bash
 docker run --rm \
@@ -186,30 +189,18 @@ docker run --rm \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8000:8000 \
     --ipc=host \
-    vllm/vllm-openai-rocm:latest \
+    vllm/vllm-openai-rocm:<tag> \
     --model Qwen/Qwen3-0.6B
 ```
 
-#### Use AMD's Docker Images
+#### Use AMD's Docker Images (Deprecated)
 
-Prior to January 20th, 2026 when the official docker images are available on [upstream vLLM docker hub](https://hub.docker.com/v2/repositories/vllm/vllm-openai-rocm/tags/), the [AMD Infinity hub for vLLM](https://hub.docker.com/r/rocm/vllm/tags) offers a prebuilt, optimized
+!!! warning "Deprecated"
+    AMD's Docker images (`rocm/vllm` and `rocm/vllm-dev`) are deprecated in favor of the official vLLM Docker images above (`vllm/vllm-openai-rocm`). Please migrate to the official images.
+
+Prior to January 20th, 2026 when the official docker images became available on [upstream vLLM docker hub](https://hub.docker.com/v2/repositories/vllm/vllm-openai-rocm/tags/), the [AMD Infinity hub for vLLM](https://hub.docker.com/r/rocm/vllm/tags) offered a prebuilt, optimized
 docker image designed for validating inference performance on the AMD Instinct MI300X™ accelerator.
-AMD also offers nightly prebuilt docker image from [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev), which has vLLM and all its dependencies installed. The entrypoint of this docker image is `/bin/bash` (different from the vLLM's Official Docker Image).
-
-```bash
-docker pull rocm/vllm-dev:nightly # to get the latest image
-docker run -it --rm \
---network=host \
---group-add=video \
---ipc=host \
---cap-add=SYS_PTRACE \
---security-opt seccomp=unconfined \
---device /dev/kfd \
---device /dev/dri \
--v <path/to/your/models>:/app/models \
--e HF_HOME="/app/models" \
-rocm/vllm-dev:nightly
-```
+AMD also offered nightly prebuilt docker image from [Docker Hub](https://hub.docker.com/r/rocm/vllm-dev), which has vLLM and all its dependencies installed. The entrypoint of this docker image is `/bin/bash` (different from the vLLM's Official Docker Image).
 
 !!! tip
     Please check [LLM inference performance validation on AMD Instinct MI300X](https://rocm.docs.amd.com/en/latest/how-to/performance-validation/mi300x/vllm-benchmark.html)

docs/getting_started/quickstart.md

@@ -56,9 +56,12 @@ This guide will help you quickly get started with vLLM to perform:
     !!! note
         It currently supports Python 3.12, ROCm 7.0 and `glibc >= 2.35`.
 
-    !!! note    
+    !!! note
         Note that, previously, docker images were published using AMD's docker release pipeline and were located `rocm/vllm-dev`. This is being deprecated by using vLLM's docker release pipeline.
 
+    !!! tip
+        A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
+
 === "Google TPU"
 
     To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.

docs/mkdocs/hooks/generate_argparse.py

@@ -153,7 +153,7 @@ class MarkdownFormatter(HelpFormatter):
             heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
             self._markdown_output.append(heading_md)
 
-            if action.choices or isinstance(action.metavar, (list, tuple)):
+            if action.choices or isinstance(action.metavar, list | tuple):
                 choices_iterable = action.choices or action.metavar
                 choices = f"`{'`, `'.join(str(c) for c in choices_iterable)}`"
                 self._markdown_output.append(f":   Possible choices: {choices}\n\n")

docs/usage/security.md

@@ -231,6 +231,18 @@ The most effective approach is to deploy vLLM behind a reverse proxy (such as ng
 - Blocks all other endpoints, including the unauthenticated inference and operational control endpoints
 - Implements additional authentication, rate limiting, and logging at the proxy layer
 
+## Request Parameter Resource Limits
+
+Certain API request parameters can have a large impact on resource consumption and may be abused to exhaust server resources. The `n` parameter in the `/v1/completions` and `/v1/chat/completions` endpoints controls how many independent output sequences are generated per request. A very large value causes the engine to allocate memory, CPU, and GPU time proportional to `n`, which can lead to out-of-memory conditions on the host and block the server from processing other requests.
+
+To mitigate this, vLLM enforces a configurable upper bound on the `n` parameter via the `VLLM_MAX_N_SEQUENCES` environment variable (default: **16384**). Requests exceeding this limit are rejected before reaching the engine.
+
+### Recommendations
+
+- **Public-facing deployments:** Consider setting `VLLM_MAX_N_SEQUENCES` to a value appropriate for your workload (e.g., `64` or `128`) to limit the blast radius of a single request.
+- **Reverse proxy layer:** In addition to vLLM's built-in limit, consider enforcing request body validation and rate limiting at your reverse proxy to further constrain abusive payloads.
+- **Monitoring:** Monitor per-request resource consumption to detect anomalous patterns that may indicate abuse.
+
 ## Tool Server and MCP Security
 
 vLLM supports connecting to external tool servers via the `--tool-server` argument. This enables models to call tools through the Responses API (`/v1/responses`). Tool server support works with all models — it is not limited to specific model architectures.

tests/distributed/test_eplb_utils.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.distributed.eplb.eplb_state import (
+    _commit_eplb_maps,
+    _commit_eplb_maps_for_layer,
+)
+
+
+def _make_model_state(
+    phy2log: torch.Tensor,
+    log2phy: torch.Tensor,
+    logcnt: torch.Tensor,
+) -> MagicMock:
+    """Build a minimal EplbModelState mock with only the three map tensors."""
+    state = MagicMock()
+    state.physical_to_logical_map = phy2log
+    state.logical_to_physical_map = log2phy
+    state.logical_replica_count = logcnt
+    return state
+
+
+def test_commit_eplb_maps_shape_change():
+    """
+    The normal path copies the physical_to_logical map in-place. When the number of
+    physical experts changes, the old map should be replaced entirely.
+    """
+    num_layers, num_logical, num_physical = 2, 4, 6
+    max_replicas = 3
+
+    # Build current state tensors
+    model_state = _make_model_state(
+        phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
+        log2phy=torch.full(
+            (num_layers, num_logical, max_replicas), -1, dtype=torch.long
+        ),
+        logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
+    )
+
+    # The new map has two more physical experts. These new physical experts will
+    # automatically map to the first two logical experts
+    new_phy2log_larger = (
+        (torch.arange(num_physical + 2, dtype=torch.long) % num_logical)
+        .unsqueeze(0)
+        .expand(num_layers, -1)
+    )
+    _commit_eplb_maps(model_state, new_phy2log_larger)
+
+    # Check that the number of physical experts has been updated and that the values
+    # match
+    assert model_state.physical_to_logical_map.shape[1] == num_physical + 2
+    assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger)
+
+
+def test_commit_eplb_maps_for_layer_logical_padding():
+    """
+    Test that logical_to_physical_map is padded with -1 to fill the
+    pre-allocated slots when the new map has fewer replicas than the max.
+    """
+    num_layers, num_logical, num_physical = 2, 4, 6
+    max_replicas = 3
+
+    model_state = _make_model_state(
+        phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
+        log2phy=torch.full(
+            (num_layers, num_logical, max_replicas), -1, dtype=torch.long
+        ),
+        logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
+    )
+
+    new_phy2log = (
+        (torch.arange(num_physical, dtype=torch.long) % num_logical)
+        .unsqueeze(0)
+        .expand(num_layers, -1)
+        .contiguous()
+    )
+    layer = 0
+    _commit_eplb_maps_for_layer(model_state, new_phy2log, layer)
+
+    assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1)
+
+
+def test_commit_eplb_maps_for_layer_shape_assert():
+    """Test that a mismatched number of physical experts triggers an assertion error."""
+    num_layers, num_logical, num_physical = 2, 4, 6
+
+    model_state = _make_model_state(
+        phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
+        log2phy=torch.full((num_layers, num_logical, 2), -1, dtype=torch.long),
+        logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
+    )
+    bad_phy2log = torch.zeros(num_layers, num_physical + 1, dtype=torch.long)
+    with pytest.raises(AssertionError):
+        _commit_eplb_maps_for_layer(model_state, bad_phy2log, layer=0)
+
+
+def test_commit_eplb_maps():
+    """Test that all values are copied correctly into model_state."""
+    num_layers, num_logical, num_physical, max_replicas = 2, 3, 4, 2
+
+    model_state = _make_model_state(
+        phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
+        log2phy=torch.full(
+            (num_layers, num_logical, max_replicas), -1, dtype=torch.long
+        ),
+        logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
+    )
+
+    new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
+    new_log2phy = torch.tensor(
+        [[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
+    )
+    new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
+
+    _commit_eplb_maps(model_state, new_phy2log)
+
+    assert torch.equal(model_state.physical_to_logical_map, new_phy2log)
+    assert torch.equal(model_state.logical_to_physical_map, new_log2phy)
+    assert torch.equal(model_state.logical_replica_count, new_logcnt)
+
+
+def test_commit_eplb_maps_for_layer():
+    """Test that only the target layer is updated"""
+    num_layers, num_logical, max_replicas = 2, 3, 2
+
+    original_phy2log = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    model_state = _make_model_state(
+        phy2log=original_phy2log.clone(),
+        log2phy=torch.full(
+            (num_layers, num_logical, max_replicas), -1, dtype=torch.long
+        ),
+        logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
+    )
+
+    new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
+    new_log2phy = torch.tensor(
+        [[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
+    )
+    new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
+
+    _commit_eplb_maps_for_layer(model_state, new_phy2log, layer=0)
+
+    # Layer 0 updated
+    assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0])
+    assert torch.equal(model_state.logical_to_physical_map[0], new_log2phy[0])
+    assert torch.equal(model_state.logical_replica_count[0], new_logcnt[0])
+
+    # Layer 1 untouched
+    assert torch.equal(model_state.physical_to_logical_map[1], original_phy2log[1])

tests/entrypoints/openai/chat_completion/test_chat.py

@@ -1020,3 +1020,114 @@ def test_chat_completion_request_n_parameter_various_values():
         assert sampling_params.n == n_value, (
             f"Expected n={n_value}, got n={sampling_params.n}"
         )
+
+
+def test_chat_completion_request_n_parameter_exceeds_default_limit(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test that n values exceeding the default limit are rejected."""
+    import vllm.envs as envs
+
+    monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
+    if hasattr(envs.__getattr__, "cache_clear"):
+        envs.__getattr__.cache_clear()
+
+    max_n = envs.VLLM_MAX_N_SEQUENCES
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Test"}],
+        n=max_n + 1,
+        max_tokens=10,
+    )
+
+    with pytest.raises(ValueError, match="n must be at most"):
+        request.to_sampling_params(
+            max_tokens=10,
+            default_sampling_params={},
+        )
+
+
+def test_chat_completion_request_n_parameter_at_limit(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test that n at exactly the limit is accepted."""
+    import vllm.envs as envs
+
+    monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
+    if hasattr(envs.__getattr__, "cache_clear"):
+        envs.__getattr__.cache_clear()
+
+    max_n = envs.VLLM_MAX_N_SEQUENCES
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Test"}],
+        n=max_n,
+        max_tokens=10,
+    )
+
+    sampling_params = request.to_sampling_params(
+        max_tokens=10,
+        default_sampling_params={},
+    )
+    assert sampling_params.n == max_n
+
+
+def test_chat_completion_request_n_parameter_custom_limit(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test that VLLM_MAX_N_SEQUENCES env var overrides the default limit."""
+    import vllm.envs as envs
+
+    monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
+    if hasattr(envs.__getattr__, "cache_clear"):
+        envs.__getattr__.cache_clear()
+
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Test"}],
+        n=128,
+        max_tokens=10,
+    )
+
+    sampling_params = request.to_sampling_params(
+        max_tokens=10,
+        default_sampling_params={},
+    )
+    assert sampling_params.n == 128
+
+    request_over = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Test"}],
+        n=129,
+        max_tokens=10,
+    )
+
+    with pytest.raises(ValueError, match="n must be at most 128"):
+        request_over.to_sampling_params(
+            max_tokens=10,
+            default_sampling_params={},
+        )
+
+
+def test_chat_completion_request_n_parameter_massive_value(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Test that astronomically large n values are rejected (CVE fix)."""
+    import vllm.envs as envs
+
+    monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
+    if hasattr(envs.__getattr__, "cache_clear"):
+        envs.__getattr__.cache_clear()
+
+    request = ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "Test"}],
+        n=100_000_000,
+        max_tokens=1,
+    )
+
+    with pytest.raises(ValueError, match="n must be at most"):
+        request.to_sampling_params(
+            max_tokens=1,
+            default_sampling_params={},
+        )

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -55,6 +55,7 @@ class MockModelConfig:
     skip_tokenizer_init = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -536,6 +536,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}

tests/entrypoints/openai/completion/test_completion_error.py

@@ -54,6 +54,7 @@ class MockModelConfig:
     skip_tokenizer_init = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}

tests/entrypoints/openai/completion/test_lora_resolvers.py

@@ -54,6 +54,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
 
     def get_diff_sampling_param(self):
         return self.diff_sampling_param or {}

tests/model_executor/model_loader/test_reload.py

@@ -148,3 +148,60 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
         mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
         add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
         assert add_perp < mul_perp
+
+
+@pytest.mark.parametrize("tp_size", [2])
+@pytest.mark.parametrize(
+    "base_model,mul_model,add_model,quantization",
+    [
+        (
+            "Qwen/Qwen3-0.6B",
+            "inference-optimization/Qwen3-0.6B-debug-multiply",
+            "inference-optimization/Qwen3-0.6B-debug-add",
+            "fp8",
+        ),
+        (
+            "inference-optimization/DeepSeek-V3-debug-empty",
+            "inference-optimization/DeepSeek-V3-debug-multiply",
+            "inference-optimization/DeepSeek-V3-debug-add",
+            "fp8",
+        ),
+        (
+            "Qwen/Qwen3-0.6B",
+            "inference-optimization/Qwen3-0.6B-debug-multiply",
+            "inference-optimization/Qwen3-0.6B-debug-add",
+            "mxfp8",
+        ),
+        # (  TODO: support mxfp4 & mla
+        #     "inference-optimization/DeepSeek-V3-debug-empty",
+        #     "inference-optimization/DeepSeek-V3-debug-multiply",
+        #     "inference-optimization/DeepSeek-V3-debug-add",
+        #     "mxfp8",
+        # ),
+    ],
+)
+def test_online_quantize_reload(
+    base_model, mul_model, add_model, quantization, tp_size, vllm_runner
+):
+    if cuda_device_count_stateless() < tp_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    if quantization == "fp8" and not current_platform.supports_fp8():
+        pytest.skip(reason="Requires FP8 support")
+
+    with vllm_runner(
+        model_name=base_model,
+        quantization=quantization,
+        tensor_parallel_size=tp_size,
+        enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
+        enable_prefix_caching=False,
+    ) as llm:
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
+        mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
+        add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
+        assert mul_perp < add_perp
+
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})
+        mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
+        add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
+        assert add_perp < mul_perp

tests/models/registry.py

@@ -791,6 +791,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
     "Ernie4_5_VLMoeForConditionalGeneration": _HfExamplesInfo(
         "baidu/ERNIE-4.5-VL-28B-A3B-PT",
         trust_remote_code=True,
+        revision="refs/pr/14",
     ),
     "FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
         "allendou/FireRedASR2-LLM-vllm",

tests/quantization/test_cpu_wna16.py

@@ -11,6 +11,7 @@ MODELS = [
     "TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
     "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ",  # with g_idx
     "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4",  # without g_idx
+    "RedHatAI/Qwen3-1.7B-quantized.w4a16",  # with zp
 ]
 DTYPE = ["bfloat16"]
 

tests/quantization/test_fp8.py

@@ -466,3 +466,26 @@ def test_fp8_reloading(
         weight_loader(param, torch.zeros(shape))  # cannot use empty
 
     method.process_weights_after_loading(layer)
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+def test_kv_cache_dtype_skip_layers(vllm_runner, monkeypatch):
+    """Test that kv_cache_dtype_skip_layers skips quantization for specified layers."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+    with vllm_runner(
+        "facebook/opt-125m",
+        kv_cache_dtype="fp8",
+        kv_cache_dtype_skip_layers=["0", "2"],
+        enforce_eager=True,
+    ) as llm:
+
+        def check_layers(model):
+            for i, layer in enumerate(model.model.decoder.layers):
+                expected = "auto" if str(i) in ["0", "2"] else "fp8"
+                assert layer.self_attn.attn.kv_cache_dtype == expected
+
+        llm.apply_model(check_layers)

tests/renderers/test_completions.py

@@ -38,6 +38,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
 
 
 @dataclass

tests/renderers/test_mistral.py

@@ -37,6 +37,7 @@ class MockModelConfig:
     skip_tokenizer_init: bool = False
     is_encoder_decoder: bool = False
     is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
 
 
 @dataclass

tests/test_config.py

@@ -1131,6 +1131,28 @@ def test_needs_dp_coordination(
     assert vllm_config.needs_dp_coordinator == expected_needs_coordinator
 
 
+def test_renderer_num_workers_with_mm_cache():
+    """Disallow renderer_num_workers > 1 when mm processor cache is enabled,
+    since neither cache type is thread-safe."""
+    mm_model = "Qwen/Qwen2-VL-2B-Instruct"
+
+    # Should raise: multi-worker + cache enabled (default cache_gb=4)
+    with pytest.raises(ValueError, match="renderer-num-workers"):
+        ModelConfig(mm_model, renderer_num_workers=4)
+
+    # Should raise: multi-worker + explicit cache size
+    with pytest.raises(ValueError, match="renderer-num-workers"):
+        ModelConfig(mm_model, renderer_num_workers=2, mm_processor_cache_gb=1.0)
+
+    # Should pass: multi-worker + cache disabled
+    config = ModelConfig(mm_model, renderer_num_workers=4, mm_processor_cache_gb=0)
+    assert config.renderer_num_workers == 4
+
+    # Should pass: single worker + cache enabled (default)
+    config = ModelConfig(mm_model, renderer_num_workers=1)
+    assert config.renderer_num_workers == 1
+
+
 def test_eagle_draft_model_config():
     """Test that EagleDraft model config is correctly set."""
     target_model_config = ModelConfig(

tests/test_envs.py

@@ -454,3 +454,55 @@ class TestVllmConfigureLogging:
 
             with pytest.raises(ValueError, match="invalid literal for int"):
                 _ = envs.VLLM_CONFIGURE_LOGGING
+
+
+class TestVllmMaxNSequences:
+    def test_default_value(self):
+        """Test that VLLM_MAX_N_SEQUENCES defaults to 64."""
+        with patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("VLLM_MAX_N_SEQUENCES", None)
+            if hasattr(envs.__getattr__, "cache_clear"):
+                envs.__getattr__.cache_clear()
+
+            assert envs.VLLM_MAX_N_SEQUENCES == 16384
+
+    def test_custom_value(self, monkeypatch: pytest.MonkeyPatch):
+        """Test that VLLM_MAX_N_SEQUENCES can be overridden."""
+        monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
+        if hasattr(envs.__getattr__, "cache_clear"):
+            envs.__getattr__.cache_clear()
+
+        assert envs.VLLM_MAX_N_SEQUENCES == 128
+
+    def test_sampling_params_respects_limit(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Test that SamplingParams rejects n above the limit."""
+        from vllm.sampling_params import SamplingParams
+
+        monkeypatch.delenv("VLLM_MAX_N_SEQUENCES", raising=False)
+        if hasattr(envs.__getattr__, "cache_clear"):
+            envs.__getattr__.cache_clear()
+
+        max_n = envs.VLLM_MAX_N_SEQUENCES
+        SamplingParams(n=max_n)
+
+        with pytest.raises(ValueError, match="n must be at most"):
+            SamplingParams(n=max_n + 1)
+
+    def test_sampling_params_respects_custom_limit(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        """Test that SamplingParams uses the overridden env var limit."""
+        from vllm.sampling_params import SamplingParams
+
+        monkeypatch.setenv("VLLM_MAX_N_SEQUENCES", "128")
+        if hasattr(envs.__getattr__, "cache_clear"):
+            envs.__getattr__.cache_clear()
+
+        SamplingParams(n=128)
+
+        with pytest.raises(ValueError, match="n must be at most 128"):
+            SamplingParams(n=129)

tests/tool_parsers/test_hermes_tool_parser.py

@@ -152,6 +152,175 @@ def test_hermes_parser_streaming(
     }
 
 
+def _simulate_streaming(
+    tokenizer: TokenizerLike,
+    parser: ToolParser,
+    request: ChatCompletionRequest,
+    text: str,
+    stream_interval: int = 1,
+) -> list:
+    """Simulate streaming with a given stream_interval.
+
+    Tokens are batched into chunks of `stream_interval` tokens,
+    mimicking how the output processor delivers them.
+    Returns a list of non-None DeltaMessages.
+    """
+    tokens = tokenizer.encode(text)
+    previous_text = ""
+    delta_messages = []
+    for i in range(0, len(tokens), stream_interval):
+        chunk_ids = tokens[i : i + stream_interval]
+        delta_text = tokenizer.decode(chunk_ids)
+        current_text = previous_text + delta_text
+        delta = parser.extract_tool_calls_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=delta_text,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=chunk_ids,
+            request=request,
+        )
+        previous_text = current_text
+        if delta is not None:
+            delta_messages.append(delta)
+    return delta_messages
+
+
+@pytest.mark.parametrize("stream_interval", [2, 3, 5, 8])
+def test_hermes_streaming_tool_call_with_stream_interval(
+    qwen_tokenizer: TokenizerLike,
+    any_chat_request: ChatCompletionRequest,
+    stream_interval: int,
+) -> None:
+    """Tool call streaming must produce correct name + args at any interval."""
+    text = (
+        '<tool_call>{"name": "get_current_temperature", '
+        '"arguments": {"location": "San Francisco", "unit": "celsius"}}'
+        "</tool_call>"
+    )
+    parser = Hermes2ProToolParser(qwen_tokenizer)
+    deltas = _simulate_streaming(
+        qwen_tokenizer, parser, any_chat_request, text, stream_interval
+    )
+
+    # Flatten all DeltaToolCalls across all deltas.
+    tool_deltas = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
+    assert tool_deltas, "Expected at least one tool call delta"
+    assert tool_deltas[0].function.name == "get_current_temperature"
+
+    # Concatenated arguments must be valid JSON matching the original.
+    args_str = "".join(tc.function.arguments or "" for tc in tool_deltas)
+    assert json.loads(args_str) == {
+        "location": "San Francisco",
+        "unit": "celsius",
+    }
+
+
+@pytest.mark.parametrize("stream_interval", [2, 3, 5, 8])
+def test_hermes_streaming_content_then_tool_call_with_stream_interval(
+    qwen_tokenizer: TokenizerLike,
+    any_chat_request: ChatCompletionRequest,
+    stream_interval: int,
+) -> None:
+    """Content before a tool call must be fully streamed, then tool call."""
+    text = (
+        "Sure, let me check the weather."
+        '<tool_call>{"name": "get_weather", '
+        '"arguments": {"city": "NYC"}}</tool_call>'
+    )
+    parser = Hermes2ProToolParser(qwen_tokenizer)
+    deltas = _simulate_streaming(
+        qwen_tokenizer, parser, any_chat_request, text, stream_interval
+    )
+
+    content_deltas = [d for d in deltas if d.content]
+    tool_deltas = [d for d in deltas if d.tool_calls]
+
+    # Content must reconstruct the prefix.
+    content_str = "".join(d.content for d in content_deltas)
+    assert content_str == "Sure, let me check the weather."
+
+    # Tool call must be correct.
+    tool_calls = [tc for d in tool_deltas for tc in d.tool_calls]
+    assert tool_calls[0].function.name == "get_weather"
+    args_str = "".join(tc.function.arguments or "" for tc in tool_calls)
+    assert json.loads(args_str) == {"city": "NYC"}
+
+
+@pytest.mark.parametrize("stream_interval", [1, 2, 4])
+def test_hermes_streaming_multiple_tool_calls_with_stream_interval(
+    qwen_tokenizer: TokenizerLike,
+    any_chat_request: ChatCompletionRequest,
+    stream_interval: int,
+) -> None:
+    """Multiple sequential tool calls must each be streamed correctly."""
+    text = (
+        '<tool_call>{"name": "search", "arguments": {"q": "cats"}}</tool_call>'
+        '<tool_call>{"name": "search", "arguments": {"q": "dogs"}}</tool_call>'
+    )
+    parser = Hermes2ProToolParser(qwen_tokenizer)
+    deltas = _simulate_streaming(
+        qwen_tokenizer, parser, any_chat_request, text, stream_interval
+    )
+
+    # Flatten all DeltaToolCalls across all deltas.
+    all_tool_calls = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
+
+    # Separate by tool index.
+    tool0 = [tc for tc in all_tool_calls if tc.index == 0]
+    tool1 = [tc for tc in all_tool_calls if tc.index == 1]
+
+    assert tool0[0].function.name == "search"
+    args0 = "".join(tc.function.arguments or "" for tc in tool0)
+    assert json.loads(args0) == {"q": "cats"}
+
+    assert tool1[0].function.name == "search"
+    args1 = "".join(tc.function.arguments or "" for tc in tool1)
+    assert json.loads(args1) == {"q": "dogs"}
+
+
+@pytest.mark.parametrize("stream_interval", [2, 5])
+def test_hermes_streaming_boolean_args_with_stream_interval(
+    qwen_tokenizer: TokenizerLike,
+    any_chat_request: ChatCompletionRequest,
+    stream_interval: int,
+) -> None:
+    """Regression test for bug #19056 with stream_interval > 1."""
+    text = (
+        "<tool_call>\n"
+        '{"name": "final_answer", "arguments": {"trigger": true}}\n'
+        "</tool_call>"
+    )
+    parser = Hermes2ProToolParser(qwen_tokenizer)
+    deltas = _simulate_streaming(
+        qwen_tokenizer, parser, any_chat_request, text, stream_interval
+    )
+
+    tool_calls = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
+    assert tool_calls[0].function.name == "final_answer"
+    args_str = "".join(tc.function.arguments or "" for tc in tool_calls)
+    assert json.loads(args_str) == {"trigger": True}
+
+
+@pytest.mark.parametrize("stream_interval", [2, 3, 5])
+def test_hermes_streaming_just_forward_text_with_stream_interval(
+    qwen_tokenizer: TokenizerLike,
+    any_chat_request: ChatCompletionRequest,
+    stream_interval: int,
+) -> None:
+    """Plain text with no tool calls must be fully forwarded."""
+    text = "This is plain text with no tool calling involved."
+    parser = Hermes2ProToolParser(qwen_tokenizer)
+    deltas = _simulate_streaming(
+        qwen_tokenizer, parser, any_chat_request, text, stream_interval
+    )
+
+    for d in deltas:
+        assert not d.tool_calls
+    assert "".join(d.content for d in deltas) == text
+
+
 def test_hermes_parser_non_streaming_no_tool_call(
     hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
@@ -218,3 +387,28 @@ def test_hermes_parser_non_streaming_tool_call_invalid_json(
 
     assert tool_call is not None
     assert not tool_call.tools_called
+
+
+def test_hermes_streaming_content_and_tool_call_in_single_chunk(
+    qwen_tokenizer: TokenizerLike,
+    any_chat_request: ChatCompletionRequest,
+) -> None:
+    """Content + complete tool call in one chunk must both be emitted."""
+    text = 'Hi!<tool_call>{"name": "f", "arguments": {"x": 1}}</tool_call>'
+    # Use a stream_interval large enough to guarantee a single chunk.
+    parser = Hermes2ProToolParser(qwen_tokenizer)
+    deltas = _simulate_streaming(
+        qwen_tokenizer,
+        parser,
+        any_chat_request,
+        text,
+        stream_interval=9999,
+    )
+
+    content_parts = [d.content for d in deltas if d.content]
+    tool_parts = [tc for d in deltas if d.tool_calls for tc in d.tool_calls]
+
+    assert "".join(content_parts) == "Hi!"
+    assert tool_parts[0].function.name == "f"
+    args_str = "".join(tc.function.arguments or "" for tc in tool_parts)
+    assert json.loads(args_str) == {"x": 1}

tests/v1/core/test_scheduler.py

@@ -2568,13 +2568,19 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
 
     # Preempted request resumed in scheduled_cached_reqs
     assert len(scheduled_cached_reqs.resumed_req_ids) == 1
-    assert len(scheduled_cached_reqs.all_token_ids) == 1
     assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
     assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
-    assert request_low.request_id in scheduled_cached_reqs.all_token_ids
-    # Resumed tokens include 30 prompt tokens and 2 decoded tokens
-    assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32
-    assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100
+    # `output_token_ids` is only propagated for the v1 model runner.
+    if not scheduler.use_v2_model_runner:
+        assert len(scheduled_cached_reqs.output_token_ids) == 1
+        assert request_low.request_id in scheduled_cached_reqs.output_token_ids
+        # request_low produced two decode tokens before being preempted.
+        assert scheduled_cached_reqs.output_token_ids[request_low.request_id] == [
+            100,
+            100,
+        ]
+    else:
+        assert len(scheduled_cached_reqs.output_token_ids) == 0
 
 
 @pytest.mark.parametrize(
@@ -3422,13 +3428,19 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
 
     # Preempted request resumed in scheduled_cached_reqs
     assert len(scheduled_cached_reqs.resumed_req_ids) == 1
-    assert len(scheduled_cached_reqs.all_token_ids) == 1
     assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
     assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
-    assert request_low.request_id in scheduled_cached_reqs.all_token_ids
-    ## Resumed tokens include 94 prompt tokens and 2 decoded tokens
-    assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 96
-    assert scheduled_cached_reqs.all_token_ids[request_low.request_id][95] == 100
+    # `output_token_ids` is only propagated for the v1 model runner.
+    if not scheduler.use_v2_model_runner:
+        assert len(scheduled_cached_reqs.output_token_ids) == 1
+        assert request_low.request_id in scheduled_cached_reqs.output_token_ids
+        # request_low produced two decode tokens before being preempted.
+        assert scheduled_cached_reqs.output_token_ids[request_low.request_id] == [
+            100,
+            100,
+        ]
+    else:
+        assert len(scheduled_cached_reqs.output_token_ids) == 0
     assert scheduler.running[0].request_id == request_low.request_id
     assert request_high.request_id in output.finished_req_ids
 

tests/v1/engine/test_abort_final_step.py

@@ -259,8 +259,25 @@ async def test_abort_during_final_step(async_scheduling: bool):
                 # Wait for generation to complete
                 await gen_task
 
-                # Give the scheduler a moment to finish cleanup
-                await asyncio.sleep(0.1)
+                # Poll for the KV connector to record the finish status
+                timeout = 5.0
+                start = time.time()
+                captured_statuses = []
+                while time.time() - start < timeout:
+                    with open(status_file) as f4:
+                        status_lines = f4.read().strip().split("\n")
+                        captured_statuses = [
+                            line
+                            for line in status_lines
+                            if line and line.startswith("FINISHED_")
+                        ]
+                    if captured_statuses:
+                        break
+                    await asyncio.sleep(0.05)
+                else:
+                    raise TimeoutError(
+                        "Timeout waiting for KV connector to record finish status."
+                    )
 
                 # Verify we got output
                 assert len(outputs) > 0, "Should have received at least one output"
@@ -275,15 +292,6 @@ async def test_abort_during_final_step(async_scheduling: bool):
                     f"'{final_output.outputs[0].finish_reason}'. "
                 )
 
-                with open(status_file) as f4:
-                    status_lines = f4.read().strip().split("\n")
-                    # Filter for actual finish statuses (not INIT or empty lines)
-                    captured_statuses = [
-                        line
-                        for line in status_lines
-                        if line and line.startswith("FINISHED_")
-                    ]
-
                 assert len(captured_statuses) >= 1, (
                     f"Expected at least 1 captured finish status, got "
                     f"{len(captured_statuses)}. File content: {status_lines}"

tests/v1/kv_connector/unit/offloading_connector/conftest.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from tests.v1.kv_connector.unit.offloading_connector.utils import (
+    request_runner,
+)
+
+__all__ = ["request_runner"]

tests/v1/kv_connector/unit/offloading_connector/test_metrics.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
+    OffloadingConnectorStats,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+    OffloadingConnector,
+)
+
+
+def test_build_kv_connector_stats_with_none():
+    """Test that build_kv_connector_stats returns empty stats when given None."""
+    stats = OffloadingConnector.build_kv_connector_stats(data=None)
+
+    assert stats is not None
+    assert isinstance(stats, OffloadingConnectorStats)
+    assert len(stats.data) == 0
+    assert stats.is_empty()
+
+
+def test_build_kv_connector_stats_with_empty_dict():
+    """Test that build_kv_connector_stats returns empty stats with empty dict."""
+    stats = OffloadingConnector.build_kv_connector_stats(data={})
+
+    assert stats is not None
+    assert isinstance(stats, OffloadingConnectorStats)
+    assert len(stats.data) == 0
+    assert stats.is_empty()
+
+
+def test_build_kv_connector_stats_reconstructs_offload_stats():
+    """Test that OffloadingConnector stats are properly reconstructed with
+    correct data."""
+    serialized_data = {
+        "CPU_to_GPU": [
+            {"op_size": 16, "op_time": 1.0},
+            {"op_size": 8, "op_time": 0.5},
+        ],
+        "GPU_to_CPU": [
+            {"op_size": 1, "op_time": 0.1},
+            {"op_size": 2, "op_time": 0.2},
+        ],
+    }
+
+    stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
+
+    offload_connector_stats = stats
+    assert isinstance(offload_connector_stats, OffloadingConnectorStats)
+    assert offload_connector_stats.data["CPU_to_GPU"] == [
+        {"op_size": 16, "op_time": 1.0},
+        {"op_size": 8, "op_time": 0.5},
+    ]
+    assert offload_connector_stats.data["GPU_to_CPU"] == [
+        {"op_size": 1, "op_time": 0.1},
+        {"op_size": 2, "op_time": 0.2},
+    ]
+
+
+def test_aggregate_same_connector():
+    """Test aggregating stats from the same connector type."""
+    stats1 = OffloadingConnectorStats(
+        data={
+            "CPU_to_GPU": [
+                {"op_size": 16, "op_time": 1.0},
+                {"op_size": 8, "op_time": 0.5},
+            ],
+            "GPU_to_CPU": [
+                {"op_size": 1, "op_time": 0.1},
+                {"op_size": 2, "op_time": 0.2},
+            ],
+        }
+    )
+
+    stats2 = OffloadingConnectorStats(
+        data={
+            "CPU_to_GPU": [
+                {"op_size": 3, "op_time": 0.2},
+                {"op_size": 7, "op_time": 0.9},
+            ],
+            "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
+        }
+    )
+
+    result = stats1.aggregate(stats2)
+
+    assert result is stats1  # Should return self
+    offload_connector_stats = result
+    assert offload_connector_stats.data["CPU_to_GPU"] == [
+        {"op_size": 16, "op_time": 1.0},
+        {"op_size": 8, "op_time": 0.5},
+        {"op_size": 3, "op_time": 0.2},
+        {"op_size": 7, "op_time": 0.9},
+    ]
+    assert offload_connector_stats.data["GPU_to_CPU"] == [
+        {"op_size": 1, "op_time": 0.1},
+        {"op_size": 2, "op_time": 0.2},
+        {"op_size": 16, "op_time": 2},
+    ]
+
+
+def test_reduce():
+    """Test that reduce() correctly reduces all nested connector stats."""
+    stats = OffloadingConnectorStats(
+        data={
+            "CPU_to_GPU": [
+                {"op_size": 16, "op_time": 1.0},
+                {"op_size": 8, "op_time": 0.5},
+                {"op_size": 3, "op_time": 0.2},
+                {"op_size": 7, "op_time": 0.9},
+            ],
+            "GPU_to_CPU": [
+                {"op_size": 1, "op_time": 0.1},
+                {"op_size": 2, "op_time": 0.2},
+                {"op_size": 16, "op_time": 2},
+            ],
+        }
+    )
+
+    reduced = stats.reduce()
+
+    assert isinstance(reduced, dict)
+    # Check that the stats were reduced (should have aggregated values)
+    assert "CPU_to_GPU_total_bytes" in reduced
+    assert "CPU_to_GPU_total_time" in reduced
+    assert "GPU_to_CPU_total_bytes" in reduced
+    assert "GPU_to_CPU_total_time" in reduced
+    assert reduced["CPU_to_GPU_total_bytes"] == 34
+    assert reduced["CPU_to_GPU_total_time"] == 2.6
+    assert reduced["GPU_to_CPU_total_time"] == 2.3
+    assert reduced["GPU_to_CPU_total_bytes"] == 19
+
+
+def test_reset():
+    """Test that reset() resets all nested connector stats."""
+    offload_connector_stats = OffloadingConnectorStats(
+        data={
+            "CPU_to_GPU": [
+                {"op_size": 3, "op_time": 0.2},
+                {"op_size": 7, "op_time": 0.9},
+            ],
+            "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
+        }
+    )
+
+    assert not offload_connector_stats.is_empty()
+
+    offload_connector_stats.reset()
+
+    # After reset, stats should be empty
+    assert offload_connector_stats.is_empty()
+    assert len(offload_connector_stats.data) == 0

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+
+import pytest
+
+from tests.v1.kv_connector.unit.offloading_connector.utils import (
+    generate_store_output,
+)
+from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
+from vllm.distributed.kv_events import BlockRemoved, BlockStored
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_offload.abstract import OffloadingEvent
+from vllm.v1.request import RequestStatus
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_offloading_connector(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+    block_size_factor = offloaded_block_size // gpu_block_size
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    # 3 blocks, store just the middle block (skip first and last)
+    # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
+    runner.new_request(token_ids=[0] * offloaded_block_size * 3)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
+    )
+    runner.run(decoded_tokens=[0])
+
+    # add block missing 1 token -> no offload
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size - 1),
+        expected_stored_gpu_block_indexes=(3, 4, 5),
+    )
+    runner.manager.prepare_store.assert_not_called()
+
+    # +1 token -> single block, fail prepare_store
+    runner.manager.prepare_store.side_effect = lambda block_hashes: None
+    runner.run(decoded_tokens=[0])
+    runner.manager.prepare_store.assert_called()
+
+    # 1 more block (+ token for async scheduling)
+    # now set block_hashes_to_store = []
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
+
+    # 1 more block (+ token for kicking off offloading)
+    # now check touch was called with all 6 blocks
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0] * (offloaded_block_size + 1),
+        expected_stored_gpu_block_indexes=(15, 16, 17),
+    )
+    runner.manager.touch.assert_called()
+    block_hashes1 = list(runner.manager.touch.call_args.args[0])
+    assert len(block_hashes1) == 6
+
+    # terminate request
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    # create a new request differing only on the last token
+    runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
+    runner.run(decoded_tokens=[0])
+    runner.manager.touch.assert_called()
+    block_hashes2 = list(runner.manager.touch.call_args.args[0])
+    assert len(block_hashes2) == 6
+
+    # verify hashes are the same, except for the last block
+    assert block_hashes1[:5] == block_hashes2[:5]
+    assert block_hashes1[5] != block_hashes2[5]
+
+    # terminate request
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
+    )
+
+    # full_block_tokens - num_computed_tokens < offloaded_block_size
+    runner.new_request(
+        token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
+    )
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+    runner.manager.lookup.assert_not_called()
+
+    # single block lookup with no hits
+    runner.new_request(token_ids=[1] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+    runner.manager.lookup.assert_called()
+    assert len(list(runner.manager.lookup.call_args.args[0])) == 1
+
+    # single block lookup with a hit
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
+    )
+
+    # single block lookup with a hit in a middle block
+    runner.new_request(
+        token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
+    )
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
+    )
+
+    # test take_events
+    def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
+        return [BlockHash(str(i).encode()) for i in int_hashes]
+
+    def take_events() -> Iterable[OffloadingEvent]:
+        yield OffloadingEvent(
+            block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
+        )
+        yield OffloadingEvent(
+            block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
+        )
+
+    runner.manager.take_events.side_effect = take_events
+    events = list(runner.scheduler_connector.take_events())
+    assert len(events) == 2
+    event = events[0]
+    assert isinstance(event, BlockStored)
+    assert event.block_hashes == to_hashes([1, 2, 3])
+    assert event.block_size == 16
+    assert event.medium == "A"
+    assert event.token_ids == []
+    assert event.parent_block_hash is None
+    assert event.lora_id is None
+    assert event.lora_name is None
+    event = events[1]
+    assert isinstance(event, BlockRemoved)
+    assert event.block_hashes == to_hashes([4, 5, 6])
+    assert event.medium == "B"
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_request_preemption(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
+    num_free_blocks_empty = free_block_queue.num_free_blocks
+
+    # 2 blocks, store all, without flushing
+    # blocks = [0, 1, 2], [3, 4, 5]
+    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0],
+        complete_transfers=False,
+    )
+
+    # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
+        complete_transfers=False,
+    )
+
+    # simulate KV cache running out of space
+    free_block_queue.num_free_blocks = 0
+
+    # request should be preempted now
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+        expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
+        expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
+    )
+
+    # restore KV cache space and reset GPU prefix cache
+    free_block_queue.num_free_blocks = num_free_blocks_empty
+    runner.scheduler.reset_prefix_cache()
+
+    # request should now return from preemption
+    # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
+    runner.manager.lookup.return_value = 3
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[0] * gpu_block_size,
+        expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
+    )
+
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(9, 10, 11),
+    )
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    # store 1 blocks
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # start a request to load the first block, but don't complete
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+    )
+
+    # request triggered a load
+    transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
+    assert transfer_jobs
+
+    # start a new request to load the same first block
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+    )
+
+    # request did not trigger a load
+    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
+
+    # complete transfers
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output([])
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_loaded_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # second request will use the GPU prefix cache
+    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_abort_loading_requests(request_runner, async_scheduling: bool):
+    offloaded_block_size = 12
+    gpu_block_size = 4
+    num_gpu_blocks = 100
+
+    runner = request_runner(
+        offloaded_block_size=offloaded_block_size,
+        gpu_block_size=gpu_block_size,
+        num_gpu_blocks=num_gpu_blocks,
+        async_scheduling=async_scheduling,
+    )
+
+    # store 1 blocks
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.prepare_store.side_effect = (
+        lambda block_hashes: generate_store_output(block_hashes)
+    )
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # start a request to load the first block, but don't complete
+    runner.scheduler.reset_prefix_cache()
+    runner.new_request(token_ids=[0] * offloaded_block_size)
+    runner.manager.lookup.return_value = 1
+    runner.run(
+        decoded_tokens=[],
+        complete_transfers=False,
+    )
+
+    # request triggered a load
+    transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
+    assert transfer_jobs
+
+    # abort request
+    req_id = str(runner.req_id)
+    runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
+
+    # verify request is not deleted
+    assert req_id in runner.scheduler.requests
+
+    # complete loading request
+    runner.run(
+        decoded_tokens=[],
+        expected_loaded_gpu_block_indexes=(0, 1, 2),
+    )
+
+    # assert request is deleted
+    assert req_id not in runner.scheduler.requests

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

@@ -0,0 +1,504 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections import defaultdict
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import get_dtype_size
+from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    MambaSpec,
+    MLAAttentionSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.kv_offload.spec import (
+    CanonicalKVCacheRef,
+    CanonicalKVCaches,
+    OffloadingSpec,
+)
+
+NUM_BLOCKS = 10
+BLOCK_SIZE = 16
+NUM_KV_HEADS = 4
+HEAD_SIZE = 64
+DTYPE = torch.float16
+
+# Attention backends to test
+ATTN_BACKENDS: list[str] = []
+if current_platform.is_cuda():
+    ATTN_BACKENDS = [
+        "FLASH_ATTN",
+        "FLEX_ATTENTION",
+        "FLASHINFER",
+        "TRITON_ATTN",
+    ]
+elif current_platform.is_rocm():
+    ATTN_BACKENDS = ["TRITON_ATTN"]
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _allocate_and_reshape_kv_caches(
+    kv_cache_config: KVCacheConfig,
+    attn_groups: list[list],
+    device: torch.device,
+):
+    """
+    Use the real GPUModelRunner allocation and reshape methods to produce
+    kv_caches, just like the model runner does during initialization.
+    """
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+    # Some backends (e.g. FlashAttention) query the KV cache layout during
+    # reshape, which ultimately calls get_current_vllm_config(). Setting
+    # the layout override avoids needing a full VllmConfig context.
+    set_kv_cache_layout("NHD")
+    try:
+        runner = object.__new__(GPUModelRunner)
+        runner.device = device
+        runner.runner_only_attn_layers = set()
+        runner.attn_groups = attn_groups
+        runner.kv_cache_config = kv_cache_config
+        runner.cache_config = MagicMock(cache_dtype="auto")
+        runner.shared_kv_cache_layers = {}
+        runner.model_config = MagicMock()
+        runner.model_config.hf_config.model_type = ""
+        runner.compilation_config = MagicMock(
+            static_forward_context=defaultdict(MagicMock)
+        )
+        runner.kv_caches = []
+
+        kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
+        return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes)
+    finally:
+        set_kv_cache_layout(None)
+
+
+def _make_mock_layer(backend_cls: type[AttentionBackend]):
+    """
+    Create a mock AttentionLayerBase whose get_attn_backend returns backend_cls.
+    """
+    layer = MagicMock()
+    layer.get_attn_backend.return_value = backend_cls
+    return layer
+
+
+def _make_worker(kv_cache_config: KVCacheConfig):
+    """
+    Create an OffloadingConnectorWorker with mocked dependencies.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
+        OffloadingConnectorWorker,
+    )
+
+    spec = MagicMock(spec=OffloadingSpec)
+    spec.kv_cache_config = kv_cache_config
+    spec.vllm_config = MagicMock()
+    spec.get_handlers.return_value = iter([])
+
+    worker = OffloadingConnectorWorker(spec=spec)
+    worker.worker = MagicMock()
+
+    return worker, spec
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("backend", ATTN_BACKENDS)
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
+    ".worker.get_layers_from_vllm_config"
+)
+def test_register_kv_caches(mock_get_layers, backend):
+    """Test register_kv_caches with multiple groups covering all layer types.
+
+    Creates one FullAttention group, one MLA group, one Mamba group, and
+    one Mamba-padded group. Each group has GROUP_SIZE layers.
+
+    KVCacheTensors are shared across all groups mirroring the real allocation
+    in kv_cache_utils.py: tensor i is shared by layer i from every group.
+    The padded-mamba group has a different page size so its layers get their
+    own dedicated tensors.
+
+    Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
+    kv_caches, which automatically applies
+    _update_hybrid_attention_mamba_layout for hybrid models.
+
+    Verifies that the canonicalized CanonicalKVCaches has the correct
+    block tensors, tensor_idx references, and page sizes across all groups.
+    """
+    from vllm.v1.attention.backends.mla.indexer import (
+        DeepseekV32IndexerBackend,
+    )
+    from vllm.v1.worker.utils import AttentionGroup
+
+    MLA_HEAD_SIZE = NUM_KV_HEADS * HEAD_SIZE * 2
+
+    # padded mamba (missing HEAD_SIZE)
+    CONV_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS, HEAD_SIZE)
+    UNALIGNED_SSM_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS - 1, HEAD_SIZE)
+
+    PAGE_SIZE_BYTES = 2 * BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * get_dtype_size(DTYPE)
+    unaligned_mamba_page_size = PAGE_SIZE_BYTES - HEAD_SIZE * get_dtype_size(DTYPE)
+
+    # unpadded mamba (fills page exactly)
+    ALIGNED_SSM_STATE_SHAPE = (BLOCK_SIZE * NUM_KV_HEADS, HEAD_SIZE)
+
+    backend_cls = AttentionBackendEnum[backend].get_class()
+
+    attn_spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    mla_spec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=1,
+        head_size=MLA_HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    unaligned_mamba_spec = MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=(CONV_STATE_SHAPE, UNALIGNED_SSM_STATE_SHAPE),
+        dtypes=(DTYPE, DTYPE),
+        page_size_padded=PAGE_SIZE_BYTES,
+    )
+    aligned_mamba_spec = MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=(CONV_STATE_SHAPE, ALIGNED_SSM_STATE_SHAPE),
+        dtypes=(DTYPE, DTYPE),
+        page_size_padded=PAGE_SIZE_BYTES,
+    )
+
+    assert attn_spec.page_size_bytes == PAGE_SIZE_BYTES
+    assert mla_spec.page_size_bytes == PAGE_SIZE_BYTES
+    assert unaligned_mamba_spec.page_size_bytes == PAGE_SIZE_BYTES
+    assert aligned_mamba_spec.page_size_bytes == PAGE_SIZE_BYTES
+
+    GROUP_SIZE = 3
+
+    # -- Build per-group layer info ----------------------------------------
+    layer_idx = 0
+
+    attn_layer_names = []
+    for _ in range(GROUP_SIZE):
+        attn_layer_names.append(f"model.layers.{layer_idx}.self_attn")
+        layer_idx += 1
+
+    mla_layer_names = []
+    for _ in range(GROUP_SIZE):
+        mla_layer_names.append(f"model.layers.{layer_idx}.self_attn")
+        layer_idx += 1
+
+    unaligned_mamba_layer_names = []
+    for _ in range(GROUP_SIZE):
+        unaligned_mamba_layer_names.append(f"model.layers.{layer_idx}.mamba_unpadded")
+        layer_idx += 1
+
+    aligned_mamba_layer_names = []
+    for _ in range(GROUP_SIZE - 1):
+        aligned_mamba_layer_names.append(f"model.layers.{layer_idx}.mamba_padded")
+        layer_idx += 1
+
+    layer_groups = [
+        attn_layer_names,
+        mla_layer_names,
+        unaligned_mamba_layer_names,
+        aligned_mamba_layer_names,
+    ]
+
+    kv_cache_tensors: list[KVCacheTensor] = []
+    for i in range(GROUP_SIZE):
+        shared_by: list[str] = []
+        for group_layer_names in layer_groups:
+            if len(group_layer_names) > i:
+                shared_by.append(group_layer_names[i])
+        kv_cache_tensors.append(
+            KVCacheTensor(
+                size=PAGE_SIZE_BYTES * NUM_BLOCKS,
+                shared_by=shared_by,
+            )
+        )
+
+    kv_cache_groups = [
+        KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
+        KVCacheGroupSpec(layer_names=mla_layer_names, kv_cache_spec=mla_spec),
+        KVCacheGroupSpec(
+            layer_names=unaligned_mamba_layer_names, kv_cache_spec=unaligned_mamba_spec
+        ),
+        KVCacheGroupSpec(
+            layer_names=aligned_mamba_layer_names, kv_cache_spec=aligned_mamba_spec
+        ),
+    ]
+
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=backend_cls,
+                layer_names=attn_layer_names,
+                kv_cache_spec=attn_spec,
+                kv_cache_group_id=0,
+            ),
+            AttentionGroup(
+                backend=DeepseekV32IndexerBackend,
+                layer_names=mla_layer_names,
+                kv_cache_spec=mla_spec,
+                kv_cache_group_id=1,
+            ),
+            AttentionGroup(
+                backend=DeepseekV32IndexerBackend,  # unused for mamba
+                layer_names=unaligned_mamba_layer_names,
+                kv_cache_spec=unaligned_mamba_spec,
+                kv_cache_group_id=2,
+            ),
+            AttentionGroup(
+                backend=DeepseekV32IndexerBackend,  # unused for mamba
+                layer_names=aligned_mamba_layer_names,
+                kv_cache_spec=aligned_mamba_spec,
+                kv_cache_group_id=3,
+            ),
+        ]
+    ]
+
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=kv_cache_tensors,
+        kv_cache_groups=kv_cache_groups,
+    )
+
+    kv_caches = _allocate_and_reshape_kv_caches(
+        kv_cache_config,
+        attn_groups,
+        device=torch.device("cuda:0"),
+    )
+
+    mock_layers: dict[str, MagicMock] = {}
+    for layer_name in attn_layer_names:
+        mock_layers[layer_name] = _make_mock_layer(backend_cls)
+    for layer_name in mla_layer_names:
+        mock_layers[layer_name] = _make_mock_layer(DeepseekV32IndexerBackend)
+    mock_get_layers.return_value = mock_layers
+
+    worker, spec = _make_worker(kv_cache_config)
+    worker.register_kv_caches(kv_caches)
+
+    canonical = spec.get_handlers.call_args[0][0]
+    assert isinstance(canonical, CanonicalKVCaches)
+
+    # -- Expected block tensors ----------------------------------------------
+    # All tensors have the same padded page size (PAGE_SIZE_BYTES).
+    # Tensor 0: shared by attn[0], mla[0], mamba_unaligned[0], mamba_aligned[0]
+    # Tensor 1: shared by attn[1], mla[1], mamba_unaligned[1], mamba_aligned[1]
+    # Tensor 2: shared by attn[2], mla[2], mamba_unaligned[2]
+    #           (mamba_aligned has only GROUP_SIZE-1 = 2 layers)
+    expected_tensors = [
+        (NUM_BLOCKS, PAGE_SIZE_BYTES),
+        (NUM_BLOCKS, PAGE_SIZE_BYTES),
+        (NUM_BLOCKS, PAGE_SIZE_BYTES),
+    ]
+
+    # -- Expected group data refs (order matches kv_cache_groups) -------------
+    ref = CanonicalKVCacheRef
+    expected_group_refs = [
+        # attn group: layers attn[0..2] → tensors 0,1,2 with full page size
+        [
+            ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
+            ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
+            ref(tensor_idx=2, page_size_bytes=PAGE_SIZE_BYTES),
+        ],
+        # mla group: layers mla[0..2] → tensors 0,1,2 with full page size
+        [
+            ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
+            ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
+            ref(tensor_idx=2, page_size_bytes=PAGE_SIZE_BYTES),
+        ],
+        # unaligned mamba group: layers [0..2] → tensors 0,1,2 with unaligned page
+        [
+            ref(tensor_idx=0, page_size_bytes=unaligned_mamba_page_size),
+            ref(tensor_idx=1, page_size_bytes=unaligned_mamba_page_size),
+            ref(tensor_idx=2, page_size_bytes=unaligned_mamba_page_size),
+        ],
+        # aligned mamba group: layers [0..1] → tensors 0,1 with full page size
+        [
+            ref(tensor_idx=0, page_size_bytes=PAGE_SIZE_BYTES),
+            ref(tensor_idx=1, page_size_bytes=PAGE_SIZE_BYTES),
+        ],
+    ]
+
+    # Verify block tensors
+    assert len(canonical.tensors) == len(expected_tensors)
+    for block_tensor, (exp_num_blocks, exp_page_size) in zip(
+        canonical.tensors, expected_tensors
+    ):
+        tensor = block_tensor.tensor
+        assert tensor.dtype == torch.int8
+        assert tensor.shape == (exp_num_blocks, exp_page_size)
+        assert block_tensor.page_size_bytes == exp_page_size
+
+    # Verify group data refs
+    assert len(canonical.group_data_refs) == len(expected_group_refs)
+    for actual_refs, exp_refs in zip(canonical.group_data_refs, expected_group_refs):
+        assert len(actual_refs) == len(exp_refs)
+        for actual, expected in zip(actual_refs, exp_refs):
+            assert actual.tensor_idx == expected.tensor_idx
+            assert actual.page_size_bytes == expected.page_size_bytes
+
+
+@pytest.mark.parametrize("backend", ATTN_BACKENDS)
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
+    ".worker.get_layers_from_vllm_config"
+)
+def test_register_kv_caches_uniform_type(mock_get_layers, backend):
+    """Test register_kv_caches with UniformTypeKVCacheSpecs.
+
+    Two attention layers use the same backend but different num_kv_heads,
+    giving them different per-layer page sizes. Each has its own
+    KVCacheTensor and are wrapped in a UniformTypeKVCacheSpecs group.
+    Verifies that each layer gets the correct tensor_idx and
+    page_size_bytes in its block data ref.
+    """
+    from vllm.v1.worker.utils import AttentionGroup
+
+    backend_cls = AttentionBackendEnum[backend].get_class()
+
+    layer_a = "model.layers.0.self_attn"
+    layer_b = "model.layers.1.self_attn"
+    spec_a = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    spec_b = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS * 2,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
+    )
+    assert spec_a.page_size_bytes != spec_b.page_size_bytes
+
+    uniform_spec = UniformTypeKVCacheSpecs(
+        block_size=BLOCK_SIZE,
+        kv_cache_specs={layer_a: spec_a, layer_b: spec_b},
+    )
+
+    kv_cache_config = KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=spec_a.page_size_bytes * NUM_BLOCKS,
+                shared_by=[layer_a],
+            ),
+            KVCacheTensor(
+                size=spec_b.page_size_bytes * NUM_BLOCKS,
+                shared_by=[layer_b],
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=[layer_a, layer_b],
+                kv_cache_spec=uniform_spec,
+            )
+        ],
+    )
+
+    attn_groups = [
+        [
+            AttentionGroup(
+                backend=backend_cls,
+                layer_names=[layer_a],
+                kv_cache_spec=spec_a,
+                kv_cache_group_id=0,
+            ),
+            AttentionGroup(
+                backend=backend_cls,
+                layer_names=[layer_b],
+                kv_cache_spec=spec_b,
+                kv_cache_group_id=0,
+            ),
+        ]
+    ]
+
+    kv_caches = _allocate_and_reshape_kv_caches(
+        kv_cache_config,
+        attn_groups,
+        device=torch.device("cuda:0"),
+    )
+
+    mock_get_layers.return_value = {
+        layer_a: _make_mock_layer(backend_cls),
+        layer_b: _make_mock_layer(backend_cls),
+    }
+
+    worker, spec = _make_worker(kv_cache_config)
+    worker.register_kv_caches(kv_caches)
+
+    canonical = spec.get_handlers.call_args[0][0]
+    assert isinstance(canonical, CanonicalKVCaches)
+
+    unbinds = backend_cls.get_name() in ("FLASH_ATTN", "FLEX_ATTENTION")
+    tensors_per_layer = 2 if unbinds else 1
+
+    for block_tensor in canonical.tensors:
+        assert block_tensor.tensor.dtype == torch.int8
+
+    # Single group with refs from both layers
+    assert len(canonical.group_data_refs) == 1
+    group_refs = canonical.group_data_refs[0]
+    assert len(group_refs) == 2 * tensors_per_layer
+
+    if unbinds:
+        half_a = spec_a.page_size_bytes // 2
+        half_b = spec_b.page_size_bytes // 2
+
+        assert len(canonical.tensors) == 4
+        assert canonical.tensors[0].page_size_bytes == half_a
+        assert canonical.tensors[1].page_size_bytes == half_a
+        assert canonical.tensors[2].page_size_bytes == half_b
+        assert canonical.tensors[3].page_size_bytes == half_b
+        assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, half_a)
+        assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, half_a)
+        assert canonical.tensors[2].tensor.shape == (NUM_BLOCKS, half_b)
+        assert canonical.tensors[3].tensor.shape == (NUM_BLOCKS, half_b)
+
+        assert group_refs[0] == CanonicalKVCacheRef(
+            tensor_idx=0, page_size_bytes=half_a
+        )
+        assert group_refs[1] == CanonicalKVCacheRef(
+            tensor_idx=1, page_size_bytes=half_a
+        )
+        assert group_refs[2] == CanonicalKVCacheRef(
+            tensor_idx=2, page_size_bytes=half_b
+        )
+        assert group_refs[3] == CanonicalKVCacheRef(
+            tensor_idx=3, page_size_bytes=half_b
+        )
+    else:
+        assert len(canonical.tensors) == 2
+        assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
+        assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
+        assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
+        assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
+
+        assert group_refs[0] == CanonicalKVCacheRef(
+            tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
+        )
+        assert group_refs[1] == CanonicalKVCacheRef(
+            tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
+        )

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -9,16 +9,17 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 
+from tests.v1.kv_connector.unit.utils import (
+    EOS_TOKEN_ID,
+    create_model_runner_output,
+    create_vllm_config,
+)
 from vllm import SamplingParams
-from vllm.config import KVTransferConfig, VllmConfig
-from vllm.distributed.kv_events import BlockRemoved, BlockStored
+from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
     OffloadingConnectorMetadata,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
-    OffloadingConnectorStats,
-)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
     OffloadingConnector,
 )
@@ -39,7 +40,6 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.kv_offload.abstract import (
     LoadStoreSpec,
-    OffloadingEvent,
     OffloadingManager,
     PrepareStoreOutput,
 )
@@ -51,15 +51,9 @@ from vllm.v1.kv_offload.worker.worker import (
     TransferSpec,
 )
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput
-from vllm.v1.request import Request, RequestStatus
+from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager
 
-from .utils import (
-    EOS_TOKEN_ID,
-    create_model_runner_output,
-    create_vllm_config,
-)
-
 
 class MockLoadStoreSpec(LoadStoreSpec):
     def __init__(self, block_hashes: Iterable[BlockHash]):
@@ -125,7 +119,7 @@ class MockOffloadingSpec(OffloadingSpec):
         return self.manager
 
     def get_handlers(
-        self, _, __
+        self, _
     ) -> Iterator[tuple[type[LoadStoreSpec], type[LoadStoreSpec], OffloadingHandler]]:
         yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler
         yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler
@@ -179,7 +173,7 @@ class RequestRunner:
             kv_role="kv_both",
             kv_connector_extra_config={
                 "spec_name": "MockOffloadingSpec",
-                "spec_module_path": "tests.v1.kv_connector.unit.test_offloading_connector",  # noqa: E501
+                "spec_module_path": "tests.v1.kv_connector.unit.offloading_connector.utils",  # noqa: E501
                 "block_size": offloaded_block_size,
             },
         )
@@ -217,10 +211,12 @@ class RequestRunner:
         )
 
         # register worker kv_caches to enable OffloadingWorker creations
-        self.worker_connector.register_cross_layers_kv_cache(
-            kv_cache=torch.empty(0),
-            attn_backend=FlashAttentionBackend,
-        )
+        # set_current_vllm_config is needed for get_kv_cache_layout() to work
+        with set_current_vllm_config(vllm_config):
+            self.worker_connector.register_cross_layers_kv_cache(
+                kv_cache=torch.empty(0),
+                attn_backend=FlashAttentionBackend,
+            )
 
         # extract connector of scheduler
         scheduler_connector = self.scheduler.connector
@@ -521,471 +517,3 @@ def generate_store_output(block_hashes: Iterable[BlockHash]):
         store_spec=MockLoadStoreSpec(block_hashes),
         block_hashes_evicted=[],
     )
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_offloading_connector(request_runner, async_scheduling: bool):
-    offloaded_block_size = 12
-    gpu_block_size = 4
-    num_gpu_blocks = 100
-    block_size_factor = offloaded_block_size // gpu_block_size
-
-    runner = request_runner(
-        offloaded_block_size=offloaded_block_size,
-        gpu_block_size=gpu_block_size,
-        num_gpu_blocks=num_gpu_blocks,
-        async_scheduling=async_scheduling,
-    )
-
-    # 3 blocks, store just the middle block (skip first and last)
-    # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8]
-    runner.new_request(token_ids=[0] * offloaded_block_size * 3)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(list(block_hashes)[1:2])
-    )
-    runner.run(decoded_tokens=[0])
-
-    # add block missing 1 token -> no offload
-    runner.run(
-        decoded_tokens=[0] * (offloaded_block_size - 1),
-        expected_stored_gpu_block_indexes=(3, 4, 5),
-    )
-    runner.manager.prepare_store.assert_not_called()
-
-    # +1 token -> single block, fail prepare_store
-    runner.manager.prepare_store.side_effect = lambda block_hashes: None
-    runner.run(decoded_tokens=[0])
-    runner.manager.prepare_store.assert_called()
-
-    # 1 more block (+ token for async scheduling)
-    # now set block_hashes_to_store = []
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output([])
-    )
-    runner.run(decoded_tokens=[0] * (offloaded_block_size + 1))
-
-    # 1 more block (+ token for kicking off offloading)
-    # now check touch was called with all 6 blocks
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(block_hashes)
-    )
-    runner.run(
-        decoded_tokens=[0] * (offloaded_block_size + 1),
-        expected_stored_gpu_block_indexes=(15, 16, 17),
-    )
-    runner.manager.touch.assert_called()
-    block_hashes1 = list(runner.manager.touch.call_args.args[0])
-    assert len(block_hashes1) == 6
-
-    # terminate request
-    runner.run(decoded_tokens=[EOS_TOKEN_ID])
-
-    # create a new request differing only on the last token
-    runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1])
-    runner.run(decoded_tokens=[0])
-    runner.manager.touch.assert_called()
-    block_hashes2 = list(runner.manager.touch.call_args.args[0])
-    assert len(block_hashes2) == 6
-
-    # verify hashes are the same, except for the last block
-    assert block_hashes1[:5] == block_hashes2[:5]
-    assert block_hashes1[5] != block_hashes2[5]
-
-    # terminate request
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=tuple(range(6 * block_size_factor)),
-    )
-
-    # full_block_tokens - num_computed_tokens < offloaded_block_size
-    runner.new_request(
-        token_ids=[0] * gpu_block_size + [1] * (offloaded_block_size - gpu_block_size)
-    )
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output([])
-    )
-    runner.run(decoded_tokens=[EOS_TOKEN_ID])
-    runner.manager.lookup.assert_not_called()
-
-    # single block lookup with no hits
-    runner.new_request(token_ids=[1] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output([])
-    )
-    runner.run(decoded_tokens=[EOS_TOKEN_ID])
-    runner.manager.lookup.assert_called()
-    assert len(list(runner.manager.lookup.call_args.args[0])) == 1
-
-    # single block lookup with a hit
-    runner.scheduler.reset_prefix_cache()
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output([])
-    )
-    runner.manager.lookup.return_value = 1
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)
-    )
-
-    # single block lookup with a hit in a middle block
-    runner.new_request(
-        token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size
-    )
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output([])
-    )
-    runner.manager.lookup.return_value = 1
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)
-    )
-
-    # test take_events
-    def to_hashes(int_hashes: list[int]) -> list[BlockHash]:
-        return [BlockHash(str(i).encode()) for i in int_hashes]
-
-    def take_events() -> Iterable[OffloadingEvent]:
-        yield OffloadingEvent(
-            block_hashes=to_hashes([1, 2, 3]), block_size=16, medium="A", removed=False
-        )
-        yield OffloadingEvent(
-            block_hashes=to_hashes([4, 5, 6]), block_size=32, medium="B", removed=True
-        )
-
-    runner.manager.take_events.side_effect = take_events
-    events = list(runner.scheduler_connector.take_events())
-    assert len(events) == 2
-    event = events[0]
-    assert isinstance(event, BlockStored)
-    assert event.block_hashes == to_hashes([1, 2, 3])
-    assert event.block_size == 16
-    assert event.medium == "A"
-    assert event.token_ids == []
-    assert event.parent_block_hash is None
-    assert event.lora_id is None
-    assert event.lora_name is None
-    event = events[1]
-    assert isinstance(event, BlockRemoved)
-    assert event.block_hashes == to_hashes([4, 5, 6])
-    assert event.medium == "B"
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_request_preemption(request_runner, async_scheduling: bool):
-    offloaded_block_size = 12
-    gpu_block_size = 4
-    num_gpu_blocks = 100
-
-    runner = request_runner(
-        offloaded_block_size=offloaded_block_size,
-        gpu_block_size=gpu_block_size,
-        num_gpu_blocks=num_gpu_blocks,
-        async_scheduling=async_scheduling,
-    )
-
-    free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
-    num_free_blocks_empty = free_block_queue.num_free_blocks
-
-    # 2 blocks, store all, without flushing
-    # blocks = [0, 1, 2], [3, 4, 5]
-    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(block_hashes)
-    )
-    runner.run(
-        decoded_tokens=[0],
-        complete_transfers=False,
-    )
-
-    # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(block_hashes)
-    )
-    runner.run(
-        decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),
-        complete_transfers=False,
-    )
-
-    # simulate KV cache running out of space
-    free_block_queue.num_free_blocks = 0
-
-    # request should be preempted now
-    runner.run(
-        decoded_tokens=[],
-        complete_transfers=False,
-        expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
-        expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
-    )
-
-    # restore KV cache space and reset GPU prefix cache
-    free_block_queue.num_free_blocks = num_free_blocks_empty
-    runner.scheduler.reset_prefix_cache()
-
-    # request should now return from preemption
-    # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
-    runner.manager.lookup.return_value = 3
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(block_hashes)
-    )
-    runner.run(
-        decoded_tokens=[0] * gpu_block_size,
-        expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),
-    )
-
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(9, 10, 11),
-    )
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
-    offloaded_block_size = 12
-    gpu_block_size = 4
-    num_gpu_blocks = 100
-
-    runner = request_runner(
-        offloaded_block_size=offloaded_block_size,
-        gpu_block_size=gpu_block_size,
-        num_gpu_blocks=num_gpu_blocks,
-        async_scheduling=async_scheduling,
-    )
-
-    # store 1 blocks
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(block_hashes)
-    )
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # start a request to load the first block, but don't complete
-    runner.scheduler.reset_prefix_cache()
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.lookup.return_value = 1
-    runner.run(
-        decoded_tokens=[],
-        complete_transfers=False,
-    )
-
-    # request triggered a load
-    transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
-    assert transfer_jobs
-
-    # start a new request to load the same first block
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.lookup.return_value = 1
-    runner.run(
-        decoded_tokens=[],
-        complete_transfers=False,
-    )
-
-    # request did not trigger a load
-    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
-
-    # complete transfers
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output([])
-    )
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_loaded_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # second request will use the GPU prefix cache
-    assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_abort_loading_requests(request_runner, async_scheduling: bool):
-    offloaded_block_size = 12
-    gpu_block_size = 4
-    num_gpu_blocks = 100
-
-    runner = request_runner(
-        offloaded_block_size=offloaded_block_size,
-        gpu_block_size=gpu_block_size,
-        num_gpu_blocks=num_gpu_blocks,
-        async_scheduling=async_scheduling,
-    )
-
-    # store 1 blocks
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.prepare_store.side_effect = (
-        lambda block_hashes: generate_store_output(block_hashes)
-    )
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        expected_stored_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # start a request to load the first block, but don't complete
-    runner.scheduler.reset_prefix_cache()
-    runner.new_request(token_ids=[0] * offloaded_block_size)
-    runner.manager.lookup.return_value = 1
-    runner.run(
-        decoded_tokens=[],
-        complete_transfers=False,
-    )
-
-    # request triggered a load
-    transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)
-    assert transfer_jobs
-
-    # abort request
-    req_id = str(runner.req_id)
-    runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
-
-    # verify request is not deleted
-    assert req_id in runner.scheduler.requests
-
-    # complete loading request
-    runner.run(
-        decoded_tokens=[],
-        expected_loaded_gpu_block_indexes=(0, 1, 2),
-    )
-
-    # assert request is deleted
-    assert req_id not in runner.scheduler.requests
-
-
-class TestOffloadingConnectorStats:
-    """Tests for OffloadingConnector stats reconstruction and operations."""
-
-    def test_build_kv_connector_stats_with_none(self):
-        """Test that build_kv_connector_stats returns empty stats when given None."""
-        stats = OffloadingConnector.build_kv_connector_stats(data=None)
-
-        assert stats is not None
-        assert isinstance(stats, OffloadingConnectorStats)
-        assert len(stats.data) == 0
-        assert stats.is_empty()
-
-    def test_build_kv_connector_stats_with_empty_dict(self):
-        """Test that build_kv_connector_stats returns empty stats with empty dict."""
-        stats = OffloadingConnector.build_kv_connector_stats(data={})
-
-        assert stats is not None
-        assert isinstance(stats, OffloadingConnectorStats)
-        assert len(stats.data) == 0
-        assert stats.is_empty()
-
-    def test_build_kv_connector_stats_reconstructs_offload_stats(self):
-        """Test that OffloadingConnector stats are properly reconstructed with
-        correct data."""
-        serialized_data = {
-            "CPU_to_GPU": [
-                {"op_size": 16, "op_time": 1.0},
-                {"op_size": 8, "op_time": 0.5},
-            ],
-            "GPU_to_CPU": [
-                {"op_size": 1, "op_time": 0.1},
-                {"op_size": 2, "op_time": 0.2},
-            ],
-        }
-
-        stats = OffloadingConnector.build_kv_connector_stats(data=serialized_data)
-
-        offload_connector_stats = stats
-        assert isinstance(offload_connector_stats, OffloadingConnectorStats)
-        assert offload_connector_stats.data["CPU_to_GPU"] == [
-            {"op_size": 16, "op_time": 1.0},
-            {"op_size": 8, "op_time": 0.5},
-        ]
-        assert offload_connector_stats.data["GPU_to_CPU"] == [
-            {"op_size": 1, "op_time": 0.1},
-            {"op_size": 2, "op_time": 0.2},
-        ]
-
-    def test_aggregate_same_connector(self):
-        """Test aggregating stats from the same connector type."""
-        stats1 = OffloadingConnectorStats(
-            data={
-                "CPU_to_GPU": [
-                    {"op_size": 16, "op_time": 1.0},
-                    {"op_size": 8, "op_time": 0.5},
-                ],
-                "GPU_to_CPU": [
-                    {"op_size": 1, "op_time": 0.1},
-                    {"op_size": 2, "op_time": 0.2},
-                ],
-            }
-        )
-
-        stats2 = OffloadingConnectorStats(
-            data={
-                "CPU_to_GPU": [
-                    {"op_size": 3, "op_time": 0.2},
-                    {"op_size": 7, "op_time": 0.9},
-                ],
-                "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
-            }
-        )
-
-        result = stats1.aggregate(stats2)
-
-        assert result is stats1  # Should return self
-        offload_connector_stats = result
-        assert offload_connector_stats.data["CPU_to_GPU"] == [
-            {"op_size": 16, "op_time": 1.0},
-            {"op_size": 8, "op_time": 0.5},
-            {"op_size": 3, "op_time": 0.2},
-            {"op_size": 7, "op_time": 0.9},
-        ]
-        assert offload_connector_stats.data["GPU_to_CPU"] == [
-            {"op_size": 1, "op_time": 0.1},
-            {"op_size": 2, "op_time": 0.2},
-            {"op_size": 16, "op_time": 2},
-        ]
-
-    def test_reduce(self):
-        """Test that reduce() correctly reduces all nested connector stats."""
-        stats = OffloadingConnectorStats(
-            data={
-                "CPU_to_GPU": [
-                    {"op_size": 16, "op_time": 1.0},
-                    {"op_size": 8, "op_time": 0.5},
-                    {"op_size": 3, "op_time": 0.2},
-                    {"op_size": 7, "op_time": 0.9},
-                ],
-                "GPU_to_CPU": [
-                    {"op_size": 1, "op_time": 0.1},
-                    {"op_size": 2, "op_time": 0.2},
-                    {"op_size": 16, "op_time": 2},
-                ],
-            }
-        )
-
-        reduced = stats.reduce()
-
-        assert isinstance(reduced, dict)
-        # Check that the stats were reduced (should have aggregated values)
-        assert "CPU_to_GPU_total_bytes" in reduced
-        assert "CPU_to_GPU_total_time" in reduced
-        assert "GPU_to_CPU_total_bytes" in reduced
-        assert "GPU_to_CPU_total_time" in reduced
-        assert reduced["CPU_to_GPU_total_bytes"] == 34
-        assert reduced["CPU_to_GPU_total_time"] == 2.6
-        assert reduced["GPU_to_CPU_total_time"] == 2.3
-        assert reduced["GPU_to_CPU_total_bytes"] == 19
-
-    def test_reset(self):
-        """Test that reset() resets all nested connector stats."""
-        offload_connector_stats = OffloadingConnectorStats(
-            data={
-                "CPU_to_GPU": [
-                    {"op_size": 3, "op_time": 0.2},
-                    {"op_size": 7, "op_time": 0.9},
-                ],
-                "GPU_to_CPU": [{"op_size": 16, "op_time": 2}],
-            }
-        )
-
-        assert not offload_connector_stats.is_empty()
-
-        offload_connector_stats.reset()
-
-        # After reset, stats should be empty
-        assert offload_connector_stats.is_empty()
-        assert len(offload_connector_stats.data) == 0

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -0,0 +1,756 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import contextlib
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+import torch
+import zmq.asyncio
+
+from vllm.config import set_current_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
+    KVConnectorRole,
+    MooncakeConnector,
+    MooncakeConnectorMetadata,
+    MooncakeXferMetadata,
+    MooncakeXferResponse,
+    MooncakeXferResponseStatus,
+    PullReqMeta,
+    SendBlockMeta,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (
+    MooncakeBootstrapServer,
+)
+from vllm.utils.network_utils import get_open_port
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+from vllm.v1.request import RequestStatus
+
+from .utils import create_request, create_scheduler, create_vllm_config
+
+
+class FakeMooncakeWrapper:
+    """Mock Mooncake TransferEngine for unit testing environments."""
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def initialize(self, local_hostname, metadata_server, protocol, device_name) -> int:
+        return 0
+
+    def get_rpc_port(self) -> int:
+        return 12345
+
+    def batch_transfer_sync_write(
+        self, target_hostname, buffers, peer_buffer_addresses, lengths
+    ) -> int:
+        return 0
+
+    def batch_register_memory(self, buffer_addresses, capacities) -> int:
+        return 0
+
+
+def test_basic_interface():
+    """Unit test for basic MooncakeConnector interface functionality."""
+
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_consumer"
+    )
+    scheduler = create_scheduler(vllm_config)
+
+    # 2 Full Blocks and 1 Half Block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_EXTERNAL_FULL_BLOCKS = 2
+    NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))
+
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_prefill=True,
+    )
+    request_id = request.request_id
+    request.kv_transfer_params.update(
+        {
+            "transfer_id": request_id,
+            "remote_bootstrap_addr": 54321,
+        }
+    )
+
+    scheduler.add_request(request)
+
+    # Remote Prefill, triggers NixlConnectorMetadata.
+    scheduler_output = scheduler.schedule()
+    kv_connector_metadata = scheduler_output.kv_connector_metadata
+    assert kv_connector_metadata is not None
+    assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
+
+    assert len(kv_connector_metadata.reqs_to_recv) == 1
+    assert request_id in kv_connector_metadata.reqs_to_recv["my-engine-id"]
+    req_meta = kv_connector_metadata.reqs_to_recv["my-engine-id"][request_id]
+
+    for block_id, block in zip(
+        req_meta.local_block_ids,
+        scheduler.kv_cache_manager.coordinator.single_type_managers[0].req_to_blocks[
+            request_id
+        ],
+    ):
+        assert block_id == block.block_id
+
+
+def test_prompt_less_than_block_size():
+    """Test that we can handle case where prompt is < block."""
+
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_consumer"
+    )
+    scheduler = create_scheduler(vllm_config)
+
+    # Half of a block.
+    BLOCK_SIZE = vllm_config.cache_config.block_size
+    NUM_TOKENS = int(BLOCK_SIZE * 0.5)
+
+    # Request will have 1 partial remote block.
+    request = create_request(
+        request_id=1,
+        block_size=BLOCK_SIZE,
+        num_tokens=NUM_TOKENS,
+        do_remote_prefill=True,
+        num_remote_blocks=1,
+    )
+    request.kv_transfer_params.update(
+        {
+            "transfer_id": request.request_id,
+            "remote_bootstrap_addr": 54321,
+        }
+    )
+
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+
+    # This request will read async.
+    kv_connector_metadata = scheduler_output.kv_connector_metadata
+    assert kv_connector_metadata is not None
+    assert isinstance(kv_connector_metadata, MooncakeConnectorMetadata)
+    assert len(kv_connector_metadata.reqs_to_recv["my-engine-id"]) == 1
+    assert len(scheduler_output.scheduled_new_reqs) == 0
+
+
+@pytest.fixture
+def bootstrap_server():
+    """Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server."""
+
+    port = get_open_port()
+    server = MooncakeBootstrapServer("127.0.0.1", port)
+    server.start()
+    yield server
+    server.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):
+    """
+    Tests the bootstrap server's api for worker registration and querying.
+
+    Validates DP/TP/PP rank indexing and error handling for duplicate registrations.
+    """
+
+    import httpx
+
+    base_url = f"http://127.0.0.1:{bootstrap_server.port}"
+
+    # Query when empty
+    async with httpx.AsyncClient() as client:
+        response = await client.get(f"{base_url}/query")
+        assert response.status_code == 200
+        assert response.json() == {}
+
+    # Register a worker
+    payload1 = {
+        "engine_id": "eng-1",
+        "dp_rank": 0,
+        "tp_rank": 0,
+        "pp_rank": 0,
+        "addr": "tcp://1.1.1.1:1111",
+    }
+    async with httpx.AsyncClient() as client:
+        response = await client.post(f"{base_url}/register", json=payload1)
+        assert response.status_code == 200
+        assert response.json() == {"status": "ok"}
+
+    # Query after registration
+    async with httpx.AsyncClient() as client:
+        response = await client.get(f"{base_url}/query")
+        assert response.status_code == 200
+        data = response.json()
+        assert "0" in data
+        assert data["0"]["engine_id"] == "eng-1"
+        assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111"
+
+    # Test failure: re-registering the same worker
+    async with httpx.AsyncClient() as client:
+        response = await client.post(f"{base_url}/register", json=payload1)
+        assert response.status_code == 400
+        assert "is already registered" in response.text
+
+    # Test failure: engine_id mismatch for same dp_rank
+    payload3_fail = {
+        "engine_id": "eng-2",
+        "dp_rank": 0,
+        "tp_rank": 1,
+        "pp_rank": 0,
+        "addr": "tcp://3.3.3.3:3333",
+    }
+    async with httpx.AsyncClient() as client:
+        response = await client.post(f"{base_url}/register", json=payload3_fail)
+        assert response.status_code == 400
+        assert "Engine ID mismatch" in response.text
+
+
+def test_scheduler_request_finished():
+    """
+    Tests the scheduler-side logic when a request finishes.
+
+    Differentiates between 'Finished' (requires transfer)
+    and 'Aborted' (immediate free).
+    """
+
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_producer"
+    )
+    scheduler = create_scheduler(vllm_config)
+    scheduler_connector = scheduler.get_kv_connector().connector_scheduler
+
+    request = create_request(request_id=1, do_remote_decode=True)
+    request.kv_transfer_params["transfer_id"] = request.request_id
+
+    # Case: Capped length (Successful prefill, need to send to decoder)
+    request.status = RequestStatus.FINISHED_LENGTH_CAPPED
+    delay_free, _ = scheduler_connector.request_finished(request, block_ids=[10, 11])
+    assert delay_free is True
+    assert "id-1" in scheduler_connector._reqs_need_send
+    assert scheduler_connector._reqs_need_send["id-1"][1] == [10, 11]
+
+    # Case: Aborted (No need to transfer, free blocks immediately)
+    scheduler_connector._reqs_need_send.clear()
+    request.status = RequestStatus.FINISHED_ABORTED
+    delay_free, _ = scheduler_connector.request_finished(request, block_ids=[12])
+    assert delay_free is False
+    assert len(scheduler_connector._reqs_need_send) == 0
+    assert "id-1" in scheduler_connector._reqs_not_processed
+
+
+@contextlib.contextmanager
+def patch_worker_dependencies():
+    """Helper to mock all distributed and network dependencies for Worker tests."""
+
+    with (
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
+            FakeMooncakeWrapper,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_ip",
+            return_value="127.0.0.1",
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_rank",
+            return_value=0,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_tensor_model_parallel_world_size",
+            return_value=1,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.get_pp_group"
+        ) as mock_pp,
+        patch("vllm.distributed.parallel_state.is_local_first_rank", return_value=True),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.should_launch_bootstrap_server",
+            return_value=False,
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.make_zmq_socket"
+        ) as mock_make_zmq,
+        patch("httpx.AsyncClient") as mock_async_client,
+    ):
+        # Mock PP group
+        mock_pp_group = MagicMock()
+        mock_pp_group.rank_in_group = 0
+        mock_pp.return_value = mock_pp_group
+
+        # Mock ZMQ socket
+        mock_socket_object = AsyncMock()
+        mock_socket_object.setsockopt = MagicMock()
+        mock_socket_ctx = MagicMock()
+        mock_socket_ctx.__enter__.return_value = mock_socket_object
+        mock_make_zmq.return_value = mock_socket_ctx
+
+        # Mock httpx client
+        mock_http_client_instance = AsyncMock()
+        mock_async_client.return_value = mock_http_client_instance
+
+        yield {
+            "mock_make_zmq": mock_make_zmq,
+            "mock_socket_object": mock_socket_object,
+            "mock_async_client": mock_async_client,
+            "mock_http_client": mock_http_client_instance,
+        }
+
+
+@pytest.mark.asyncio
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.TransferEngine",
+    FakeMooncakeWrapper,
+)
+async def test_kv_producer(monkeypatch):
+    """
+    Simulates a Producer Worker (Prefiller) receiving a transfer request
+    from a Consumer (Decoder).
+
+    Verifies memory offset calculation: ptr = base_addr + block_id * block_len.
+    """
+
+    monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_producer"
+    )
+
+    with set_current_vllm_config(vllm_config), patch_worker_dependencies():
+        prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
+        prefill_worker = prefill_connector.connector_worker
+        prefill_worker.kv_caches_base_addr = [0x1000]
+        block_len = 4096
+        prefill_worker.block_len_per_layer = [block_len]
+
+        # Override loop to use current test loop
+        origin_sender_loop = prefill_worker.sender_loop
+        prefill_worker.sender_loop = asyncio.get_event_loop()
+
+        # A request is finished on Producer and ready to be sent.
+        transfer_id = "xfer-req-1"
+        send_meta = SendBlockMeta(
+            p_req_id="p-req-1",
+            transfer_id=transfer_id,
+            local_block_ids=[10, 11],
+            ready=asyncio.Event(),
+        )
+        prefill_worker.reqs_need_send[transfer_id] = send_meta
+        send_meta.ready.set()
+
+        # Remote consumer request metadata
+        xfer_meta = MooncakeXferMetadata(
+            remote_hostname="consumer-host",
+            remote_port=54321,
+            remote_tp_size=1,
+            remote_tp_rank=0,
+            req_blocks={"d-req-1": (transfer_id, [20, 21])},
+            kv_caches_base_addr=[0x2000],
+            block_lens=[block_len],
+        )
+
+        mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
+        mock_socket.send_multipart = AsyncMock()
+        identity = b"consumer-id"
+
+        with patch.object(
+            prefill_worker, "_send_blocks", return_value=0
+        ) as mock_send_blocks:
+            # Normal case: 2 blocks to 2 blocks
+            # Worker processes the consumer's request
+            await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
+            # Verify transfer parameters are correct
+            src_ptr = 0x1000 + 10 * block_len
+            dst_ptr = 0x2000 + 20 * block_len
+            length = 2 * block_len
+            mock_send_blocks.assert_called_once_with(
+                "consumer-host:54321", [src_ptr], [dst_ptr], [length]
+            )
+            mock_socket.send_multipart.assert_called_once()
+
+            # Verify the response sent back to the consumer
+            sent_call = mock_socket.send_multipart.call_args[0][0]
+            sent_identity, sent_payload = sent_call
+            assert sent_identity == identity
+            response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
+            assert response.status == MooncakeXferResponseStatus.FINISH
+            assert response.ok_reqs == ["d-req-1"]
+
+            # Verify internal state cleanup
+            assert transfer_id not in prefill_worker.reqs_need_send
+            assert "p-req-1" in prefill_worker.finished_sending_reqs
+
+            # More cases:
+            # Consumer only needs 1 block (less than P)
+            mock_send_blocks.reset_mock()
+            mock_socket.send_multipart.reset_mock()
+            prefill_worker.reqs_need_send[transfer_id] = send_meta
+            send_meta.sent = 0
+            send_meta.ready.set()
+            xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20])
+            # Worker processes the consumer's request
+            await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
+            # Verify transfer parameters are correct: 11 to 20
+            src_ptr = 0x1000 + 11 * block_len
+            dst_ptr = 0x2000 + 20 * block_len
+            length = 1 * block_len
+            mock_send_blocks.assert_called_once_with(
+                "consumer-host:54321", [src_ptr], [dst_ptr], [length]
+            )
+            mock_socket.send_multipart.assert_called_once()
+
+            # Consumer needs 3 blocks (more than P, error case)
+            mock_send_blocks.reset_mock()
+            mock_socket.send_multipart.reset_mock()
+            prefill_worker.reqs_need_send[transfer_id] = send_meta
+            send_meta.sent = 0
+            send_meta.ready.set()
+            xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21, 22])
+            # Worker processes the consumer's request
+            await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
+            # This should not be called because error.
+            mock_send_blocks.assert_not_called()
+            mock_socket.send_multipart.assert_called_once()
+            _, sent_payload = mock_socket.send_multipart.call_args[0][0]
+            response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
+            assert response.err_msg == "P num blocks less than D"
+            assert response.err_reqs == ["d-req-1"]
+
+            # Timeout
+            mock_send_blocks.reset_mock()
+            mock_socket.send_multipart.reset_mock()
+            prefill_worker.reqs_need_send[transfer_id] = send_meta
+            send_meta.sent = 0
+            send_meta.ready.clear()
+            xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
+            # Worker processes the consumer's request
+            await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
+            # This should not be called because timeout.
+            mock_send_blocks.assert_not_called()
+            mock_socket.send_multipart.assert_called_once()
+            _, sent_payload = mock_socket.send_multipart.call_args[0][0]
+            response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
+            assert response.err_msg == "Timeout waiting for P side ready."
+            assert response.err_reqs == ["d-req-1"]
+
+        # Transfer error
+        with patch.object(
+            prefill_worker, "_send_blocks", return_value=123
+        ) as mock_send_blocks:
+            mock_socket.send_multipart.reset_mock()
+            prefill_worker.reqs_need_send[transfer_id] = send_meta
+            send_meta.sent = 0
+            send_meta.ready.set()
+            xfer_meta.req_blocks["d-req-1"] = (transfer_id, [20, 21])
+            # Worker processes the consumer's request
+            await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
+            mock_send_blocks.assert_called_once()
+            mock_socket.send_multipart.assert_called_once()
+            _, sent_payload = mock_socket.send_multipart.call_args[0][0]
+            response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
+            assert response.err_msg == "Mooncake transfer engine returned 123"
+            assert response.err_reqs == ["d-req-1"]
+
+        # Clean up
+        prefill_worker.sender_loop = origin_sender_loop
+        prefill_worker.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_kv_consumuer(monkeypatch):
+    """
+    Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.
+
+    Verifies that MooncakeXferMetadata is correctly serialized and sent via ZMQ.
+    """
+
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_consumer"
+    )
+
+    with set_current_vllm_config(vllm_config), patch_worker_dependencies() as mocks:
+        decode_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
+        decode_worker = decode_connector.connector_worker
+        decode_worker.kv_caches_base_addr = [0x1000]
+        decode_worker.rpc_port = 54321
+
+        # A request to pull data arrives.
+        pull_metas = {
+            "d-req-1": PullReqMeta(
+                d_req_id="d-req-1",
+                transfer_id="xfer-req-1",
+                local_block_ids=[100, 101],
+                remote_engine_id="p-engine",
+                remote_bootstrap_addr="http://bootstrap:33333",
+                pull_tasks_count=1,
+            )
+        }
+        decode_worker._remote_agents = {"p-engine": {0: {0: "tcp://producer:1234"}}}
+        decode_worker._tp_size["p-engine"] = 1
+
+        # Mock the response from the producer.
+        mock_response = MooncakeXferResponse(
+            status=MooncakeXferResponseStatus.FINISH, ok_reqs=["d-req-1"]
+        )
+        encoded_response = decode_worker._encoder.encode(mock_response)
+        mocks["mock_socket_object"].recv.return_value = encoded_response
+
+        # Trigger the receive logic.
+        decode_worker.receive_kv("p-engine", pull_metas)
+        await asyncio.sleep(1)  # Allow async task to run
+
+        # Verify the metadata sent to the producer.
+        mocks["mock_make_zmq"].assert_called_with(
+            decode_worker.async_zmq_ctx,
+            "tcp://producer:1234",
+            zmq.DEALER,
+            bind=False,
+            linger=0,
+        )
+        sent_payload = mocks["mock_socket_object"].send.call_args[0][0]
+        sent_meta = decode_worker._xfer_meta_decoder.decode(sent_payload)
+
+        assert sent_meta.remote_hostname == "127.0.0.1"
+        assert sent_meta.remote_port == 54321
+        assert sent_meta.req_blocks["d-req-1"] == ("xfer-req-1", [100, 101])
+
+        # Verify internal state is updated correctly.
+        assert "d-req-1" in decode_worker.finished_recving_reqs
+
+        # Clean up
+        decode_worker.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_worker_get_finished_timeout(monkeypatch):
+    """Tests the cleanup mechanism for requests."""
+
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_producer"
+    )
+    with set_current_vllm_config(vllm_config), patch_worker_dependencies():
+        prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
+        prefill_worker = prefill_connector.connector_worker
+
+        # Add an expired request (expire_time is in the past).
+        prefill_worker.reqs_need_send["tx-expired"] = SendBlockMeta(
+            p_req_id="p-req-expired",
+            transfer_id="tx-expired",
+            local_block_ids=[1, 2],
+            ready=MagicMock(),
+            expire_time=time.perf_counter() - 100,
+        )
+
+        # Add a non-expired request.
+        prefill_worker.reqs_need_send["tx-active"] = SendBlockMeta(
+            p_req_id="p-req-active",
+            transfer_id="tx-active",
+            local_block_ids=[3, 4],
+            ready=MagicMock(),
+            expire_time=time.perf_counter() + 100,
+        )
+
+        finished_reqs = await prefill_worker.fetch_finished_sending_reqs()
+
+        assert "p-req-expired" in finished_reqs
+        assert "p-req-active" not in finished_reqs
+        assert "tx-expired" not in prefill_worker.reqs_need_send
+        assert "tx-active" in prefill_worker.reqs_need_send
+
+
+def test_register_kv_caches():
+    """Tests the memory registration logic with the underlying Mooncake engine."""
+
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_consumer"
+    )
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch_worker_dependencies(),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
+        ),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
+        ) as mock_thread,
+    ):
+        connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
+        worker = connector.connector_worker
+        mock_thread.return_value.is_alive.return_value = False
+
+        kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
+            num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+        )
+        tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+        tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+        kv_caches = {"layer0": tensor1, "layer1": tensor2}
+
+        with patch.object(
+            worker.engine, "batch_register_memory", return_value=0
+        ) as mock_batch_register:
+            connector.register_kv_caches(kv_caches)
+
+            mock_batch_register.assert_called_once()
+            registered_ptrs, registered_lens = mock_batch_register.call_args[0]
+            expected_ptrs = {
+                tensor.data_ptr()
+                for kv_pair in kv_caches.values()
+                for tensor in kv_pair
+            }
+            assert set(registered_ptrs) == expected_ptrs
+            assert set(registered_lens) == {tensor1[0].nbytes}
+
+            # Verify block_len_per_layer is set correctly.
+            assert len(worker.block_len_per_layer) == len(registered_ptrs)
+            for bl in worker.block_len_per_layer:
+                assert bl == tensor1[0].nbytes // tensor1.shape[1]
+
+
+@pytest.mark.asyncio
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
+    "mooncake_connector.TransferEngine",
+    FakeMooncakeWrapper,
+)
+@pytest.mark.parametrize("d_tp_size", [1, 4], ids=["p_tp2_d_tp1", "p_tp2_d_tp4"])
+async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
+    """
+    Tests heterogeneous TP support in the producer transfer path.
+
+    Verifies correct pointer and offset calculation when producer TP=2
+    sends to consumer with TP=1 (P>D) or TP=4 (P<D).
+
+    Parametrized cases:
+    - P TP=2 > D TP=1: one D rank receives; dst_offset based on P rank
+    - P TP=2 < D TP=4: two D ranks receive; src_offset based on D rank
+    """
+
+    P_TP_SIZE = 2
+    P_TP_RANK = 0
+    LOCAL_BLOCK_LEN = 4096
+
+    local_block_len = LOCAL_BLOCK_LEN
+    remote_block_len = LOCAL_BLOCK_LEN * P_TP_SIZE // d_tp_size
+
+    monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
+    vllm_config = create_vllm_config(
+        kv_connector="MooncakeConnector", kv_role="kv_producer"
+    )
+
+    with set_current_vllm_config(vllm_config), patch_worker_dependencies():
+        prefill_connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
+        prefill_worker = prefill_connector.connector_worker
+
+        # Override TP rank/size to simulate P TP=2
+        prefill_worker.tp_rank = P_TP_RANK
+        prefill_worker.tp_size = P_TP_SIZE
+        # Update shared dict so kv_topo sees correct TP size
+        prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE
+        prefill_worker.kv_topo.tp_rank = P_TP_RANK
+
+        prefill_worker.kv_caches_base_addr = [0x1000]
+        prefill_worker.block_len_per_layer = [local_block_len]
+
+        origin_sender_loop = prefill_worker.sender_loop
+        prefill_worker.sender_loop = asyncio.get_event_loop()
+
+        transfer_id = "xfer-hetero-1"
+        local_block_ids = [10, 11]
+        send_meta = SendBlockMeta(
+            p_req_id="p-req-h1",
+            transfer_id=transfer_id,
+            local_block_ids=local_block_ids,
+            ready=asyncio.Event(),
+        )
+        prefill_worker.reqs_need_send[transfer_id] = send_meta
+        send_meta.ready.set()
+
+        # Compute target D ranks using the production code path
+        target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size)
+
+        mock_socket = AsyncMock(spec=zmq.asyncio.Socket)
+        mock_socket.send_multipart = AsyncMock()
+        identity = b"consumer-hetero"
+
+        # Assign different remote block IDs per D rank
+        d_rank_remote_blocks = {
+            rank: [20 + i * 10, 21 + i * 10] for i, rank in enumerate(target_d_ranks)
+        }
+
+        with patch.object(
+            prefill_worker, "_send_blocks", return_value=0
+        ) as mock_send_blocks:
+            for d_rank in target_d_ranks:
+                remote_block_ids = d_rank_remote_blocks[d_rank]
+                xfer_meta = MooncakeXferMetadata(
+                    remote_hostname="consumer-host",
+                    remote_port=54321,
+                    remote_tp_size=d_tp_size,
+                    remote_tp_rank=d_rank,
+                    req_blocks={
+                        f"d-req-h1-r{d_rank}": (
+                            transfer_id,
+                            remote_block_ids,
+                        )
+                    },
+                    kv_caches_base_addr=[0x2000],
+                    block_lens=[remote_block_len],
+                )
+
+                mock_send_blocks.reset_mock()
+                mock_socket.send_multipart.reset_mock()
+
+                await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
+
+                # Verify _send_blocks was called
+                mock_send_blocks.assert_called_once()
+                call_args = mock_send_blocks.call_args[0]
+                src_ptrs = call_args[1]
+                dst_ptrs = call_args[2]
+                lengths = call_args[3]
+
+                # Heterogeneous TP: blocks cannot be coalesced because
+                # local and remote block_lens differ
+                assert len(src_ptrs) == len(local_block_ids)
+                assert len(dst_ptrs) == len(local_block_ids)
+                assert len(lengths) == len(local_block_ids)
+
+                # Compute expected offsets based on TP ratio
+                if d_tp_size <= P_TP_SIZE:
+                    tp_ratio = P_TP_SIZE // d_tp_size
+                    expected_src_off = 0
+                    expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
+                    expected_xfer_len = local_block_len
+                else:
+                    ratio_abs = d_tp_size // P_TP_SIZE
+                    expected_src_off = (d_rank % ratio_abs) * remote_block_len
+                    expected_dst_off = 0
+                    expected_xfer_len = remote_block_len
+
+                for idx, (lblk, rblk) in enumerate(
+                    zip(local_block_ids, remote_block_ids)
+                ):
+                    assert src_ptrs[idx] == (
+                        0x1000 + lblk * local_block_len + expected_src_off
+                    )
+                    assert dst_ptrs[idx] == (
+                        0x2000 + rblk * remote_block_len + expected_dst_off
+                    )
+                    assert lengths[idx] == expected_xfer_len
+
+                # Verify successful response sent back to consumer
+                mock_socket.send_multipart.assert_called_once()
+                _, sent_payload = mock_socket.send_multipart.call_args[0][0]
+                response = prefill_worker._xfer_resp_decoder.decode(sent_payload)
+                assert response.status == MooncakeXferResponseStatus.FINISH
+                assert response.ok_reqs == [f"d-req-h1-r{d_rank}"]
+
+        # After serving all D ranks, the request should be complete
+        assert transfer_id not in prefill_worker.reqs_need_send
+        assert "p-req-h1" in prefill_worker.finished_sending_reqs
+
+        prefill_worker.sender_loop = origin_sender_loop
+        prefill_worker.shutdown()

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -91,6 +91,9 @@ def clear_kv_transfer():
     yield
     if has_kv_transfer_group():
         ensure_kv_transfer_shutdown()
+    # Reset any KV cache layout override set during tests so it doesn't
+    # leak into tests in other modules.
+    set_kv_cache_layout(None)
 
 
 def get_default_xfer_telemetry(

tests/v1/kv_connector/unit/utils.py

@@ -100,6 +100,8 @@ def create_vllm_config(
     hf_overrides: dict[str, Any] | None = None,
     attention_backend: str | None = None,
     kv_load_failure_policy: Literal["recompute", "fail"] = "fail",
+    kv_connector: str = "NixlConnector",
+    kv_role: str = "kv_both",
 ) -> VllmConfig:
     """Initialize VllmConfig For Testing."""
     model_config = ModelConfig(
@@ -124,8 +126,8 @@ def create_vllm_config(
         enable_prefix_caching=True,
     )
     kv_transfer_config = KVTransferConfig(
-        kv_connector="NixlConnector",
-        kv_role="kv_both",
+        kv_connector=kv_connector,
+        kv_role=kv_role,
         enable_permute_local_kv=enable_permute_local_kv,
         kv_connector_extra_config=kv_connector_extra_config or {},
         kv_load_failure_policy=kv_load_failure_policy,

tests/v1/kv_offload/test_cpu_gpu.py

@@ -6,32 +6,20 @@ import time
 import pytest
 import torch
 
-from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
+from vllm.v1.kv_offload.spec import (
+    CanonicalKVCacheRef,
+    CanonicalKVCaches,
+    CanonicalKVCacheTensor,
+)
 from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandlers
 
-BACKENDS_TO_TEST = [FlashAttentionBackend]
-
-if not current_platform.is_rocm():
-    from vllm.v1.attention.backends.flashinfer import FlashInferBackend
-
-    BACKENDS_TO_TEST.append(FlashInferBackend)
-
-    from vllm.v1.attention.backends.mla.flashattn_mla import FlashAttnMLABackend
-
-    BACKENDS_TO_TEST.append(FlashAttnMLABackend)
-
 NUM_GPU_BLOCKS = [64]
 NUM_CPU_BLOCKS = [256]
-KERNEL_BLOCK_SIZES = [16]
-LOGICAL_BLOCK_SIZES = [16, 32]
-LOGICAL_BLOCKS_PER_CPU_BLOCK = [1, 3]
-HEAD_SIZES = [64]
-NUM_HEADS = [8]
-NUM_LAYERS = [4]
-DTYPES = [torch.bfloat16]
+GPU_PAGE_SIZES = [512, 1024]
+BLOCK_SIZE_FACTORS = [1, 3]
+NUM_TENSORS = [4]
 SEEDS = [0]
 CUDA_DEVICES = ["cuda:0"]
 NUM_MAPPINGS = [3]
@@ -39,15 +27,11 @@ NUM_MAPPINGS = [3]
 
 @pytest.mark.parametrize("gpu_to_cpu", [True, False])
 @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS)
-@pytest.mark.parametrize("head_size", HEAD_SIZES)
-@pytest.mark.parametrize("num_heads", NUM_HEADS)
-@pytest.mark.parametrize("kernel_block_size", KERNEL_BLOCK_SIZES)
-@pytest.mark.parametrize("logical_block_size", LOGICAL_BLOCK_SIZES)
-@pytest.mark.parametrize("logical_blocks_per_cpu_block", LOGICAL_BLOCKS_PER_CPU_BLOCK)
+@pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES)
+@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS)
 @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS)
 @pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS)
-@pytest.mark.parametrize("num_layers", NUM_LAYERS)
-@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("num_tensors", NUM_TENSORS)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @torch.inference_mode()
@@ -55,113 +39,89 @@ def test_transfer(
     default_vllm_config,
     gpu_to_cpu: bool,
     num_mappings: int,
-    head_size: int,
-    num_heads: int,
-    kernel_block_size: int,
-    logical_block_size: int,
-    logical_blocks_per_cpu_block: int,
+    gpu_page_size_bytes: int,
+    block_size_factor: int,
     num_gpu_blocks: int,
     num_cpu_blocks: int,
-    num_layers: int,
-    dtype: torch.dtype,
+    num_tensors: int,
     seed: int,
     device: str,
 ) -> None:
     set_random_seed(seed)
 
-    # create per-layer GPU KV caches based on available attn_backends
-    attn_backends_list = BACKENDS_TO_TEST
-
-    assert logical_block_size % kernel_block_size == 0
-    kernel_blocks_per_gpu_block = logical_block_size // kernel_block_size
-    num_gpu_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
-
-    gpu_caches = {}
-    attn_backends = {}
-    for i in range(num_layers):
-        layer_name = f"layer {i}"
-
-        attn_backend = attn_backends_list[i % len(attn_backends_list)]
-        attn_backends[layer_name] = attn_backend
-
-        gpu_cache_shape = attn_backend.get_kv_cache_shape(
-            num_gpu_kernel_blocks, kernel_block_size, num_heads, head_size
+    # build CanonicalKVCacheTensor list: one per tensor
+    kv_cache_tensors: list[CanonicalKVCacheTensor] = []
+    for i in range(num_tensors):
+        gpu_tensor = torch.randint(
+            -128,
+            127,
+            (num_gpu_blocks, gpu_page_size_bytes),
+            dtype=torch.int8,
+            device=device,
+        )
+        kv_cache_tensors.append(
+            CanonicalKVCacheTensor(
+                tensor=gpu_tensor,
+                page_size_bytes=gpu_page_size_bytes,
+            )
         )
-        gpu_caches[layer_name] = torch.rand(gpu_cache_shape, dtype=dtype, device=device)
 
-    # create handler
-    cpu_block_size = logical_blocks_per_cpu_block * logical_block_size
-    kernel_blocks_per_cpu_block = cpu_block_size // kernel_block_size
+    # one group containing all tensors, one data ref per tensor
+    kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]] = [
+        [
+            CanonicalKVCacheRef(
+                tensor_idx=i,
+                page_size_bytes=gpu_page_size_bytes,
+            )
+            for i in range(num_tensors)
+        ]
+    ]
+
+    kv_caches = CanonicalKVCaches(
+        tensors=kv_cache_tensors,
+        group_data_refs=kv_cache_groups_data_refs,
+    )
     handlers = CpuGpuOffloadingHandlers(
-        attn_backends=attn_backends,
-        gpu_block_size=logical_block_size,
-        cpu_block_size=cpu_block_size,
+        kv_caches=kv_caches,
+        block_size_factor=block_size_factor,
         num_cpu_blocks=num_cpu_blocks,
-        gpu_caches=gpu_caches,
     )
 
     # select block mappings
-    gpu_blocks = random.sample(
-        range(num_gpu_blocks), num_mappings * logical_blocks_per_cpu_block
-    )
+    gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * block_size_factor)
     cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings)
 
-    # convert gpu blocks to kernel block size
-    gpu_blocks_in_kernel_block_size = []
-    for gpu_block in gpu_blocks:
-        base_block_id = gpu_block * kernel_blocks_per_gpu_block
-        for i in range(kernel_blocks_per_gpu_block):
-            gpu_blocks_in_kernel_block_size.append(i + base_block_id)
+    # expand cpu blocks to gpu-page granularity for uniform comparison:
+    # each cpu block maps to block_size_factor consecutive sub-blocks
+    cpu_blocks_expanded = [
+        cpu_block * block_size_factor + j
+        for cpu_block in cpu_blocks
+        for j in range(block_size_factor)
+    ]
 
-    # convert cpu blocks to gpu block size
-    cpu_blocks_in_kernel_block_size = []
-    for cpu_block in cpu_blocks:
-        base_block_id = cpu_block * kernel_blocks_per_cpu_block
-        for i in range(kernel_blocks_per_cpu_block):
-            cpu_blocks_in_kernel_block_size.append(i + base_block_id)
-
-    # maybe skip some GPU block to test reading from the middle of a CPU block
+    # maybe skip some GPU blocks to test reading from the middle of a CPU block
     if not gpu_to_cpu:
-        gpu_blocks_to_skip = logical_blocks_per_cpu_block - 1
-        gpu_blocks = gpu_blocks[gpu_blocks_to_skip:]
-        kernel_blocks_to_skip = gpu_blocks_to_skip * kernel_blocks_per_gpu_block
-        gpu_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size[
-            kernel_blocks_to_skip:
-        ]
-        cpu_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size[
-            kernel_blocks_to_skip:
-        ]
+        blocks_to_skip = block_size_factor - 1
+        gpu_blocks = gpu_blocks[blocks_to_skip:]
+        cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:]
 
     # set transfer direction
     if gpu_to_cpu:
         handler = handlers.gpu_to_cpu_handler
-        src_blocks = gpu_blocks
-        dst_blocks = cpu_blocks
-        src_spec = GPULoadStoreSpec(src_blocks, group_sizes=(len(src_blocks),))
-        dst_spec = CPULoadStoreSpec(dst_blocks)
-        src_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
-        dst_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
-        dst_size_in_kernel_blocks = num_cpu_blocks * kernel_blocks_per_cpu_block
+        src_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))
+        dst_spec = CPULoadStoreSpec(cpu_blocks)
+        dst_to_src = dict(zip(cpu_blocks_expanded, gpu_blocks))
+        num_dst_sub_blocks = num_cpu_blocks * block_size_factor
     else:
         handler = handlers.cpu_to_gpu_handler
-        src_blocks = cpu_blocks
-        dst_blocks = gpu_blocks
-        src_spec = CPULoadStoreSpec(src_blocks)
-        dst_spec = GPULoadStoreSpec(dst_blocks, group_sizes=(len(dst_blocks),))
-        src_blocks_in_kernel_block_size = cpu_blocks_in_kernel_block_size
-        dst_blocks_in_kernel_block_size = gpu_blocks_in_kernel_block_size
-        dst_size_in_kernel_blocks = num_gpu_blocks * kernel_blocks_per_gpu_block
-
-    # build dst -> src mapping
-    dst_to_src = {}
-    for src_block, dst_block in zip(
-        src_blocks_in_kernel_block_size, dst_blocks_in_kernel_block_size
-    ):
-        dst_to_src[dst_block] = src_block
+        src_spec = CPULoadStoreSpec(cpu_blocks)
+        dst_spec = GPULoadStoreSpec(gpu_blocks, group_sizes=(len(gpu_blocks),))
+        dst_to_src = dict(zip(gpu_blocks, cpu_blocks_expanded))
+        num_dst_sub_blocks = num_gpu_blocks
 
     # clone src and dst tensors before transfer
-    orig_src_caches = [x.clone() for x in handler.src_tensors]
-    orig_dst_caches = [x.clone() for x in handler.dst_tensors]
+    orig_src_tensors = [x.clone() for x in handler.src_tensors]
+    orig_dst_tensors = [x.clone() for x in handler.dst_tensors]
 
     # call transfer function
     start_time = time.time()
@@ -180,11 +140,8 @@ def test_transfer(
                 if gpu_to_cpu
                 else ("CPU", "GPU")
             )
-            assert (
-                finished[0].transfer_size
-                == handler.total_block_size_in_bytes
-                * handler.dst_block_size_factor
-                * len(dst_blocks)
+            assert finished[0].transfer_size == (
+                len(gpu_blocks) * handler.group_block_size_in_bytes[0]
             )
             assert finished[0].transfer_time > 0
             assert finished[0].transfer_time < (time.time() - start_time)
@@ -192,19 +149,23 @@ def test_transfer(
         time.sleep(0.1)
 
     # verify src tensors did not change
-    for orig_tensor, tensor in zip(orig_src_caches, handler.src_tensors):
+    for orig_tensor, tensor in zip(orig_src_tensors, handler.src_tensors):
         assert torch.equal(orig_tensor, tensor)
 
-    # verify dst tensors
-    for dst_block in range(dst_size_in_kernel_blocks):
-        src_block_candidate = dst_to_src.get(dst_block)
-        for src_cache, dst_cache, orig_dst_cache in zip(
-            handler.src_tensors,
-            handler.dst_tensors,
-            orig_dst_caches,
-        ):
-            if src_block_candidate is not None:
-                expected_value = src_cache[src_block_candidate]
+    # verify dst tensors at gpu-page granularity.
+    for src_tensor, dst_tensor, orig_dst_tensor in zip(
+        handler.src_tensors,
+        handler.dst_tensors,
+        orig_dst_tensors,
+    ):
+        # view both GPU and CPU tensors as (n, gpu_page_size_bytes) for comparison.
+        src_view = src_tensor.view(-1, gpu_page_size_bytes)
+        dst_view = dst_tensor.view(-1, gpu_page_size_bytes)
+        orig_dst_view = orig_dst_tensor.view(-1, gpu_page_size_bytes)
+        for dst_sub_block in range(num_dst_sub_blocks):
+            src_sub_block = dst_to_src.get(dst_sub_block)
+            if src_sub_block is not None:
+                expected = src_view[src_sub_block]
             else:
-                expected_value = orig_dst_cache[dst_block]
-            torch.testing.assert_close(dst_cache[dst_block].cpu(), expected_value.cpu())
+                expected = orig_dst_view[dst_sub_block]
+            torch.testing.assert_close(dst_view[dst_sub_block].cpu(), expected.cpu())

tests/v1/worker/test_gpu_model_runner.py

@@ -300,7 +300,7 @@ def test_update_states_request_resumed(model_runner, dist_init):
         req_ids=[req_id],
         resumed_req_ids=set(),
         new_token_ids=[[]],
-        all_token_ids={},
+        output_token_ids={},
         new_block_ids=[([0],)],
         num_computed_tokens=[0],
         num_output_tokens=[0],

tools/vllm-rocm/generate-rocm-wheels-root-index.sh

@@ -17,14 +17,14 @@
 #
 # Environment variables:
 #   S3_BUCKET   - Bucket name (default: vllm-wheels)
-#   VARIANT     - ROCm variant (default: rocm700)
+#   VARIANT     - ROCm variant (default: rocm721)
 #   DRY_RUN     - Set to 1 for preview mode (same as --dry-run)
 
 set -euo pipefail
 
 # ======== Configuration ========
 BUCKET="${S3_BUCKET:-vllm-wheels}"
-VARIANT="${VARIANT:-rocm700}"
+VARIANT="${VARIANT:-rocm721}"
 DRY_RUN="${DRY_RUN:-0}"
 FORCE_VERSION=""
 

vllm/config/cache.py

@@ -66,27 +66,29 @@ class CacheConfig:
     enable_prefix_caching: bool = True
     """Whether to enable prefix caching."""
     prefix_caching_hash_algo: PrefixCachingHashAlgo = "sha256"
-    """Set the hash algorithm for prefix caching:\n
-    - "sha256" uses Pickle for object serialization before hashing. This is the
-    current default, as SHA256 is the most secure choice to avoid potential
-    hash collisions.\n
+    """Set the hash algorithm for prefix caching:
+
+    - "sha256" uses Pickle for object serialization before hashing. This is the current
+      default, as SHA256 is the most secure choice to avoid potential hash collisions.
     - "sha256_cbor" provides a reproducible, cross-language compatible hash. It
-    serializes objects using canonical CBOR and hashes them with SHA-256.\n
+      serializes objects using canonical CBOR and hashes them with SHA-256.
     - "xxhash" uses Pickle serialization with xxHash (128-bit) for faster,
-    non-cryptographic hashing. Requires the optional ``xxhash`` package.
-    IMPORTANT: Use of a hashing algorithm that is not considered 
-    cryptographically secure theoretically increases the risk of hash collisions,
-    which can cause undefined behavior or even leak private information in
-    multi-tenant environments. Even if collisions are still very unlikely, it is
-    important to consider your security risk tolerance against the performance
-    benefits before turning this on.\n
+      non-cryptographic hashing. Requires the optional ``xxhash`` package.
+      IMPORTANT: Use of a hashing algorithm that is not considered  cryptographically
+      secure theoretically increases the risk of hash collisions, which can cause
+      undefined behavior or even leak private information in multi-tenant environments.
+      Even if collisions are still very unlikely, it is important to consider your
+      security risk tolerance against the performance benefits before turning this on.
     - "xxhash_cbor" combines canonical CBOR serialization with xxHash for
-    reproducible hashing. Requires the optional ``xxhash`` package."""
+      reproducible hashing. Requires the optional ``xxhash`` package."""
     calculate_kv_scales: bool = False
     """Deprecated: This option is deprecated and will be removed in v0.19.
     It enables dynamic calculation of `k_scale` and `v_scale` when
     kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model
     checkpoint if available. Otherwise, the scales will default to 1.0."""
+    kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
+    """Layer patterns to skip KV cache quantization. Accepts layer indices
+    (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
     cpu_kvcache_space_bytes: int | None = None
     """(CPU backend only) CPU key-value cache space."""
     mamba_page_size_padded: int | None = None

vllm/config/kernel.py

@@ -32,14 +32,14 @@ class KernelConfig:
     moe_backend: MoEBackend = "auto"
     """Backend for MoE expert computation kernels. Available options:
 
-    - "auto": Automatically select the best backend based on model and hardware\n
-    - "triton": Use Triton-based fused MoE kernels\n
-    - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)\n
-    - "cutlass": Use vLLM CUTLASS kernels\n
-    - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels\n
-    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels\n
-    - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)\n
-    - "marlin": Use Marlin kernels (weight-only quantization)\n
+    - "auto": Automatically select the best backend based on model and hardware
+    - "triton": Use Triton-based fused MoE kernels
+    - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
+    - "cutlass": Use vLLM CUTLASS kernels
+    - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
+    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
+    - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
+    - "marlin": Use Marlin kernels (weight-only quantization)
     - "aiter": Use AMD AITer kernels (ROCm only)"""
 
     @field_validator("moe_backend", mode="before")

vllm/config/load.py

@@ -51,7 +51,7 @@ class LoadConfig:
     - "gguf" will load weights from GGUF format files (details specified in
       https://github.com/ggml-org/ggml/blob/master/docs/gguf.md).
     - "mistral" will load weights from consolidated safetensors files used by
-      Mistral models.\n
+      Mistral models.
     - Other custom values can be supported via plugins.
     """
     download_dir: str | None = None

vllm/config/model.py

@@ -125,26 +125,28 @@ class ModelConfig:
     """Name or path of the Hugging Face tokenizer to use. If unspecified, model
     name or path will be used."""
     tokenizer_mode: TokenizerMode | str = "auto"
-    """Tokenizer mode:\n
+    """Tokenizer mode:
+
     - "auto" will use the tokenizer from `mistral_common` for Mistral models
-    if available, otherwise it will use the "hf" tokenizer.\n
-    - "hf" will use the fast tokenizer if available.\n
-    - "slow" will always use the slow tokenizer.\n
-    - "mistral" will always use the tokenizer from `mistral_common`.\n
-    - "deepseek_v32" will always use the tokenizer from `deepseek_v32`.\n
-    - "qwen_vl" will always use the tokenizer from `qwen_vl`.\n
+      if available, otherwise it will use the "hf" tokenizer.
+    - "hf" will use the fast tokenizer if available.
+    - "slow" will always use the slow tokenizer.
+    - "mistral" will always use the tokenizer from `mistral_common`.
+    - "deepseek_v32" will always use the tokenizer from `deepseek_v32`.
+    - "qwen_vl" will always use the tokenizer from `qwen_vl`.
     - Other custom values can be supported via plugins."""
     trust_remote_code: bool = False
     """Trust remote code (e.g., from HuggingFace) when downloading the model
     and tokenizer."""
     dtype: ModelDType | torch.dtype = "auto"
-    """Data type for model weights and activations:\n
+    """Data type for model weights and activations:
+
     - "auto" will use FP16 precision for FP32 and FP16 models, and BF16
-    precision for BF16 models.\n
-    - "half" for FP16. Recommended for AWQ quantization.\n
-    - "float16" is the same as "half".\n
-    - "bfloat16" for a balance between precision and range.\n
-    - "float" is shorthand for FP32 precision.\n
+      precision for BF16 models.
+    - "half" for FP16. Recommended for AWQ quantization.
+    - "float16" is the same as "half".
+    - "bfloat16" for a balance between precision and range.
+    - "float" is shorthand for FP32 precision.
     - "float32" for FP32 precision."""
     seed: int = 0
     """Random seed for reproducibility.
@@ -182,13 +184,14 @@ class ModelConfig:
     automatically derived from the model config.
 
     When passing via `--max-model-len`, supports k/m/g/K/M/G in human-readable
-    format. Examples:\n
-    - 1k -> 1000\n
-    - 1K -> 1024\n
-    - 25.6k -> 25,600\n
+    format. Examples:
+
+    - 1k -> 1000
+    - 1K -> 1024
+    - 25.6k -> 25,600
     - -1 or 'auto' -> Automatically choose the maximum model length that fits in
-    GPU memory. This will use the model's maximum context length if it fits,
-    otherwise it will find the largest length that can be accommodated."""
+      GPU memory. This will use the model's maximum context length if it fits,
+      otherwise it will find the largest length that can be accommodated."""
     spec_target_max_model_len: int | None = None
     """Specify the maximum length for spec decoding draft models."""
     quantization: QuantizationMethods | str | None = None
@@ -248,10 +251,11 @@ class ModelConfig:
     prometheus metrics, if multiple names provided, metrics tag will take the
     first one."""
     config_format: str | ConfigFormat = "auto"
-    """The format of the model config to load:\n
+    """The format of the model config to load:
+
     - "auto" will try to load the config in hf format if available after trying
-    to load in mistral format.\n
-    - "hf" will load the config in hf format.\n
+      to load in mistral format.
+    - "hf" will load the config in hf format.
     - "mistral" will load the config in mistral format."""
     hf_token: bool | str | None = None
     """The token to use as HTTP bearer authorization for remote files . If
@@ -276,12 +280,12 @@ class ModelConfig:
     """Enable sleep mode for the engine (only cuda and
     hip platforms are supported)."""
     model_impl: str | ModelImpl = "auto"
-    """Which implementation of the model to use:\n
-    - "auto" will try to use the vLLM implementation, if it exists, and fall
-    back to the Transformers implementation if no vLLM implementation is
-    available.\n
-    - "vllm" will use the vLLM model implementation.\n
-    - "transformers" will use the Transformers model implementation.\n
+    """Which implementation of the model to use:
+
+    - "auto" will try to use the vLLM implementation, if it exists, and fall back to the
+      Transformers implementation if no vLLM implementation is available.
+    - "vllm" will use the vLLM model implementation.
+    - "transformers" will use the Transformers model implementation.
     - "terratorch" will use the TerraTorch model implementation.
     """
     override_attention_dtype: str | None = None
@@ -291,6 +295,10 @@ class ModelConfig:
     definitions"""
     io_processor_plugin: str | None = None
     """IOProcessor plugin name to load at model startup"""
+    renderer_num_workers: int = 1
+    """Number of worker threads in the renderer thread pool. This pool
+    handles async tokenization, chat template rendering, and multimodal
+    preprocessing."""
 
     # Pooler config
     pooler_config: PoolerConfig | None = None
@@ -639,6 +647,19 @@ class ModelConfig:
 
             self.multimodal_config = MultiModalConfig(**mm_config_kwargs)  # type: ignore[arg-type]
 
+            if (
+                self.renderer_num_workers > 1
+                and self.multimodal_config.mm_processor_cache_gb > 0
+            ):
+                raise ValueError(
+                    "Cannot use --renderer-num-workers > 1 with the "
+                    "multimodal processor cache enabled. The cache is "
+                    "not thread-safe and does not support concurrent "
+                    "renderer workers. Please set "
+                    "--renderer-num-workers 1 (the default), or "
+                    "disable the cache with --mm-processor-cache-gb 0."
+                )
+
         # Multimodal GGUF models must use original repo for mm processing
         if is_gguf(self.tokenizer) and self.is_multimodal_model:
             raise ValueError(
@@ -1508,10 +1529,11 @@ class ModelConfig:
     @property
     def score_type(self) -> ScoreType:
         """
-        Scoring API handles score/rerank for:\n
-        - "classify" task (score_type: cross-encoder models)\n
-        - "embed" task (score_type: bi-encoder models)\n
-        - "token_embed" task (score_type: late interaction models)\n
+        Scoring API handles score/rerank for:
+
+        - "classify" task (score_type: cross-encoder models)
+        - "embed" task (score_type: bi-encoder models)
+        - "token_embed" task (score_type: late interaction models)
         """
         # fixme: self._model_info.score_type is the score type before
         #  as_seq_cls_model, which is "bi-encoder", rather than the
@@ -1589,9 +1611,10 @@ class ModelConfig:
         such as the lm_head in a generation model,
         or the score or classifier in a classification model.
 
-        `head_dtype` currently only supports pooling models.\n
-        - The pooling model defaults to using fp32 head,
-        you can use --hf-overrides '{"head_dtype": "model"}' to disable it.
+        `head_dtype` currently only supports pooling models.
+
+        - The pooling model defaults to using fp32 head, you can use
+          --hf-overrides '{"head_dtype": "model"}' to disable it.
         """
 
         head_dtype = _get_head_dtype(

vllm/config/multimodal.py

@@ -146,14 +146,14 @@ class MultiModalConfig:
     parallelism (TP).
 
     - `"weights"`: Within the same vLLM engine, split the weights of
-        each layer across TP ranks. (default TP behavior)\n
+      each layer across TP ranks. (default TP behavior)
     - `"data"`: Within the same vLLM engine, split the batched input data
-        across TP ranks to process the data in parallel, while hosting
-        the full weights on each TP rank.
-        This batch-level DP is not to be confused with API request-level
-        DP (which is controlled by `--data-parallel-size`).
-        This is only supported on a per-model basis and falls back to
-        `"weights"` if the encoder does not support DP."""
+      across TP ranks to process the data in parallel, while hosting
+      the full weights on each TP rank.
+      This batch-level DP is not to be confused with API request-level
+      DP (which is controlled by `--data-parallel-size`).
+      This is only supported on a per-model basis and falls back to
+      `"weights"` if the encoder does not support DP."""
     mm_encoder_attn_backend: AttentionBackendEnum | None = None
     """Optional override for the multi-modal encoder attention backend when
     using vision transformers. Accepts any value from

vllm/config/parallel.py

@@ -148,10 +148,11 @@ class ParallelConfig:
     eplb_config: EPLBConfig = Field(default_factory=EPLBConfig)
     """Expert parallelism configuration."""
     expert_placement_strategy: ExpertPlacementStrategy = "linear"
-    """The expert placement strategy for MoE layers:\n
+    """The expert placement strategy for MoE layers:
+
     - "linear": Experts are placed in a contiguous manner. For example, with 4
       experts and 2 ranks, rank 0 will have experts [0, 1] and rank 1 will have
-      experts [2, 3].\n
+      experts [2, 3].
     - "round_robin": Experts are placed in a round-robin manner. For example,
       with 4 experts and 2 ranks, rank 0 will have experts [0, 2] and rank 1
       will have experts [1, 3]. This strategy can help improve load balancing
@@ -159,11 +160,11 @@ class ParallelConfig:
     all2all_backend: All2AllBackend = "allgather_reducescatter"
     """All2All backend for MoE expert parallel communication. Available options:
 
-    - "allgather_reducescatter": All2all based on allgather and reducescatter\n
-    - "deepep_high_throughput": Use deepep high-throughput kernels\n
-    - "deepep_low_latency": Use deepep low-latency kernels\n
-    - "mori": Use mori kernels\n
-    - "nixl_ep": Use nixl-ep kernels\n
+    - "allgather_reducescatter": All2all based on allgather and reducescatter
+    - "deepep_high_throughput": Use deepep high-throughput kernels
+    - "deepep_low_latency": Use deepep low-latency kernels
+    - "mori": Use mori kernels
+    - "nixl_ep": Use nixl-ep kernels
     - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl
     - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels"""
 

vllm/config/profiler.py

@@ -37,7 +37,7 @@ class ProfilerConfig:
     profiler: ProfilerKind | None = None
     """Which profiler to use. Defaults to None. Options are:
 
-    - 'torch': Use PyTorch profiler.\n
+    - 'torch': Use PyTorch profiler.
     - 'cuda': Use CUDA profiler."""
 
     torch_profiler_dir: str = ""

vllm/config/scheduler.py

@@ -106,11 +106,12 @@ class SchedulerConfig:
     max_num_batched_tokens in case max multimodal embedding size is larger."""
 
     policy: SchedulerPolicy = "fcfs"
-    """The scheduling policy to use:\n
-    - "fcfs" means first come first served, i.e. requests are handled in order
-    of arrival.\n
+    """The scheduling policy to use:
+
+    - "fcfs" means first come first served, i.e. requests are handled in order 
+      of arrival.
     - "priority" means requests are handled based on given priority (lower
-    value means earlier handling) and time of arrival deciding any ties)."""
+      value means earlier handling) and time of arrival deciding any ties)."""
 
     disable_chunked_mm_input: bool = False
     """If set to true and chunked prefill is enabled, we do not want to

vllm/config/utils.py

@@ -37,10 +37,12 @@ ConfigT = TypeVar("ConfigT", bound=DataclassInstance)
 
 
 @overload
+@dataclass_transform(field_specifiers=(PydanticField,))
 def config(cls: type[ConfigT]) -> type[ConfigT]: ...
 
 
 @overload
+@dataclass_transform(field_specifiers=(PydanticField,))
 def config(
     *, config: ConfigDict | None = None, **kwargs: Any
 ) -> Callable[[type[ConfigT]], type[ConfigT]]: ...

vllm/config/vllm.py

@@ -152,13 +152,11 @@ def enable_rope_kvcache_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
-    """Enable if using AITER RMSNorm and AITER Triton GEMMs
-    and hidden size is 2880 i.e. gpt-oss; otherwise Inductor handles fusion."""
+    """Enable if using AITER RMSNorm and hidden size is 2880 i.e. gpt-oss."""
     from vllm._aiter_ops import rocm_aiter_ops
 
     return (
         rocm_aiter_ops.is_rmsnorm_enabled()
-        and not rocm_aiter_ops.is_triton_gemm_enabled()
         and cfg.model_config is not None
         and cfg.model_config.get_hidden_size() == 2880
     )

vllm/distributed/eplb/eplb_state.py

@@ -729,13 +729,6 @@ class EplbState:
                     eplb_model_state.physical_to_logical_map.cpu(),
                 )
 
-                num_logical_experts = global_expert_load_window.shape[-1]
-                (new_logical_to_physical_map, new_logical_replica_count) = (
-                    compute_logical_maps(
-                        new_physical_to_logical_map, num_logical_experts
-                    )
-                )
-
                 # Update expert weights
                 rearrange_expert_weights_inplace(
                     eplb_model_state.physical_to_logical_map,
@@ -747,39 +740,11 @@ class EplbState:
                 )
 
                 if not is_profile:
-                    if (
-                        eplb_model_state.physical_to_logical_map.shape[1]
-                        != new_physical_to_logical_map.shape[1]
-                    ):
-                        eplb_model_state.physical_to_logical_map = (
-                            new_physical_to_logical_map.to(
-                                eplb_model_state.physical_to_logical_map.device
-                            )
-                        )
-                    else:
-                        eplb_model_state.physical_to_logical_map.copy_(
-                            new_physical_to_logical_map
-                        )
-                    max_physical_slots = new_logical_to_physical_map.shape[-1]
-                    assert (
-                        max_physical_slots
-                        <= eplb_model_state.logical_to_physical_map.shape[-1]
-                    )
-                    new_logical_to_physical_map = torch.nn.functional.pad(
-                        new_logical_to_physical_map,
-                        (
-                            0,
-                            eplb_model_state.logical_to_physical_map.shape[-1]
-                            - max_physical_slots,
-                        ),
-                        value=-1,
-                    )
-                    eplb_model_state.logical_to_physical_map.copy_(
-                        new_logical_to_physical_map
-                    )
-                    eplb_model_state.logical_replica_count.copy_(
-                        new_logical_replica_count
+                    _commit_eplb_maps(
+                        eplb_model_state,
+                        new_physical_to_logical_map=new_physical_to_logical_map,
                     )
+
                 if is_main_rank:
                     assert start_event is not None
                     assert end_event is not None
@@ -829,42 +794,6 @@ class EplbState:
                 is_profile=is_profile,
             )
 
-    def _update_layer_mapping_from_new(
-        self, model_state: EplbModelState, layer: int
-    ) -> None:
-        if model_state.new_physical_to_logical_map is None:
-            return
-
-        target_device = model_state.physical_to_logical_map.device
-        new_physical = model_state.new_physical_to_logical_map
-        # If the number of physical experts has changed, then the new map needs to
-        # be copied synchronously to avoid a race condition with the async worker
-        if model_state.physical_to_logical_map.shape[1] != new_physical.shape[1]:
-            model_state.physical_to_logical_map = new_physical.to(target_device)
-        else:
-            model_state.physical_to_logical_map[layer].copy_(
-                new_physical[layer].to(target_device, non_blocking=True)
-            )
-
-        num_logical_experts = model_state.logical_to_physical_map.shape[1]
-        new_logical, new_replica_count = compute_logical_maps(
-            new_physical[layer], num_logical_experts
-        )
-
-        logical_device = model_state.logical_to_physical_map.device
-        max_slots = model_state.logical_to_physical_map.shape[-1]
-        slot_delta = max_slots - new_logical.shape[-1]
-        if slot_delta > 0:
-            new_logical = torch.nn.functional.pad(
-                new_logical, (0, slot_delta), value=-1
-            )
-        model_state.logical_to_physical_map[layer].copy_(new_logical.to(logical_device))
-
-        replica_device = model_state.logical_replica_count.device
-        model_state.logical_replica_count[layer].copy_(
-            new_replica_count.to(replica_device)
-        )
-
     def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
         parallel_state = get_ep_group()
         cpu_group = getattr(parallel_state, "cpu_group", None)
@@ -936,7 +865,12 @@ class EplbState:
             model_state.buffer_consumed_event = consumed_event
 
             transferred_layer = model_state.layer_to_transfer
-            self._update_layer_mapping_from_new(model_state, transferred_layer)
+            assert model_state.new_physical_to_logical_map is not None
+            _commit_eplb_maps_for_layer(
+                model_state,
+                new_physical_to_logical_map=model_state.new_physical_to_logical_map,
+                layer=transferred_layer,
+            )
             # After the main thread consumes, advance layer_to_transfer
             model_state.layer_to_transfer += 1
             model_state.ep_buffer_ready = 0
@@ -1175,3 +1109,84 @@ def compute_logical_maps(
     if per_layer:
         return logical_to_physical_map_out.squeeze(0), logical_replica_count.squeeze(0)
     return logical_to_physical_map_out, logical_replica_count
+
+
+def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None:
+    src_padding = dst.shape[-1] - src.shape[-1]
+    assert src_padding >= 0
+    new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1)
+    dst.copy_(new_src)
+
+
+def _commit_eplb_maps_for_layer(
+    model_state: EplbModelState,
+    new_physical_to_logical_map: torch.Tensor,
+    layer: int,
+) -> None:
+    """
+    Per-layer version of _commit_eplb_maps that's used by the sync portion of EPLB
+    when running async EPLB. Copies all of the new_* maps into model_state. After this
+    function completes, the new mappings will become the current mappings and will be
+    visible to the model.
+    """
+
+    # Commit physical_to_logical_map
+    src = new_physical_to_logical_map[layer]
+    dst = model_state.physical_to_logical_map[layer]
+    assert src.shape == dst.shape, (
+        "The number of physical experts must stay the same while running Async EPLB. "
+        f"Current number of physical experts: {dst.shape[0]}. New number of physical "
+        f"experts {src.shape[0]}."
+    )
+    dst.copy_(src, non_blocking=True)
+
+    num_logical_experts = model_state.logical_to_physical_map.shape[1]
+    new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)
+    # Commit logical_to_physical_map
+    _pad_out_tensor(
+        src=new_logical,
+        dst=model_state.logical_to_physical_map[layer],
+    )
+
+    # Commit logical_replica_count
+    src = new_replica_count
+    dst = model_state.logical_replica_count[layer]
+    assert src.shape == dst.shape
+    dst.copy_(src, non_blocking=True)
+
+
+def _commit_eplb_maps(
+    model_state: EplbModelState,
+    new_physical_to_logical_map: torch.Tensor,
+) -> None:
+    """
+    Copies all of the new_* maps into model_state. After this function completes,
+    the new mappings will become the current mappings and will be visible to the
+    model.
+    """
+
+    # Commit physical_to_logical_map
+    src = new_physical_to_logical_map
+    dst = model_state.physical_to_logical_map
+
+    # Rare Case: When the number of physical experts has changed, discard the old
+    # physical to logical expert map and use the new one. This only happens when the
+    # number of GPUs available to vLLM changes while vLLM is running. Otherwise copy the
+    # new map into the old one.
+    if src.shape[1] != dst.shape[1]:
+        model_state.physical_to_logical_map = src.to(dst.device)
+    else:
+        dst.copy_(src, non_blocking=True)
+
+    num_logical_experts = model_state.logical_to_physical_map.shape[1]
+    new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)
+    # Commit logical_to_physical_map
+    _pad_out_tensor(
+        src=new_logical,
+        dst=model_state.logical_to_physical_map,
+    )
+
+    # Commit logical_replica_count
+    src = new_replica_count
+    dst = model_state.logical_replica_count
+    dst.copy_(src, non_blocking=True)

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

@@ -47,14 +47,17 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.request import RequestStatus
 
+logger = init_logger(__name__)
+
 try:
     from mooncake.engine import TransferEngine
-except ImportError as e:
-    raise ImportError(
+except ImportError:
+    logger.warning(
         "Please install mooncake by following the instructions at "
         "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "
         "to run VLLM with MooncakeTransferEngine."
-    ) from e
+    )
+    TransferEngine = None
 
 if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -64,8 +67,6 @@ if TYPE_CHECKING:
 ReqId = str  # Internal scheduler request ID
 TransferId = str  # KV transfer coordination ID (shared by P/D)
 
-logger = init_logger(__name__)
-
 
 @dataclass(frozen=True)
 class TransferRegion:
@@ -638,6 +639,9 @@ class MooncakeConnectorWorker:
     """Implementation of Worker side methods"""
 
     def __init__(self, vllm_config: VllmConfig, engine_id: str):
+        if TransferEngine is None:
+            logger.error("Mooncake is not available")
+            raise RuntimeError("Mooncake is not available")
         logger.info("Initializing Mooncake Transfer Engine worker %s", engine_id)
 
         self.vllm_config = vllm_config
@@ -721,9 +725,7 @@ class MooncakeConnectorWorker:
             # Start bootstrap server on global rank 0.
             if should_launch_bootstrap_server(vllm_config):
                 _, port = get_mooncake_bootstrap_addr(vllm_config)
-                self.bootstrap_server = MooncakeBootstrapServer(
-                    vllm_config, "0.0.0.0", port
-                )
+                self.bootstrap_server = MooncakeBootstrapServer("0.0.0.0", port)
                 self.bootstrap_server.start()
 
         if not self.is_kv_producer:
@@ -778,7 +780,9 @@ class MooncakeConnectorWorker:
             if self.sender_loop.is_running():
                 self.sender_loop.call_soon_threadsafe(self.sender_loop.stop)
                 self._sender_listener_t.join()
-            if should_launch_bootstrap_server(self.vllm_config):
+            if should_launch_bootstrap_server(self.vllm_config) and hasattr(
+                self, "bootstrap_server"
+            ):
                 self.bootstrap_server.shutdown()
         if not self.is_kv_producer and self.receiver_loop.is_running():
             self.receiver_loop.call_soon_threadsafe(self.receiver_loop.stop)

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

@@ -8,7 +8,6 @@ import uvicorn
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
-from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
 from vllm.logger import init_logger
 
@@ -38,7 +37,7 @@ class MooncakeBootstrapServer:
     Prefiller workers register their connection info (IP, port, ranks) here.
     """
 
-    def __init__(self, vllm_config: VllmConfig, host: str, port: int):
+    def __init__(self, host: str, port: int):
         self.workers: dict[int, EngineEntry] = {}
 
         self.host = host

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import defaultdict
+from dataclasses import replace
 
 import torch
 
@@ -18,7 +19,17 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.v1.attention.backend import AttentionBackend
-from vllm.v1.kv_offload.spec import OffloadingSpec
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    MambaSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.kv_offload.spec import (
+    CanonicalKVCacheRef,
+    CanonicalKVCaches,
+    CanonicalKVCacheTensor,
+    OffloadingSpec,
+)
 from vllm.v1.kv_offload.worker.worker import (
     OffloadingWorker,
     TransferSpec,
@@ -53,17 +64,13 @@ class OffloadingConnectorWorker:
         self._job_counter = job_id + 1
         return job_id
 
-    def _register_handlers(
-        self,
-        kv_caches: dict[str, torch.Tensor],
-        attn_backends: dict[str, type[AttentionBackend]],
-    ):
-        for src_cls, dst_cls, handler in self.spec.get_handlers(
-            kv_caches, attn_backends
-        ):
+    def _register_handlers(self, kv_caches: CanonicalKVCaches):
+        for src_cls, dst_cls, handler in self.spec.get_handlers(kv_caches):
             self.worker.register_handler(src_cls, dst_cls, handler)
 
-    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+    def register_kv_caches(
+        self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
+    ):
         layer_names = list(kv_caches.keys())
         layers = get_layers_from_vllm_config(
             self.spec.vllm_config,
@@ -73,16 +80,221 @@ class OffloadingConnectorWorker:
         attn_backends = {
             layer_name: layers[layer_name].get_attn_backend()
             for layer_name in layer_names
+            if layer_name in layers
         }
-        self._register_handlers(kv_caches, attn_backends)
+
+        # layer_name -> list of matching KV cache tensors
+        # such that each tensor starts with the num_blocks dimension.
+        # FlashAttention layers which use the (2, num_blocks, ...) layout
+        # will possibly map to 2 tensors, one per K and one per V.
+        # All other layers will probably map to a single tensor.
+        tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
+        # layer_name -> size of (un-padded) page in bytes
+        unpadded_page_size_bytes: dict[str, int] = {}
+        # layer_name -> size of page in bytes
+        page_size_bytes: dict[str, int] = {}
+        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+            group_layer_names = kv_cache_group.layer_names
+            group_kv_cache_spec = kv_cache_group.kv_cache_spec
+            if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
+                per_layer_specs = group_kv_cache_spec.kv_cache_specs
+            else:
+                per_layer_specs = {}
+            for layer_name in group_layer_names:
+                layer_kv_cache_spec = per_layer_specs.get(
+                    layer_name, group_kv_cache_spec
+                )
+                if isinstance(layer_kv_cache_spec, AttentionSpec):
+                    layer_kv_cache = kv_caches[layer_name]
+                    assert isinstance(layer_kv_cache, torch.Tensor)
+                    assert layer_kv_cache.storage_offset() == 0
+
+                    # get the logical dimension for num_blocks
+                    test_shape = attn_backends[layer_name].get_kv_cache_shape(
+                        num_blocks=1234,
+                        block_size=16,
+                        num_kv_heads=1,
+                        head_size=256,
+                    )
+                    num_blocks_logical_dim = test_shape.index(1234)
+
+                    # sort the logical dimensions by stride (high to low)
+                    # to get a physical-to-logical mapping:
+                    # physical_to_logical[physical_pos] = logical_dim
+                    logical_strides = layer_kv_cache.stride()
+                    physical_to_logical = sorted(
+                        range(len(logical_strides)),
+                        key=lambda idx: logical_strides[idx],
+                        reverse=True,
+                    )
+
+                    num_blocks_physical_dim = physical_to_logical.index(
+                        num_blocks_logical_dim
+                    )
+                    if num_blocks_physical_dim == 0:
+                        num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
+                        storage = layer_kv_cache.untyped_storage()
+                        page = layer_kv_cache_spec.page_size_bytes
+                        tensors_per_block[layer_name] = (
+                            torch.tensor(
+                                [],
+                                dtype=torch.int8,
+                                device=layer_kv_cache.device,
+                            )
+                            .set_(storage)
+                            .view(num_blocks, page),
+                        )
+                        page_size_bytes[layer_name] = (
+                            layer_kv_cache_spec.page_size_bytes
+                        )
+                    else:
+                        # Flash Attention case: (2, num_blocks, ...)
+                        assert test_shape[0] == 2
+                        assert physical_to_logical[0] == 0
+                        assert num_blocks_physical_dim == 1
+
+                        # unbind the tensor to separate K and V tensors
+                        num_blocks = layer_kv_cache.shape[num_blocks_logical_dim]
+                        half_page_size = layer_kv_cache_spec.page_size_bytes // 2
+                        storage = layer_kv_cache.untyped_storage()
+                        raw = (
+                            torch.tensor(
+                                [],
+                                dtype=torch.int8,
+                                device=layer_kv_cache.device,
+                            )
+                            .set_(storage)
+                            .view(2, num_blocks, half_page_size)
+                        )
+                        tensors_per_block[layer_name] = tuple(raw.unbind(0))
+
+                        page_size_bytes[layer_name] = half_page_size
+
+                    unpadded_page_size_bytes[layer_name] = page_size_bytes[layer_name]
+
+                elif isinstance(layer_kv_cache_spec, MambaSpec):
+                    state_tensors = kv_caches[layer_name]
+                    assert isinstance(state_tensors, list)
+
+                    # re-construct the raw (num_blocks, page_size) tensor
+                    # from the first state tensor
+                    assert len(state_tensors) > 0
+                    first_state_tensor = state_tensors[0]
+                    assert first_state_tensor.storage_offset() == 0
+                    num_blocks = first_state_tensor.shape[0]
+                    tensor = (
+                        torch.tensor(
+                            [],
+                            dtype=torch.int8,
+                            device=first_state_tensor.device,
+                        )
+                        .set_(first_state_tensor.untyped_storage())
+                        .view((num_blocks, layer_kv_cache_spec.page_size_bytes))
+                    )
+                    tensors_per_block[layer_name] = (tensor,)
+
+                    page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
+                    unpadded_page_size_bytes[layer_name] = replace(
+                        layer_kv_cache_spec, page_size_padded=None
+                    ).page_size_bytes
+
+                else:
+                    raise NotImplementedError
+
+        block_tensors: list[CanonicalKVCacheTensor] = []
+        block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
+        for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:
+            tensor_layer_names = kv_cache_tensor.shared_by
+
+            # verify all layers in the group reference the exact same tensors
+            assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1
+            assert (
+                len({tensors_per_block[n][0].data_ptr() for n in tensor_layer_names})
+                == 1
+            )
+            assert (
+                len({tensors_per_block[n][0].stride() for n in tensor_layer_names}) == 1
+            )
+
+            # pick the first layer to represent the group
+            first_layer_name = tensor_layer_names[0]
+            for tensor in tensors_per_block[first_layer_name]:
+                block_tensors.append(
+                    CanonicalKVCacheTensor(
+                        tensor=tensor,
+                        page_size_bytes=page_size_bytes[first_layer_name],
+                    )
+                )
+
+                curr_tensor_idx = len(block_tensors) - 1
+                for layer_name in tensor_layer_names:
+                    block_data_refs[layer_name].append(
+                        CanonicalKVCacheRef(
+                            tensor_idx=curr_tensor_idx,
+                            page_size_bytes=(unpadded_page_size_bytes[layer_name]),
+                        )
+                    )
+
+        group_data_refs: list[list[CanonicalKVCacheRef]] = []
+        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+            group_refs: list[CanonicalKVCacheRef] = []
+            for layer_name in kv_cache_group.layer_names:
+                group_refs += block_data_refs[layer_name]
+            group_data_refs.append(group_refs)
+
+        canonical_kv_caches = CanonicalKVCaches(
+            tensors=block_tensors,
+            group_data_refs=group_data_refs,
+        )
+
+        self._register_handlers(canonical_kv_caches)
 
     def register_cross_layers_kv_cache(
         self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
     ):
-        cross_layer_name = "ALL_LAYERS"
-        kv_caches = {cross_layer_name: kv_cache}
-        attn_backends = {cross_layer_name: attn_backend}
-        self._register_handlers(kv_caches, attn_backends)
+        # verify that num_blocks is at physical position 0 in the cross-layers
+        # tensor layout.
+        test_shape = attn_backend.get_kv_cache_shape(
+            num_blocks=1234, block_size=16, num_kv_heads=1, head_size=256
+        )
+        num_blocks_logical_dim = test_shape.index(1234) + 1
+        physical_to_logical = attn_backend.get_kv_cache_stride_order(
+            include_num_layers_dimension=True
+        )
+        num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim)
+        assert num_blocks_physical_dim == 0
+
+        kv_cache_groups = self.spec.kv_cache_config.kv_cache_groups
+        assert len(kv_cache_groups) == 1
+        kv_cache_spec = kv_cache_groups[0].kv_cache_spec
+        num_layers = len(kv_cache_groups[0].layer_names)
+        page_size_bytes = kv_cache_spec.page_size_bytes * num_layers
+
+        assert kv_cache.storage_offset() == 0
+        storage = kv_cache.untyped_storage()
+        assert len(storage) % page_size_bytes == 0
+        num_blocks = len(storage) // page_size_bytes
+        tensor = (
+            torch.tensor(
+                [],
+                dtype=torch.int8,
+                device=kv_cache.device,
+            )
+            .set_(storage)
+            .view(num_blocks, page_size_bytes)
+        )
+        kv_cache_tensor = CanonicalKVCacheTensor(
+            tensor=tensor, page_size_bytes=page_size_bytes
+        )
+        # in cross layers layout, there's currently only a single group
+        kv_cache_data_ref = CanonicalKVCacheRef(
+            tensor_idx=0, page_size_bytes=page_size_bytes
+        )
+        canonical_kv_caches = CanonicalKVCaches(
+            tensors=[kv_cache_tensor], group_data_refs=[[kv_cache_data_ref]]
+        )
+
+        self._register_handlers(canonical_kv_caches)
 
     def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
         for job_id, transfer_spec in self._unsubmitted_store_jobs:

vllm/engine/arg_utils.py

@@ -508,6 +508,7 @@ class EngineArgs:
         MultiModalConfig.mm_encoder_attn_backend
     )
     io_processor_plugin: str | None = None
+    renderer_num_workers: int = 1
     skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
     video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate
     mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc
@@ -596,6 +597,9 @@ class EngineArgs:
     attention_backend: AttentionBackendEnum | None = AttentionConfig.backend
 
     calculate_kv_scales: bool = CacheConfig.calculate_kv_scales
+    kv_cache_dtype_skip_layers: list[str] = get_field(
+        CacheConfig, "kv_cache_dtype_skip_layers"
+    )
     mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype
     mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype
     mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
@@ -767,6 +771,10 @@ class EngineArgs:
         model_group.add_argument(
             "--io-processor-plugin", **model_kwargs["io_processor_plugin"]
         )
+        model_group.add_argument(
+            "--renderer-num-workers",
+            **model_kwargs["renderer_num_workers"],
+        )
 
         # Model loading arguments
         load_kwargs = get_kwargs(LoadConfig)
@@ -998,6 +1006,9 @@ class EngineArgs:
         cache_group.add_argument(
             "--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"]
         )
+        cache_group.add_argument(
+            "--kv-cache-dtype-skip-layers", **cache_kwargs["kv_cache_dtype_skip_layers"]
+        )
         cache_group.add_argument(
             "--kv-sharing-fast-prefill", **cache_kwargs["kv_sharing_fast_prefill"]
         )
@@ -1290,7 +1301,7 @@ class EngineArgs:
         # delay the Pydantic validation that comes with SpeculativeConfig.
         vllm_kwargs["speculative_config"]["type"] = optional_type(json.loads)
         vllm_group.add_argument(
-            "--speculative-config", **vllm_kwargs["speculative_config"]
+            "--speculative-config", "-sc", **vllm_kwargs["speculative_config"]
         )
         vllm_group.add_argument(
             "--kv-transfer-config", **vllm_kwargs["kv_transfer_config"]
@@ -1438,6 +1449,7 @@ class EngineArgs:
             video_pruning_rate=self.video_pruning_rate,
             mm_tensor_ipc=self.mm_tensor_ipc,
             io_processor_plugin=self.io_processor_plugin,
+            renderer_num_workers=self.renderer_num_workers,
         )
 
     def validate_tensorizer_args(self):
@@ -1572,6 +1584,7 @@ class EngineArgs:
             enable_prefix_caching=self.enable_prefix_caching,
             prefix_caching_hash_algo=self.prefix_caching_hash_algo,
             calculate_kv_scales=self.calculate_kv_scales,
+            kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers,
             kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,
             mamba_cache_dtype=self.mamba_cache_dtype,
             mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,

vllm/entrypoints/openai/responses/protocol.py

@@ -105,8 +105,8 @@ def serialize_message(msg):
     elif hasattr(msg, "to_dict"):
         return msg.to_dict()
     else:
-        # fallback to pyandic dump
-        return msg.model_dump_json()
+        # fallback to pydantic dump
+        return msg.model_dump_json(by_alias=True)
 
 
 def serialize_messages(msgs):

vllm/envs.py

@@ -86,6 +86,7 @@ if TYPE_CHECKING:
     VLLM_ALLOW_LONG_MAX_MODEL_LEN: bool = False
     VLLM_RPC_TIMEOUT: int = 10000  # ms
     VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5  # seconds
+    VLLM_MAX_N_SEQUENCES: int = 16384
     VLLM_PLUGINS: list[str] | None = None
     VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None
     VLLM_LORA_RESOLVER_HF_REPO_LIST: str | None = None
@@ -870,6 +871,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_HTTP_TIMEOUT_KEEP_ALIVE": lambda: int(
         os.environ.get("VLLM_HTTP_TIMEOUT_KEEP_ALIVE", "5")
     ),
+    # Maximum allowed value for the `n` sampling parameter (number of output
+    # sequences per request). Limits resource consumption to prevent
+    # denial-of-service via excessively large fan-out. Default: 16384.
+    "VLLM_MAX_N_SEQUENCES": lambda: int(
+        os.environ.get("VLLM_MAX_N_SEQUENCES", "16384")
+    ),
     # a list of plugin names to load, separated by commas.
     # if this is not set, it means all plugins will be loaded
     # if this is set to an empty string, no plugins will be loaded

vllm/model_executor/kernels/linear/mixed_precision/cpu.py

@@ -58,28 +58,29 @@ class CPUWNA16LinearKernel(MPLinearKernel):
         return True, None
 
     # note assumes that
-    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
-    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
-    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1}
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} (marlin)
+    #                  or: {input_dim = 1, output_dim = 0, packed_dim = 1} (CT)
+    #  `weight_scale`  is: {input_dim = 0, output_dim = 1} (marlin)
+    #                  or: {input_dim = 1, output_dim = 0} (CT)
+    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1} (marlin)
+    #                  or: {input_dim = 1, output_dim = 0, packed_dim = 0} (CT)
     def _process_gptq_weights(self, layer: torch.nn.Module):
-        packed_weight = layer.qweight.data
+        packed_weight = getattr(layer, self.w_q_name)
+        assert packed_weight.input_dim == packed_weight.packed_dim
+        is_ct_format = packed_weight.input_dim == 1
+        if is_ct_format:
+            packed_weight = packed_weight.t()
         bits = self.config.weight_type.mantissa
         pack_factor = 32 // bits
-        p_w_k, p_w_n = packed_weight.size()
+        p_w_k, _ = packed_weight.size()
         input_size = p_w_k * pack_factor
-        output_size = p_w_n
-        isa_hint = _get_isa_hint(layer.scales.dtype)
+        isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype)
         layer.isa_hint = isa_hint
 
-        layer.qzeros = None
-        if not self.config.has_g_idx:
-            layer.g_idx = None
-
         # convert input dim packed to output dim packed
         weight = unpack_quantized_values_into_int32(
-            packed_weight, self.config.weight_type, 1
-        ).view(p_w_k, p_w_n, pack_factor)
-        weight = weight.permute(0, 2, 1).reshape(input_size, output_size).contiguous()
+            packed_weight, self.config.weight_type, 0
+        )
         weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1)
         # make 16 output channel as a block and transpose to the make
         # the block contiguous
@@ -89,10 +90,29 @@ class CPUWNA16LinearKernel(MPLinearKernel):
             .reshape(-1, input_size * 16 // pack_factor)
             .contiguous()
         )
-        layer.qweight.data = weight
+        getattr(layer, self.w_q_name).data = weight
+
+        # transpose scale, zp for CT format
+        if is_ct_format:
+            scales = getattr(layer, self.w_s_name)
+            scales.data = scales.t().contiguous()
+            if self.config.zero_points:
+                assert self.w_zp_name
+                zp = getattr(layer, self.w_zp_name)
+                zp.data = zp.t().contiguous()
 
     def process_weights_after_loading(self, layer: torch.nn.Module):
-        if not self.config.zero_points:
+        if (not self.config.zero_points) and (self.w_zp_name is not None):
+            setattr(layer, self.w_zp_name, None)
+
+        if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
+            setattr(layer, self.w_gidx_name, None)
+
+        w_input_dim = getattr(layer, self.w_q_name).input_dim
+        w_pack_dim = getattr(layer, self.w_q_name).packed_dim
+        quant_method = "gptq" if w_pack_dim == w_input_dim else "awq"
+
+        if quant_method == "gptq":
             # GPTQ
             self._process_gptq_weights(layer)
         else:
@@ -105,12 +125,13 @@ class CPUWNA16LinearKernel(MPLinearKernel):
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
         x = ops.cpu_gemm_wna16(
             input=x,
-            q_weight=layer.qweight,
-            scales=layer.scales,
-            zeros=layer.qzeros,
-            g_idx=layer.g_idx,
+            q_weight=w_q,
+            scales=w_s,
+            zeros=w_zp,
+            g_idx=w_gidx,
             bias=bias,
             pack_factor=8,  # 32 // 4
             isa_hint=layer.isa_hint,

vllm/model_executor/layers/attention/attention.py

@@ -240,6 +240,31 @@ class Attention(nn.Module, AttentionLayerBase):
             and kv_cache_scheme.get("strategy") == "attn_head"
         )
 
+        # Skip quantization for specified layers
+        if cache_config is not None and cache_config.kv_cache_dtype_skip_layers:
+            from vllm.model_executor.models.utils import extract_layer_index
+
+            skip = False
+            # Check attention type
+            if (
+                sliding_window is not None
+                and "sliding_window" in cache_config.kv_cache_dtype_skip_layers
+            ):
+                skip = True
+            # Check layer index
+            layer_idx = extract_layer_index(prefix)
+            if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers:
+                skip = True
+            if skip:
+                kv_cache_dtype = "auto"
+                calculate_kv_scales = False
+            logger.info(
+                "Layer %s: kv_cache_dtype=%s, sliding_window=%s",
+                prefix,
+                kv_cache_dtype,
+                sliding_window,
+            )
+
         self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
             kv_cache_dtype, vllm_config.model_config
         )

vllm/model_executor/layers/fused_moe/config.py

@@ -229,9 +229,6 @@ class FusedMoEQuantConfig:
     _w1: FusedMoEQuantDesc
     _w2: FusedMoEQuantDesc
     is_nvfp4_scale_swizzled: bool = True
-    # CK MXFP4 (gfx950) padding info for rocm_aiter_ops.fused_moe()
-    hidden_pad: int = 0
-    intermediate_pad: int = 0
 
     def __post_init__(self):
         assert not self.per_act_token_quant or self.block_shape is None, (
@@ -1172,6 +1169,11 @@ class FusedMoEConfig:
     # Defaults to in_dtype if not specified.
     router_logits_dtype: torch.dtype | None = None
 
+    # Defaults to hidden_dim if not specified.
+    hidden_dim_unpadded: int | None = None
+    # Defaults to intermediate_size_per_partition if not specified.
+    intermediate_size_per_partition_unpadded: int | None = None
+
     moe_backend: str = "auto"
     max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
     has_bias: bool = False
@@ -1195,6 +1197,13 @@ class FusedMoEConfig:
         if self.router_logits_dtype is None:
             self.router_logits_dtype = self.in_dtype
 
+        if self.hidden_dim_unpadded is None:
+            self.hidden_dim_unpadded = self.hidden_dim
+        if self.intermediate_size_per_partition_unpadded is None:
+            self.intermediate_size_per_partition_unpadded = (
+                self.intermediate_size_per_partition
+            )
+
     @property
     def tp_size(self):
         return self.moe_parallel_config.tp_size

vllm/model_executor/layers/fused_moe/configs/E=128,N=192,device_name=NVIDIA_H800,dtype=fp8_w8a8.json

@@ -0,0 +1,147 @@
+{
+    "triton_version": "3.5.0",
+    "1": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "2": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 5
+    },
+    "4": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "8": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "16": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "24": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "32": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "48": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "96": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "128": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "256": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "512": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 8,
+        "num_stages": 2
+    },
+    "1024": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "1536": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "2048": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "3072": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "4096": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 8,
+        "num_stages": 3
+    }
+}

vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=NVIDIA_H100_80GB_HBM3.json

@@ -0,0 +1,147 @@
+{
+    "triton_version": "3.5.1",
+    "1": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "2": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 32,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 4
+    },
+    "4": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "8": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "16": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "24": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 8,
+        "num_stages": 4
+    },
+    "32": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "48": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "64": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 8,
+        "num_stages": 2
+    },
+    "96": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "128": {
+        "BLOCK_SIZE_M": 16,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "256": {
+        "BLOCK_SIZE_M": 32,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 32,
+        "num_warps": 4,
+        "num_stages": 2
+    },
+    "512": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3
+    },
+    "1024": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "1536": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 8,
+        "num_stages": 3
+    },
+    "2048": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 1,
+        "num_warps": 8,
+        "num_stages": 4
+    },
+    "3072": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 8,
+        "num_stages": 4
+    },
+    "4096": {
+        "BLOCK_SIZE_M": 128,
+        "BLOCK_SIZE_N": 256,
+        "BLOCK_SIZE_K": 64,
+        "GROUP_SIZE_M": 16,
+        "num_warps": 8,
+        "num_stages": 4
+    }
+}

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

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
+from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     activation_to_flashinfer_int,
 )
@@ -152,11 +153,8 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
         import flashinfer
         from flashinfer.fused_moe import Fp8QuantizationType
 
-        # Pack topk_ids and topk_weights into single tensor
-        # Format: (expert_id << 16) | (weight_bf16.view(int16))
-        packed_topk_ids = (topk_ids << 16) | topk_weights.to(torch.bfloat16).view(
-            torch.int16
-        )
+        # Pack topk ids and weights into format expected by the kernel.
+        packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
 
         # trtllm_fp8_block_scale_routed_moe does not support autotuning
         # so skip this kernel during dummy run for autotuning.

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

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
+from vllm.model_executor.layers.fused_moe.utils import trtllm_moe_pack_topk_ids_weights
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     activation_to_flashinfer_int,
 )
@@ -183,9 +184,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
         assert self.quant_config.w2_scale is not None
 
         # Pack topk ids and weights into format expected by the kernel.
-        packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
-            torch.bfloat16
-        ).view(torch.int16)
+        packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
 
         # trtllm_fp4_block_scale_routed_moe does not support autotuning
         # so skip this kernel during dummy run for autotuning.

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

@@ -9,6 +9,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
+    FusedMoEParallelConfig,
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -65,6 +66,38 @@ class FusedMoEMethodBase(QuantizeMethodBase):
         """
         return False
 
+    def maybe_roundup_sizes(
+        self,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        act_dtype: torch.dtype,
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> tuple[int, int]:
+        """
+        Given layer hidden size and intermediate size per partition and MoE
+        configurations, round up hidden_size and intermediate_size_per_partition
+        if necessary.
+
+        Args:
+            hidden_size: Layer hidden-size
+            intermediate_size_per_partition: Intermediate size per partition for
+                the layer.
+            act_dtype: Data type of the layer activations.
+            moe_parallel_config: Fused MoE parallelization strategy configuration.
+
+        Return:
+            A tuple of (rounded_hidden_size, rounded_intermediate_size_per_partition),
+            where:
+                - rounded_hidden_size is the possibly rounded up hidden size.
+                - rounded_intermediate_size_per_partition is the possibly rounded
+                  up intermediate size per partition.
+        """
+        from .all2all_utils import maybe_roundup_layer_hidden_size
+
+        return maybe_roundup_layer_hidden_size(
+            hidden_size, act_dtype, moe_parallel_config
+        ), intermediate_size_per_partition
+
     def maybe_make_prepare_finalize(
         self,
         routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

@@ -428,13 +428,9 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
     assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
     assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
 
-    # Shape check: when weights are padded (e.g. hidden_size padded for
-    # GFX950 swizzle), unpadded_K_w1 carries the original dimension.
-    expected_K_w1 = unpadded_K_w1 if unpadded_K_w1 is not None else w1.shape[-2]
-    assert hidden_states.shape[-1] == expected_K_w1, (
-        f"hidden_states K={hidden_states.shape[-1]} != "
-        f"expected K={expected_K_w1} (w1 K={w1.shape[-2]})"
-    )
+    # Shape check: weights are padded (e.g. hidden_size padded for
+    # GFX950 swizzle).
+    assert hidden_states.shape[-1] == w1.shape[-2]
     assert w2.shape[-1] == w1.shape[1]
 
     E, _, N = w1.shape
@@ -494,12 +490,6 @@ def triton_kernel_fused_mxfp4_w4a8_experts(
         unpadded_K=unpadded_K_w2,
     )
 
-    # When hidden_size was padded for alignment (e.g. GFX950 swizzle),
-    # the kernel output has the padded dimension. Slice back to the
-    # original hidden_size so downstream layers see the expected shape.
-    if unpadded_N_w2 is not None and intermediate_cache3.shape[-1] != unpadded_N_w2:
-        intermediate_cache3 = intermediate_cache3[..., :unpadded_N_w2].contiguous()
-
     return intermediate_cache3
 
 

vllm/model_executor/layers/fused_moe/layer.py

@@ -210,42 +210,6 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
     )
 
 
-# TODO(rob): move this down to the kernel.
-def maybe_roundup_hidden_size(
-    hidden_size: int,
-    act_dtype: torch.dtype,
-    moe_parallel_config: FusedMoEParallelConfig,
-    is_lora_enabled: bool,
-    model_type: str | None,
-) -> int:
-    """
-    Given layer hidden size and MoE configurations, round up hidden_size
-    if necessary.
-
-    Args:
-        hidden_size: Layer hidden-size
-        act_dtype: Data type of the layer activations.
-        moe_parallel_config: Fused MoE parallelization strategy configuration.
-        is_lora_enabled: True if the engine is enabled with LoRA. This
-            is used in the case of mxfp4 quantization in selecting the
-            MxFP4Backend.
-        model_type: for checking if gpt-oss
-
-    Return:
-        Rounded up hidden_size if rounding up is required based on the configs.
-        Original hidden size otherwise.
-    """
-    from vllm.model_executor.layers.fused_moe.all2all_utils import (
-        maybe_roundup_layer_hidden_size,
-    )
-
-    hidden_size = maybe_roundup_layer_hidden_size(
-        hidden_size, act_dtype, moe_parallel_config
-    )
-
-    return hidden_size
-
-
 # --8<-- [start:fused_moe]
 @CustomOp.register("fused_moe")
 class FusedMoE(CustomOp):
@@ -459,7 +423,7 @@ class FusedMoE(CustomOp):
             ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."
 
         assert intermediate_size % self.tp_size == 0
-        self.intermediate_size_per_partition = intermediate_size // self.tp_size
+        intermediate_size_per_partition = intermediate_size // self.tp_size
         self.reduce_results = reduce_results
         self.renormalize = renormalize
 
@@ -501,28 +465,13 @@ class FusedMoE(CustomOp):
         )
         self.routing_method_type: RoutingMethodType = self.router.routing_method_type
 
-        # Round up hidden size before creating moe_config.
-        # This way moe_config is created with the correct hidden_size from the start.
-        unpadded_hidden_size = hidden_size
-        self.model_type = (
-            self.vllm_config.model_config.hf_config.model_type
-            if self.vllm_config.model_config is not None
-            else None
-        )
-        hidden_size = maybe_roundup_hidden_size(
-            hidden_size=hidden_size,
-            act_dtype=moe_in_dtype,
-            moe_parallel_config=self.moe_parallel_config,
-            is_lora_enabled=vllm_config.lora_config is not None,
-            model_type=self.model_type,
-        )
-        self.hidden_size = hidden_size
-
         self.moe_config: FusedMoEConfig = FusedMoEConfig(
             num_experts=self.global_num_experts,
             experts_per_token=top_k,
             hidden_dim=hidden_size,
-            intermediate_size_per_partition=self.intermediate_size_per_partition,
+            hidden_dim_unpadded=hidden_size,
+            intermediate_size_per_partition=intermediate_size_per_partition,
+            intermediate_size_per_partition_unpadded=intermediate_size_per_partition,
             num_local_experts=self.local_num_experts,
             num_logical_experts=self.logical_num_experts,
             moe_parallel_config=self.moe_parallel_config,
@@ -567,13 +516,6 @@ class FusedMoE(CustomOp):
         # for heuristic purposes, so it must be initialized first.
         self.quant_method: FusedMoEMethodBase = _get_quant_method()
 
-        # Quant methods (e.g. Mxfp4MoEMethod) may round up hidden_dim
-        # and intermediate_size in moe_config during __init__. Sync
-        # self.hidden_size so downstream consumers (e.g. LoRA) see the
-        # padded value.
-        if self.moe_config.hidden_dim != self.hidden_size:
-            self.hidden_size = self.moe_config.hidden_dim
-
         if not self.moe_config.is_act_and_mul and not current_platform.is_cuda_alike():
             raise NotImplementedError(
                 "is_act_and_mul=False is supported only for CUDA and ROCm for now"
@@ -591,11 +533,24 @@ class FusedMoE(CustomOp):
                 f"EPLB is not supported {self.quant_method.__class__.__name__}."
             )
 
+        # Round up hidden size and update moe_config.
+        hidden_size, intermediate_size_per_partition = (
+            self.quant_method.maybe_roundup_sizes(
+                hidden_size,
+                intermediate_size_per_partition,
+                moe_in_dtype,
+                self.moe_parallel_config,
+            )
+        )
+        self.moe_config.hidden_dim = hidden_size
+        self.moe_config.intermediate_size_per_partition = (
+            intermediate_size_per_partition
+        )
+
         moe_quant_params = {
             "num_experts": self.local_num_experts,
-            "hidden_size": self.hidden_size,
-            "unpadded_hidden_size": unpadded_hidden_size,
-            "intermediate_size_per_partition": self.intermediate_size_per_partition,
+            "hidden_size": hidden_size,
+            "intermediate_size_per_partition": intermediate_size_per_partition,
             "params_dtype": params_dtype,
             "weight_loader": self.weight_loader,
             "global_num_experts": self.global_num_experts,
@@ -933,9 +888,17 @@ class FusedMoE(CustomOp):
         # Only narrow if the loaded_weight is not a scalar (0-dim tensor)
         # and we're not loading the full weight
         if not load_full and loaded_weight.ndim > 0:
-            loaded_weight = loaded_weight.narrow(
-                shard_dim, shard_size * tp_rank, shard_size
-            )
+            # Handle padding: loaded_weight might be smaller than shard_size on last
+            # TP rank
+            start_offset = shard_size * tp_rank
+            available = loaded_weight.shape[shard_dim] - start_offset
+            if available <= 0:
+                # If there is no available weight to load for this TP rank
+                # (can happen on last TP rank with padding), we can skip
+                # loading and return early
+                return
+            narrow_size = min(shard_size, available)
+            loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)
         # Narrow parameter and load.
         # w1, gate_proj: Load into first logical weight of w13.
         if shard_id == "w1":
@@ -944,6 +907,13 @@ class FusedMoE(CustomOp):
         else:
             assert shard_id == "w3"
             expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)
+
+        # Handle padding: if loaded_weight is smaller than expert_data (can happen
+        # on last TP shard with padding), copy to top-left corner
+        if expert_data.shape != loaded_weight.shape:
+            expert_data = expert_data[
+                : loaded_weight.shape[0], : loaded_weight.shape[1]
+            ]
         expert_data.copy_(loaded_weight)
 
     def _load_w2(
@@ -961,10 +931,24 @@ class FusedMoE(CustomOp):
         # Only narrow if the loaded_weight is not a scalar (0-dim tensor)
         # and we're not loading the full weight
         if not load_full and loaded_weight.ndim > 0:
-            loaded_weight = loaded_weight.narrow(
-                shard_dim, shard_size * tp_rank, shard_size
-            )
+            # Handle padding: loaded_weight might be smaller than shard_size on last
+            # TP rank
+            start_offset = shard_size * tp_rank
+            available = loaded_weight.shape[shard_dim] - start_offset
+            if available <= 0:
+                # If there is no available weight to load for this TP rank
+                # (can happen on last TP rank with padding), we can skip
+                # loading and return early
+                return
+            narrow_size = min(shard_size, available)
+            loaded_weight = loaded_weight.narrow(shard_dim, start_offset, narrow_size)
         # w2, down_proj: Load into only logical weight of w2.
+        # Handle padding: if loaded_weight is smaller than expert_data (can happen
+        # on last TP shard with padding), copy to top-left corner
+        if expert_data.shape != loaded_weight.shape:
+            expert_data = expert_data[
+                : loaded_weight.shape[0], : loaded_weight.shape[1]
+            ]
         expert_data.copy_(loaded_weight)
 
     def _load_single_value(
@@ -1549,6 +1533,14 @@ class FusedMoE(CustomOp):
             ]
         ]
 
+    @property
+    def hidden_size(self) -> int:
+        return self.moe_config.hidden_dim
+
+    @property
+    def intermediate_size_per_partition(self) -> int:
+        return self.moe_config.intermediate_size_per_partition
+
     def extra_repr(self) -> str:
         s = (
             f"global_num_experts={self.global_num_experts}, "

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

@@ -20,10 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     mxfp4_w4a16_moe_quant_config,
     ocp_mx_moe_quant_config,
 )
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
-    _swizzle_mxfp4,
-    get_padding_alignment,
-)
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kMxfp4Static,
@@ -396,9 +393,8 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
         intermediate_size = round_up(intermediate_size, 128)
         hidden_size = round_up(hidden_size, 128)
     elif current_platform.is_rocm():
-        pad_align = get_padding_alignment()
-        intermediate_size = round_up(intermediate_size, pad_align)
-        hidden_size = round_up(hidden_size, pad_align)
+        intermediate_size = round_up(intermediate_size, 256)
+        hidden_size = round_up(hidden_size, 256)
     else:
         intermediate_size = round_up(intermediate_size, 64)
     return hidden_size, intermediate_size
@@ -779,8 +775,6 @@ def make_mxfp4_moe_quant_config(
     w2_scale: Union[torch.Tensor, "PrecisionConfig"],
     w1_bias: torch.Tensor | None = None,
     w2_bias: torch.Tensor | None = None,
-    hidden_pad: int = 0,
-    intermediate_pad: int = 0,
 ) -> FusedMoEQuantConfig | None:
     """Create a FusedMoEQuantConfig for the given MXFP4 backend."""
     if mxfp4_backend in (
@@ -802,16 +796,12 @@ def make_mxfp4_moe_quant_config(
         Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
         Mxfp4MoeBackend.CK,
     ):
-        config = mxfp4_w4a16_moe_quant_config(
+        return mxfp4_w4a16_moe_quant_config(
             w1_bias=w1_bias,
             w2_bias=w2_bias,
             w1_scale=w1_scale,
             w2_scale=w2_scale,
         )
-        if mxfp4_backend == Mxfp4MoeBackend.CK:
-            config.hidden_pad = hidden_pad
-            config.intermediate_pad = intermediate_pad
-        return config
     else:
         return ocp_mx_moe_quant_config(
             quant_dtype="mxfp4",

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

@@ -10,6 +10,7 @@ from vllm._aiter_ops import rocm_aiter_ops
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FUSED_MOE_UNQUANTIZED_CONFIG,
+    FusedMoEConfig,
     FusedMoEParallelConfig,
     FusedMoEQuantConfig,
 )
@@ -186,6 +187,7 @@ def rocm_aiter_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    moe_config: FusedMoEConfig,
     activation: MoEActivation = MoEActivation.SILU,
     apply_router_weight_on_input: bool = False,
     expert_map: torch.Tensor | None = None,
@@ -276,6 +278,17 @@ def rocm_aiter_fused_experts(
                 "Only support topk=1 when `apply_router_weight_on_input` is True"
             )
 
+        # Compute padding on-the-fly for CK MXFP4 kernels
+        hidden_pad = 0
+        intermediate_pad = 0
+        assert moe_config.hidden_dim_unpadded is not None
+        assert moe_config.intermediate_size_per_partition_unpadded is not None
+        hidden_pad = hidden_states.shape[1] - moe_config.hidden_dim_unpadded
+        intermediate_pad = (
+            moe_config.intermediate_size_per_partition
+            - moe_config.intermediate_size_per_partition_unpadded
+        )
+
         return rocm_aiter_ops.fused_moe(
             hidden_states,
             w1,
@@ -292,8 +305,8 @@ def rocm_aiter_fused_experts(
             doweight_stage1=apply_router_weight_on_input,
             num_local_tokens=num_local_tokens,
             output_dtype=output_dtype,
-            hidden_pad=quant_config.hidden_pad,
-            intermediate_pad=quant_config.intermediate_pad,
+            hidden_pad=hidden_pad,
+            intermediate_pad=intermediate_pad,
             bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
             bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
         )
@@ -419,6 +432,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):
             apply_router_weight_on_input=apply_router_weight_on_input,
             expert_map=expert_map,
             quant_config=self.quant_config,
+            moe_config=self.moe_config,
             a1q_scale=a1q_scale,
             num_local_tokens=num_local_tokens,
             output_dtype=output.dtype,

vllm/model_executor/layers/fused_moe/utils.py

@@ -323,3 +323,16 @@ def normalize_batched_scales_shape(
 @functools.cache
 def disable_inplace() -> bool:
     return is_torch_equal_or_newer("2.9")
+
+
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+def trtllm_moe_pack_topk_ids_weights(
+    topk_ids: torch.Tensor, topk_weights: torch.Tensor
+) -> torch.Tensor:
+    """
+    Pack topk_ids and topk_weights into a single int32 tensor.
+    Format: (expert_id << 16) | weight_bf16.view(int16)
+    """
+    return (topk_ids.to(torch.int32) << 16) | topk_weights.to(torch.bfloat16).view(
+        torch.int16
+    )

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

@@ -715,8 +715,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        layer.intermediate_size_per_partition = intermediate_size_per_partition
-        layer.hidden_size = hidden_size
         layer.num_experts = num_experts
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None
@@ -2274,8 +2272,6 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        layer.intermediate_size_per_partition = intermediate_size_per_partition
-        layer.hidden_size = hidden_size
         layer.num_experts = num_experts
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None

vllm/model_executor/layers/quantization/fp8.py

@@ -73,7 +73,9 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     cutlass_fp8_supported,
     normalize_e4m3fn_to_e4m3fnuz,
 )
-from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
+from vllm.model_executor.model_loader.reload.layerwise import (
+    initialize_online_processing,
+)
 from vllm.model_executor.parameter import (
     BlockQuantScaleParameter,
     ModelWeightParameter,
@@ -496,8 +498,8 @@ class Fp8LinearMethod(LinearMethodBase):
 
 
 class Fp8OnlineLinearMethod(Fp8LinearMethod):
-    """Online version of Fp8LinearMethod, loads the fp16/bf16 checkpoint
-    and quantized the weights during loading."""
+    """Online version of Fp8LinearMethod which loads a full precision checkpoint
+    and quantizes weights during loading."""
 
     uses_meta_device: bool = True
 
@@ -519,84 +521,25 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None
 
-        # WEIGHT
-        def patched_weight_loader(param, loaded_weight, *args, **kwargs):
-            # track how many elements we have updated
-            if not hasattr(layer, "_loaded_numel"):
-                layer._loaded_numel = 0
-
-                # when the first `loaded_weight` is about to be
-                # loaded to `param`, materialize `param` just-in-time
-                weight = ModelWeightParameter(
-                    data=torch.empty_like(layer.weight, device=layer._load_device),
-                    input_dim=1,
-                    output_dim=0,
-                    weight_loader=patched_weight_loader,
-                )
-                _copy_missing_attrs(layer.weight, weight)
-                layer.register_parameter("weight", weight)
-                del layer._load_device
-
-            # refresh the reference to `param` to reflect just-in-time
-            # materialization
-            param = layer.weight
-
-            # load the current weight chunk
-            copy_numel_counter = CopyNumelCounter()
-            with copy_numel_counter:
-                res = weight_loader(param, loaded_weight, *args, **kwargs)  # type: ignore[misc]
-            layer._loaded_numel += copy_numel_counter.copied_numel
-
-            # if we have loaded all of the elements, call
-            # process_weights_after_loading
-            target_loaded_numel = layer.weight.numel()
-            if layer._loaded_numel == target_loaded_numel:
-                self.process_weights_after_loading(layer)
-
-                # Prevent the usual `process_weights_after_loading` call from doing
-                # anything
-                layer._already_called_process_weights_after_loading = True
-
-                # Note that we keep `layer._loaded_numel` around just in case
-                # there is logic added to vllm in the future which calls a
-                # weight loader twice - we do not want to re-initialize in
-                # that case.
-
-            return res
-
         weight = ModelWeightParameter(
             data=torch.empty(
                 output_size_per_partition,
                 input_size_per_partition,
-                # materialized just-in-time in `patched_weight_loader`
-                device="meta",
+                device="meta",  # materialized and processed during loading
                 dtype=params_dtype,
             ),
             input_dim=1,
             output_dim=0,
-            weight_loader=patched_weight_loader,
+            weight_loader=weight_loader,
         )
-        # stash the correct device for `patched_weight_loader`
-        layer._load_device = torch.get_default_device()
         layer.register_parameter("weight", weight)
 
+        initialize_online_processing(layer)
+
     def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return
 
-        # deferred initialization of randomly initialized weights for the
-        # `--load_format dummy` feature
-        if layer.weight.device == torch.device("meta"):
-            weight = ModelWeightParameter(
-                data=torch.empty_like(layer.weight, device=layer._load_device),
-                input_dim=1,
-                output_dim=0,
-                weight_loader=layer.weight.weight_loader,
-            )
-            _copy_missing_attrs(layer.weight, weight)
-            layer.register_parameter("weight", weight)
-            initialize_single_dummy_weight(layer.weight)
-
         # TODO(future): support block_quant in online quant path
         assert not self.block_quant
 
@@ -672,8 +615,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        layer.intermediate_size_per_partition = intermediate_size_per_partition
-        layer.hidden_size = hidden_size
         layer.num_experts = num_experts
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None
@@ -847,9 +788,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             )
 
     def process_weights_after_loading(self, layer: Module) -> None:
-        if getattr(layer, "_already_called_process_weights_after_loading", False):
-            return
-
         # Allow for accessing weights and scales in standard way.
         w13 = layer.w13_weight
         w2 = layer.w2_weight
@@ -894,9 +832,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
         )
 
-        # Prevent duplicate processing (e.g., during weight reload)
-        layer._already_called_process_weights_after_loading = True
-
     def maybe_make_prepare_finalize(
         self,
         routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
@@ -1011,92 +946,16 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        layer.intermediate_size_per_partition = intermediate_size_per_partition
-        layer.hidden_size = hidden_size
         layer.num_experts = num_experts
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None
 
-        # We are doing online quantization, patch the weight loaded
-        # to call `process_weights_after_loading` in a streaming fashion
-        # as soon as the last weight chunk is loaded.
-        weight_loader = extra_weight_attrs["weight_loader"]
-        # create a new holder to prevent modifying behavior of any other
-        # objects which might depend on the old one
-        new_extra_weight_attrs = extra_weight_attrs
-
-        def patched_weight_loader(param, loaded_weight, *args, **kwargs):
-            # add a counter to track how many elements we have updated
-            if not hasattr(layer, "_loaded_numel"):
-                layer._loaded_numel = 0
-
-                # save the ids of original w13 and w2 so that we can
-                # distinguish which one `param` should map to further
-                # down in this file
-                layer._w13_weight_orig_id = id(layer.w13_weight)
-                layer._w2_weight_orig_id = id(layer.w2_weight)
-
-                # when the first `loaded_weight` is about to be
-                # loaded to `param`, materialize `param` just-in-time
-
-                w13_weight = torch.nn.Parameter(
-                    torch.empty_like(layer.w13_weight, device=layer._load_device),
-                    requires_grad=False,
-                )
-                set_weight_attrs(w13_weight, extra_weight_attrs)
-                _copy_missing_attrs(layer.w13_weight, w13_weight)
-                layer.register_parameter("w13_weight", w13_weight)
-
-                w2_weight = torch.nn.Parameter(
-                    torch.empty_like(layer.w2_weight, device=layer._load_device),
-                    requires_grad=False,
-                )
-                set_weight_attrs(w2_weight, extra_weight_attrs)
-                _copy_missing_attrs(layer.w2_weight, w2_weight)
-                layer.register_parameter("w2_weight", w2_weight)
-                del layer._load_device
-
-            # refresh the reference to `param` to reflect just-in-time
-            # materialization
-            if id(param) == layer._w13_weight_orig_id:
-                param = layer.w13_weight
-            elif id(param) == layer._w2_weight_orig_id:
-                param = layer.w2_weight
-
-            # load the current weight chunk
-            copy_numel_counter = CopyNumelCounter()
-            with copy_numel_counter:
-                res = weight_loader(param, loaded_weight, *args, **kwargs)  # type: ignore[misc]
-            layer._loaded_numel += copy_numel_counter.copied_numel
-
-            # if we have loaded all of the elements, call
-            # process_weights_after_loading
-            target_loaded_numel = layer.w13_weight.numel() + layer.w2_weight.numel()
-            if layer._loaded_numel == target_loaded_numel:
-                self.process_weights_after_loading(layer)
-
-                # Prevent the usual `process_weights_after_loading` call
-                # from doing anything
-                layer._already_called_process_weights_after_loading = True
-
-                # Note that we keep `layer._loaded_numel`,
-                # `layer._w13_weight_orig_id` and `layer._w2_weight_orig_id`
-                # around because if EP is on, weight loaders for non-local
-                # experts will run but not actually copy any elements, and we
-                # need to not re-initialize in that case.
-
-            return res
-
-        new_extra_weight_attrs["weight_loader"] = patched_weight_loader
-        extra_weight_attrs = new_extra_weight_attrs
-
         # WEIGHTS
         w13_weight = torch.nn.Parameter(
             torch.empty(
                 num_experts,
                 2 * intermediate_size_per_partition,
                 hidden_size,
-                # materialized just-in-time in `patched_weight_loader`
                 device="meta",
                 dtype=params_dtype,
             ),
@@ -1110,91 +969,53 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
                 num_experts,
                 hidden_size,
                 intermediate_size_per_partition,
-                # materialized just-in-time in `patched_weight_loader`
-                device="meta",
+                device="meta",  # materialized and processed during loading
                 dtype=params_dtype,
             ),
             requires_grad=False,
         )
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
-        # stash the correct device for `patched_weight_loader`
-        layer._load_device = torch.get_default_device()
 
         # BIASES (for models like GPT-OSS that have biased MoE)
         if self.moe.has_bias:
-            # Use the original weight_loader (not patched) for biases
-            orig_extra_weight_attrs = dict(extra_weight_attrs)
-            orig_extra_weight_attrs["weight_loader"] = weight_loader
             w13_bias = torch.nn.Parameter(
                 torch.zeros(
                     num_experts,
                     2 * intermediate_size_per_partition,
+                    device="meta",  # materialized and processed during loading
                     dtype=layer.orig_dtype,
                 ),
                 requires_grad=False,
             )
             layer.register_parameter("w13_bias", w13_bias)
-            set_weight_attrs(w13_bias, orig_extra_weight_attrs)
+            set_weight_attrs(w13_bias, extra_weight_attrs)
+
             w2_bias = torch.nn.Parameter(
-                torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype),
+                torch.zeros(
+                    num_experts,
+                    hidden_size,
+                    device="meta",  # materialized and processed during loading
+                    dtype=layer.orig_dtype,
+                ),
                 requires_grad=False,
             )
             layer.register_parameter("w2_bias", w2_bias)
-            set_weight_attrs(w2_bias, orig_extra_weight_attrs)
+            set_weight_attrs(w2_bias, extra_weight_attrs)
 
-        # WEIGHT_SCALES
-        # Allocate 2 scales for w1 and w3 respectively.
-        # They will be combined to a single scale after weight loading.
-        w13_weight_scale = torch.nn.Parameter(
-            torch.ones(num_experts, dtype=torch.float32), requires_grad=False
-        )
-        w2_weight_scale = torch.nn.Parameter(
-            torch.ones(num_experts, dtype=torch.float32), requires_grad=False
-        )
-        layer.register_parameter("w13_weight_scale", w13_weight_scale)
-        layer.register_parameter("w2_weight_scale", w2_weight_scale)
-        set_weight_attrs(w13_weight_scale, extra_weight_attrs)
-        set_weight_attrs(w2_weight_scale, extra_weight_attrs)
-
-        layer.w13_input_scale = None
-        layer.w2_input_scale = None
+        initialize_online_processing(layer)
 
     def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return
 
-        # deferred initialization of randomly initialized weights for the
-        # `--load_format dummy` feature
-        if layer.w13_weight.device == torch.device("meta"):
-            w13_weight = torch.nn.Parameter(
-                torch.empty_like(layer.w13_weight, device=layer._load_device),
-                requires_grad=False,
-            )
-            set_weight_attrs(
-                w13_weight, {"weight_loader": layer.w13_weight.weight_loader}
-            )
-            _copy_missing_attrs(layer.w13_weight, w13_weight)
-            layer.register_parameter("w13_weight", w13_weight)
-            initialize_single_dummy_weight(layer.w13_weight)
-        if layer.w2_weight.device == torch.device("meta"):
-            w2_weight = torch.nn.Parameter(
-                torch.empty_like(layer.w2_weight, device=layer._load_device),
-                requires_grad=False,
-            )
-            set_weight_attrs(
-                w2_weight, {"weight_loader": layer.w2_weight.weight_loader}
-            )
-            _copy_missing_attrs(layer.w2_weight, w2_weight)
-            layer.register_parameter("w2_weight", w2_weight)
-            initialize_single_dummy_weight(layer.w2_weight)
-
-        # If checkpoint is fp16, quantize in place.
         fp8_dtype = current_platform.fp8_dtype()
         w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
         w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
-        w13_scale = layer.w13_weight_scale
-        w2_scale = layer.w2_weight_scale
+        w13_scale = torch.ones(layer.num_experts, dtype=torch.float32)
+        w2_scale = torch.ones(layer.num_experts, dtype=torch.float32)
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
 
         for expert in range(layer.local_num_experts):
             w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
@@ -1211,8 +1032,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
             w2,
             w13_scale,
             w2_scale,
-            layer.w13_input_scale,
-            layer.w2_input_scale,
+            w13_input_scale=layer.w13_input_scale,
+            w2_input_scale=layer.w2_input_scale,
         )
 
         # Prevent duplicate processing (e.g., during weight reload)

vllm/model_executor/layers/quantization/mxfp4.py

@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe import (
 )
 from vllm.model_executor.layers.fused_moe import modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEParallelConfig,
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
@@ -107,18 +108,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
         self.moe_kernel: mk.FusedMoEKernel | None = None
 
-        # Round up dims once based on backend. This mutates the shared
-        # FusedMoEConfig in-place so that create_weights() and all
-        # downstream code see the padded dimensions. This must happen
-        # before create_weights() is called.
-        self.moe.hidden_dim, self.moe.intermediate_size_per_partition = (
-            mxfp4_round_up_hidden_size_and_intermediate_size(
-                self.mxfp4_backend,
-                self.moe.hidden_dim,
-                self.moe.intermediate_size_per_partition,
-            )
-        )
-
         # Used for triton kernel precision configs
         self.w13_precision_config = None
         self.w2_precision_config = None
@@ -129,6 +118,23 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         # so can skip the padding in the forward before applying the moe method
         return self.mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
 
+    def maybe_roundup_sizes(
+        self,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        act_dtype: torch.dtype,
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> tuple[int, int]:
+        hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
+            hidden_size=hidden_size,
+            intermediate_size_per_partition=intermediate_size_per_partition,
+            act_dtype=act_dtype,
+            moe_parallel_config=moe_parallel_config,
+        )
+        return mxfp4_round_up_hidden_size_and_intermediate_size(
+            self.mxfp4_backend, hidden_size, intermediate_size_per_partition
+        )
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -143,32 +149,16 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         scale_dtype = torch.uint8
         mxfp4_block = 32
 
-        # Use pre-rounded sizes from config
-        self.intermediate_size = intermediate_size_per_partition_after_pad = (
-            self.moe.intermediate_size_per_partition
-        )
-        self.hidden_size = hidden_size = self.moe.hidden_dim
-
-        # Expose padded dimensions on the layer for LoRA and Marlin code
-        # that reads layer.hidden_size / layer.intermediate_size_per_partition.
         layer.params_dtype = params_dtype
         layer.num_experts = num_experts
-        layer.hidden_size = hidden_size
-        layer.intermediate_size_per_partition = (
-            intermediate_size_per_partition_after_pad
-        )
-
-        # CK (gfx950) padding info for rocm_aiter_ops.fused_moe()
-        self.hidden_pad = extra_weight_attrs.get("hidden_pad", 0)
-        self.intermediate_pad = (
-            intermediate_size_per_partition_after_pad - intermediate_size_per_partition
-        )
+        self.intermediate_size = intermediate_size_per_partition
+        self.hidden_size = hidden_size
 
         # Fused gate_up_proj (column parallel)
         w13_weight = torch.nn.Parameter(
             torch.zeros(
                 num_experts,
-                2 * intermediate_size_per_partition_after_pad,
+                2 * intermediate_size_per_partition,
                 hidden_size // 2,
                 dtype=weight_dtype,
             ),
@@ -180,7 +170,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         w13_weight_scale = torch.nn.Parameter(
             torch.zeros(
                 num_experts,
-                2 * intermediate_size_per_partition_after_pad,
+                2 * intermediate_size_per_partition,
                 hidden_size // mxfp4_block,
                 dtype=scale_dtype,
             ),
@@ -194,7 +184,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             torch.zeros(
                 num_experts,
                 hidden_size,
-                intermediate_size_per_partition_after_pad // 2,
+                intermediate_size_per_partition // 2,
                 dtype=weight_dtype,
             ),
             requires_grad=False,
@@ -206,7 +196,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             torch.zeros(
                 num_experts,
                 hidden_size,
-                intermediate_size_per_partition_after_pad // mxfp4_block,
+                intermediate_size_per_partition // mxfp4_block,
                 dtype=scale_dtype,
             ),
             requires_grad=False,
@@ -218,7 +208,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             w13_bias = torch.nn.Parameter(
                 torch.zeros(
                     num_experts,
-                    2 * intermediate_size_per_partition_after_pad,
+                    2 * intermediate_size_per_partition,
                     dtype=torch.bfloat16,
                 ),
                 requires_grad=False,
@@ -368,8 +358,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             w2_scale=w2_scale,
             w1_bias=w1_bias,
             w2_bias=w2_bias,
-            hidden_pad=self.hidden_pad,
-            intermediate_pad=self.intermediate_pad,
         )
 
     def select_gemm_impl(

vllm/model_executor/layers/quantization/mxfp8.py

@@ -337,6 +337,8 @@ class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):
         w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
         w13_scale = layer.w13_weight_scale
         w2_scale = layer.w2_weight_scale
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
 
         w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight)
         w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight)

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

@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe import (
     MoEActivation,
 )
 from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEParallelConfig,
     FusedMoEQuantConfig,
     fp8_w8a8_moe_quant_config,
     mxfp4_w4a8_moe_quant_config,
@@ -27,13 +28,13 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
     Mxfp4MoeBackend,
+    mxfp4_round_up_hidden_size_and_intermediate_size,
     select_mxfp4_moe_backend,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
     prepare_fp8_moe_layer_for_marlin,
 )
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
-    CK_MXFP4_MOE_DIM_ALIGNMENT,
     _swizzle_mxfp4,
 )
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -49,7 +50,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
-from vllm.utils.math_utils import round_up
 
 logger = init_logger(__name__)
 
@@ -173,8 +173,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        layer.intermediate_size_per_partition = intermediate_size_per_partition
-        layer.hidden_size = hidden_size
         layer.num_experts = num_experts
         layer.orig_dtype = params_dtype
         layer.weight_block_size = None
@@ -182,7 +180,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
 
         # WEIGHTS
         w13_weight = torch.nn.Parameter(
-            torch.empty(
+            torch.zeros(
                 num_experts,
                 2 * intermediate_size_per_partition,
                 hidden_size,
@@ -194,7 +192,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
         set_weight_attrs(w13_weight, extra_weight_attrs)
 
         w2_weight = torch.nn.Parameter(
-            torch.empty(
+            torch.zeros(
                 num_experts,
                 hidden_size,
                 intermediate_size_per_partition,
@@ -461,6 +459,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
                 activation=layer.activation,
                 apply_router_weight_on_input=layer.apply_router_weight_on_input,
                 quant_config=self.moe_quant_config,
+                moe_config=layer.moe_config,
                 expert_map=layer.expert_map,
             )
         elif self.use_marlin:
@@ -527,7 +526,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
     ):
         params_dtype = torch.uint32
         w13_weight = torch.nn.Parameter(
-            torch.empty(
+            torch.zeros(
                 num_experts,
                 2 * intermediate_size_per_partition,
                 hidden_size // 8,  # INT32 packing for W4
@@ -536,7 +535,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
             requires_grad=False,
         )
         w2_weight = torch.nn.Parameter(
-            torch.empty(
+            torch.zeros(
                 num_experts,
                 hidden_size,
                 intermediate_size_per_partition // 8,  # INT32 packing for W4
@@ -649,6 +648,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
             activation=layer.activation,
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
             quant_config=self.moe_quant_config,
+            moe_config=layer.moe_config,
             expert_map=layer.expert_map,
         )
 
@@ -702,6 +702,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
         self.mxfp4_backend: Mxfp4MoeBackend | None = None
         if self.ocp_mx_scheme == "w_mxfp4":
             self.mxfp4_backend, _ = select_mxfp4_moe_backend(moe)
+        elif self.ocp_mx_scheme.startswith("w_mxfp4"):
+            # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes.
+            self.mxfp4_backend = Mxfp4MoeBackend.NONE
 
         if self.input_quant is not None:
             self.static_input_scales = not self.input_quant.get("is_dynamic")
@@ -734,36 +737,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
         self.emulate = (
             not current_platform.supports_mx()
             or not self.ocp_mx_scheme.startswith("w_mxfp4")
-        ) and (self.mxfp4_backend is None or not self.use_rocm_aiter_moe)
-
-        # CK's pre-compiled MXFP4 MoE GEMM kernel instances have dimension
-        # alignment requirements. When violated (e.g. MiniMax-M2.1 with
-        # TP=4 yields intermediate_size_per_partition=384), AITER raises:
-        # "device_gemm ... does not support this GEMM problem".
-        # Fall back to emulation in that case.
-        # For gpt_oss models, create_weights rounds up the dimensions
-        # internally, so the alignment check is skipped.
-        if (
-            not self.emulate
-            and self.use_rocm_aiter_moe
-            and self.ocp_mx_scheme is not None
-            and self.ocp_mx_scheme.startswith("w_mxfp4")
-            and self.model_type != "gpt_oss"
-            and moe.intermediate_size_per_partition % CK_MXFP4_MOE_DIM_ALIGNMENT != 0
-        ):
-            logger.warning_once(
-                "AITER CK MXFP4 MoE GEMM does not support "
-                "intermediate_size_per_partition=%d (not a multiple of %d). "
-                "This typically happens when intermediate_size / "
-                "tensor_parallel_size produces an incompatible dimension. "
-                "Falling back to emulation mode. To avoid this overhead, "
-                "use a compatible tensor_parallel_size or set "
-                "VLLM_ROCM_USE_AITER_MOE=0.",
-                moe.intermediate_size_per_partition,
-                CK_MXFP4_MOE_DIM_ALIGNMENT,
-            )
-            self.use_rocm_aiter_moe = False
-            self.emulate = True
+        ) and (
+            self.mxfp4_backend is None
+            or self.mxfp4_backend is Mxfp4MoeBackend.NONE
+            or not self.use_rocm_aiter_moe
+        )
 
         if self.emulate:
             logger.warning_once(
@@ -780,6 +758,27 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
                 "The current mode supports native MoE MXFP4 computation"
             )
 
+    def maybe_roundup_sizes(
+        self,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        act_dtype: torch.dtype,
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> tuple[int, int]:
+        hidden_size, intermediate_size_per_partition = super().maybe_roundup_sizes(
+            hidden_size=hidden_size,
+            intermediate_size_per_partition=intermediate_size_per_partition,
+            act_dtype=act_dtype,
+            moe_parallel_config=moe_parallel_config,
+        )
+        if self.mxfp4_backend is not None:
+            hidden_size, intermediate_size_per_partition = (
+                mxfp4_round_up_hidden_size_and_intermediate_size(
+                    self.mxfp4_backend, hidden_size, intermediate_size_per_partition
+                )
+            )
+        return hidden_size, intermediate_size_per_partition
+
     def get_packed_dim(self, dim: int, quant_dtype: str):
         if quant_dtype == "mxfp4":
             assert dim % 2 == 0
@@ -805,40 +804,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
         )
 
         params_dtype = torch.uint8
-        self.intermediate_size_per_partition = intermediate_size_per_partition
-        if self.model_type == "gpt_oss":
-            if current_platform.is_rocm():
-                intermediate_size_per_partition_after_pad = round_up(
-                    intermediate_size_per_partition, 256
-                )
-            else:
-                intermediate_size_per_partition_after_pad = round_up(
-                    intermediate_size_per_partition, 64
-                )
-        else:
-            intermediate_size_per_partition_after_pad = intermediate_size_per_partition
-
-        self.unpadded_hidden_size = extra_weight_attrs.get(
-            "unpadded_hidden_size", hidden_size
-        )
-
-        # On GFX950, the GFX950MXScaleLayout swizzle requires
-        # hidden_size to be a multiple of 256 (SCALE_K = hidden_size / 32
-        # must be divisible by 8). Pad hidden_size for weight/scale
-        # allocation; the original value is preserved in unpadded_hidden_size.
-        # Only applies to the native (non-emulated) CK path on GFX950.
-        if (
-            self.model_type == "gpt_oss"
-            and current_platform.is_rocm()
-            and not self.emulate
-        ):
-            hidden_size = round_up(hidden_size, 256)
 
         # WEIGHTS
         w13_weight = torch.nn.Parameter(
-            torch.empty(
+            torch.zeros(
                 num_experts,
-                2 * intermediate_size_per_partition_after_pad,
+                2 * intermediate_size_per_partition,
                 self.get_packed_dim(hidden_size, self.weight_dtype),
                 dtype=params_dtype,
             ),
@@ -849,12 +820,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
         set_weight_attrs(w13_weight, extra_weight_attrs)
 
         w2_weight = torch.nn.Parameter(
-            torch.empty(
+            torch.zeros(
                 num_experts,
                 hidden_size,
-                self.get_packed_dim(
-                    intermediate_size_per_partition_after_pad, self.weight_dtype
-                ),
+                self.get_packed_dim(intermediate_size_per_partition, self.weight_dtype),
                 dtype=params_dtype,
             ),
             requires_grad=False,
@@ -867,7 +836,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
         w13_weight_scale = torch.nn.Parameter(
             torch.ones(
                 num_experts,
-                2 * intermediate_size_per_partition_after_pad,
+                2 * intermediate_size_per_partition,
                 hidden_size // OCP_MX_BLOCK_SIZE,
                 dtype=params_dtype,
             ),
@@ -877,7 +846,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
             torch.ones(
                 num_experts,
                 hidden_size,
-                intermediate_size_per_partition_after_pad // OCP_MX_BLOCK_SIZE,
+                intermediate_size_per_partition // OCP_MX_BLOCK_SIZE,
                 dtype=params_dtype,
             ),
             requires_grad=False,
@@ -892,7 +861,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
             w13_bias = torch.nn.Parameter(
                 torch.zeros(
                     num_experts,
-                    2 * intermediate_size_per_partition_after_pad,
+                    2 * intermediate_size_per_partition,
                     dtype=torch.float32,
                 ),
                 requires_grad=False,
@@ -1072,6 +1041,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
                 topk_ids=topk_ids,
                 activation=layer.activation,
                 quant_config=self.moe_quant_config,
+                moe_config=layer.moe_config,
                 expert_map=layer.expert_map,
             )
         else:
@@ -1204,6 +1174,8 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
             triton_kernel_moe_forward,
         )
 
+        assert self.moe.hidden_dim_unpadded is not None
+        assert self.moe.intermediate_size_per_partition_unpadded is not None
         return triton_kernel_moe_forward(
             hidden_states=x,
             w1=self.w13_weight_triton_tensor,
@@ -1215,8 +1187,8 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
             expert_map=expert_map,
             quant_config=self.moe_quant_config,
             apply_router_weight_on_input=layer.apply_router_weight_on_input,
-            unpadded_N_w1=self.intermediate_size_per_partition * 2,
-            unpadded_K_w1=self.unpadded_hidden_size,
-            unpadded_N_w2=self.unpadded_hidden_size,
-            unpadded_K_w2=self.intermediate_size_per_partition,
+            unpadded_N_w1=self.moe.intermediate_size_per_partition_unpadded * 2,
+            unpadded_K_w1=self.moe.hidden_dim_unpadded,
+            unpadded_N_w2=self.moe.hidden_dim_unpadded,
+            unpadded_K_w2=self.moe.intermediate_size_per_partition_unpadded,
         )

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

@@ -254,7 +254,6 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
                 w13, w13_scale, w2, w2_scale, is_act_and_mul, min_alignment
             )
         )
-        layer.intermediate_size_per_partition = padded_intermediate
         layer.moe_config.intermediate_size_per_partition = padded_intermediate
 
         w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe(

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

@@ -439,7 +439,6 @@ def prepare_fp8_moe_layer_for_fi(
             layer.moe_config.is_act_and_mul,
             min_alignment,
         )
-        layer.intermediate_size_per_partition = new_intermediate
         layer.moe_config.intermediate_size_per_partition = new_intermediate
 
     # FI kernels require W31 layout rather than W13.

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

@@ -27,10 +27,19 @@ def is_fp4_marlin_supported():
     return current_platform.has_device_capability(75)
 
 
-def _nvfp4_compute_scale_factor(marlin_scales: torch.Tensor) -> float:
+def _nvfp4_compute_scale_factor(
+    marlin_scales: torch.Tensor,
+    a_dtype: torch.dtype | None = None,
+) -> float:
     """Compute the power-of-2 scale_factor needed so that all non-zero
     values in marlin_scales * 2^7 are >= 2 after rescaling.
     Returns a Python float (power of 2, >= 1.0)."""
+
+    # Since half has a smaller dynamic range compared to bfloat16,
+    # no rescaling is applied here if active dtype is half.
+    if a_dtype is not None and a_dtype == torch.half:
+        return 1.0
+
     ws_float = marlin_scales.float() * (2**7)
     nonzero_mask = ws_float > 0
     if nonzero_mask.any():
@@ -44,6 +53,7 @@ def _nvfp4_compute_scale_factor(marlin_scales: torch.Tensor) -> float:
 def nvfp4_marlin_process_scales(
     marlin_scales: torch.Tensor,
     scale_factor: float | None = None,
+    a_dtype: torch.dtype | None = None,
 ) -> tuple[torch.Tensor, float]:
     """Process NVFP4 weight scales into the special S0E5M3 format for Marlin.
 
@@ -91,7 +101,7 @@ def nvfp4_marlin_process_scales(
     # to fully utilize the E4M3 dynamic range (e.g., global_scale=1).
     # The caller must compensate by dividing global_scale by scale_factor.
     if scale_factor is None:
-        scale_factor = _nvfp4_compute_scale_factor(marlin_scales)
+        scale_factor = _nvfp4_compute_scale_factor(marlin_scales, a_dtype)
     if scale_factor > 1.0:
         marlin_scales = (marlin_scales.float() * scale_factor).to(torch.half)
 
@@ -119,12 +129,14 @@ def mxfp4_marlin_process_scales(marlin_scales, input_dtype=None):
     return marlin_scales
 
 
-def nvfp4_marlin_process_global_scale(global_scale):
-    assert global_scale.dtype in [torch.half, torch.bfloat16]
+def nvfp4_marlin_process_global_scale(global_scale, a_dtype: torch.dtype | None = None):
+    if a_dtype is None:
+        a_dtype = global_scale.dtype
+    assert a_dtype in [torch.half, torch.bfloat16]
     fp4_exponent = 2
-    if global_scale.dtype == torch.half:
+    if a_dtype == torch.half:
         target_exponent = 5
-    elif global_scale.dtype == torch.bfloat16:
+    elif a_dtype == torch.bfloat16:
         target_exponent = 8
     # exponent_bias_fp16 = 2 ** 4 - 2 ** 1 = 14
     # exponent_bias_bf16 = 2 ** 7 - 2 ** 1 = 126
@@ -244,11 +256,15 @@ def prepare_fp4_layer_for_marlin(
     )
 
     if is_nvfp4:
-        weight_scale, scale_factor = nvfp4_marlin_process_scales(weight_scale)
+        weight_scale, scale_factor = nvfp4_marlin_process_scales(
+            weight_scale, a_dtype=param_dtype
+        )
         layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
 
-        weight_global_scale = layer.weight_global_scale.to(param_dtype)
-        weight_global_scale = nvfp4_marlin_process_global_scale(weight_global_scale)
+        weight_global_scale = layer.weight_global_scale.to(torch.float32)
+        weight_global_scale = nvfp4_marlin_process_global_scale(
+            weight_global_scale, param_dtype
+        )
         weight_global_scale = weight_global_scale / scale_factor
         layer.weight_global_scale = torch.nn.Parameter(
             weight_global_scale, requires_grad=False
@@ -339,7 +355,6 @@ def prepare_nvfp4_moe_layer_for_marlin(
         scales: torch.Tensor, g_scales: torch.Tensor, name: str
     ) -> tuple[torch.Tensor, torch.Tensor]:
         scales = scales.to(param_dtype)
-        g_scales = g_scales.to(param_dtype)
 
         tensor_list = []
         num_shards = 2 if is_act_and_mul else 1
@@ -350,7 +365,7 @@ def prepare_nvfp4_moe_layer_for_marlin(
 
         # All experts share one global_scale, so compute the max
         # scale_factor across all experts first, then apply uniformly.
-        combined_scale_factor = _nvfp4_compute_scale_factor(scales)
+        combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype)
 
         for i in range(E):
             scale = scales[i].T
@@ -362,12 +377,12 @@ def prepare_nvfp4_moe_layer_for_marlin(
                 is_a_8bit=is_a_8bit,
             )
             marlin_scales, _ = nvfp4_marlin_process_scales(
-                marlin_scales, scale_factor=combined_scale_factor
+                marlin_scales, scale_factor=combined_scale_factor, a_dtype=param_dtype
             )
             tensor_list.append(marlin_scales)
 
         scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
-        g_scales = nvfp4_marlin_process_global_scale(g_scales)
+        g_scales = nvfp4_marlin_process_global_scale(g_scales, param_dtype)
         g_scales = g_scales / combined_scale_factor
         return scales, g_scales
 
@@ -438,7 +453,7 @@ def prepare_moe_fp4_layer_for_marlin(
             scales = scales.view(torch.float8_e8m0fnu)
         scales = scales.to(param_dtype)
         if is_nvfp4:
-            global_scale = getattr(layer, name + "_weight_scale_2").to(param_dtype)
+            global_scale = getattr(layer, name + "_weight_scale_2")
 
         tensor_list = []
         if "w13" in name:
@@ -449,7 +464,7 @@ def prepare_moe_fp4_layer_for_marlin(
         # For NVFP4: compute unified scale_factor across all experts
         combined_scale_factor = None
         if is_nvfp4:
-            combined_scale_factor = _nvfp4_compute_scale_factor(scales)
+            combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype)
 
         for i in range(e):
             scale = scales[i].T
@@ -463,7 +478,9 @@ def prepare_moe_fp4_layer_for_marlin(
             )
             if is_nvfp4:
                 marlin_scales, _ = nvfp4_marlin_process_scales(
-                    marlin_scales, scale_factor=combined_scale_factor
+                    marlin_scales,
+                    scale_factor=combined_scale_factor,
+                    a_dtype=param_dtype,
                 )
             else:
                 marlin_scales = mxfp4_marlin_process_scales(
@@ -477,7 +494,7 @@ def prepare_moe_fp4_layer_for_marlin(
 
         if is_nvfp4:
             assert combined_scale_factor is not None
-            global_scale = nvfp4_marlin_process_global_scale(global_scale)
+            global_scale = nvfp4_marlin_process_global_scale(global_scale, param_dtype)
             global_scale = global_scale / combined_scale_factor
             global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
             setattr(layer, name + "_weight_scale_2", global_scale)
@@ -665,7 +682,7 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
     )
     marlin_scales, scale_factor = nvfp4_marlin_process_scales(marlin_scales)
 
-    global_scale = nvfp4_marlin_process_global_scale(global_scale)
+    global_scale = nvfp4_marlin_process_global_scale(global_scale).to(torch.float32)
     global_scale = global_scale / scale_factor
 
     return weight_ref.T, marlin_qweight, marlin_scales, global_scale

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

@@ -6,7 +6,6 @@ import torch
 
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from vllm.triton_utils import triton
 from vllm.utils.import_utils import has_triton_kernels
 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
 
@@ -49,9 +48,16 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
 
         value_layout = StridedLayout
         if on_gfx950():
-            from triton_kernels.tensor_details.layout import GFX950MXScaleLayout
+            try:
+                # triton < 3.6
+                from triton_kernels.tensor_details.layout import GFX950MXScaleLayout
 
-            scale_layout = GFX950MXScaleLayout
+                scale_layout = GFX950MXScaleLayout
+            except ImportError:
+                # triton >= 3.6
+                from triton_kernels.tensor_details.layout import CDNA4MXScaleLayout
+
+                scale_layout = CDNA4MXScaleLayout
         else:
             scale_layout = StridedLayout
     else:
@@ -85,14 +91,6 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
     return quant_tensor, InFlexData(), scale
 
 
-def get_padding_alignment():
-    return (
-        256
-        if triton.runtime.driver.active.get_current_target().arch in ("gfx950",)
-        else 128
-    )
-
-
 def _dequant_mxfp4(
     x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
 ) -> torch.Tensor:

vllm/model_executor/model_loader/base_loader.py

@@ -9,6 +9,7 @@ import vllm.envs as envs
 from vllm.config import ModelConfig, VllmConfig
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.model_loader.reload import finalize_layerwise_processing
 from vllm.model_executor.model_loader.utils import (
     initialize_model,
     process_weights_after_loading,
@@ -49,16 +50,13 @@ class BaseModelLoader(ABC):
             device_config.device if load_config.device is None else load_config.device
         )
         target_device = torch.device(load_device)
-        with set_default_torch_dtype(model_config.dtype):
-            with target_device:
-                model = initialize_model(
-                    vllm_config=vllm_config, model_config=model_config, prefix=prefix
-                )
-
+        with set_default_torch_dtype(model_config.dtype), target_device:
+            model = initialize_model(
+                vllm_config=vllm_config, model_config=model_config, prefix=prefix
+            )
             log_model_inspection(model)
 
             logger.debug("Loading weights on %s ...", load_device)
-            # Quantization does not happen in `load_weights` but after it
             self.load_weights(model, model_config)
 
             # Log peak GPU memory after loading weights. This is needed
@@ -71,6 +69,11 @@ class BaseModelLoader(ABC):
                     scope="local",
                 )
 
+            # Process weights into kernel format. Note that when using online
+            # quantization, weights are (typically) quantized as they are loaded.
+            if _has_online_quant(model):
+                finalize_layerwise_processing(model, model_config)
+
             process_weights_after_loading(model, model_config, target_device)
 
         return model.eval()
@@ -84,3 +87,12 @@ def log_model_inspection(model: nn.Module) -> None:
     from vllm.model_inspection import format_model_inspection
 
     logger.info("vLLM model structure:\n%s", format_model_inspection(model))
+
+
+def _has_online_quant(model: nn.Module):
+    for module in model.modules():
+        quant_method = getattr(module, "quant_method", None)
+        if getattr(quant_method, "uses_meta_device", False):
+            return True
+
+    return False

vllm/model_executor/model_loader/dummy_loader.py

@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
 import torch.nn as nn
 
 from vllm.config import ModelConfig
 from vllm.config.load import LoadConfig
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+from vllm.model_executor.model_loader.reload.meta import materialize_meta_tensor
+from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
 from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights
 
 
@@ -23,6 +26,12 @@ class DummyModelLoader(BaseModelLoader):
         pass  # Nothing to download
 
     def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
+        # materialize meta tensors as part of online quantization lifecycle
+        for layer in model.modules():
+            for name, param in get_layer_tensors(layer).items():
+                if param.device == torch.device("meta"):
+                    setattr(layer, name, materialize_meta_tensor(param))
+
         # NOTE(woosuk): For accurate performance evaluation, we assign
         # random values to the weights.
         initialize_dummy_weights(model, model_config)