[Bugfix] Fix DeepSeek v3.2 tokenizer outputting None issue (#33832)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

.buildkite/hardware_tests/arm.yaml

@@ -1,8 +0,0 @@
-group: Hardware
-steps:
-  - label: "Arm CPU Test"
-    soft_fail: true
-    device: arm_cpu
-    no_plugin: true
-    commands: 
-    - bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

.buildkite/hardware_tests/cpu.yaml

@@ -0,0 +1,100 @@
+group: CPU
+depends_on: []
+steps:
+- label: CPU-Kernel Tests
+  depends_on: []
+  soft_fail: true
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - csrc/cpu/
+  - cmake/cpu_extension.cmake
+  - CMakeLists.txt
+  - vllm/_custom_ops.py
+  - tests/kernels/attention/test_cpu_attn.py
+  - tests/kernels/moe/test_cpu_fused_moe.py
+  - tests/kernels/test_onednn.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
+      pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
+      pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
+      pytest -x -v -s tests/kernels/test_onednn.py"
+
+- label: CPU-Language Generation and Pooling Model Tests
+  depends_on: []
+  soft_fail: true
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - csrc/cpu/
+  - vllm/
+  - tests/models/language/generation/
+  - tests/models/language/pooling/
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      pytest -x -v -s tests/models/language/generation -m cpu_model
+      pytest -x -v -s tests/models/language/pooling -m cpu_model"
+
+- label: CPU-Quantization Model Tests
+  depends_on: []
+  soft_fail: true
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - csrc/cpu/
+  - vllm/model_executor/layers/quantization/cpu_wna16.py
+  - vllm/model_executor/layers/quantization/gptq_marlin.py
+  - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
+  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
+  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
+  - tests/quantization/test_compressed_tensors.py
+  - tests/quantization/test_cpu_wna16.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
+      pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
+      pytest -x -v -s tests/quantization/test_cpu_wna16.py"
+      
+- label: CPU-Distributed Tests
+  depends_on: []
+  soft_fail: true
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  - csrc/cpu/shm.cpp
+  - vllm/v1/worker/cpu_worker.py
+  - vllm/v1/worker/gpu_worker.py
+  - vllm/v1/worker/cpu_model_runner.py
+  - vllm/v1/worker/gpu_model_runner.py
+  - vllm/platforms/cpu.py
+  - vllm/distributed/parallel_state.py
+  - vllm/distributed/device_communicators/cpu_communicator.py
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"
+
+- label: CPU-Multi-Modal Model Tests %N
+  depends_on: []
+  soft_fail: true
+  device: intel_cpu
+  no_plugin: true
+  source_file_dependencies:
+  # - vllm/
+  - vllm/model_executor/layers/rotary_embedding
+  - tests/models/multimodal/generation/
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
+      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
+  parallelism: 2
+
+- label: "Arm CPU Test"
+  depends_on: []
+  soft_fail: true
+  device: arm_cpu
+  no_plugin: true
+  commands: 
+  - bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh

.buildkite/hardware_tests/intel.yaml

@@ -1,13 +1,6 @@
 group: Hardware
 depends_on: ~
 steps:
-  - label: "Intel CPU Test"
-    soft_fail: true
-    device: intel_cpu
-    no_plugin: true
-    commands: 
-    - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh
-  
   - label: "Intel HPU Test"
     soft_fail: true
     device: intel_hpu

.buildkite/scripts/annotate-release.sh

@@ -27,7 +27,7 @@ aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-
 To download and upload the image:
 
 \`\`\`
-Download images:
+# Download images:
 
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
@@ -35,8 +35,12 @@ docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
+docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
+docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
 
-Tag and push images:
+# Tag and push images:
+
+## CUDA
 
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
 docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
@@ -62,34 +66,21 @@ docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-a
 docker push vllm/vllm-openai:latest-aarch64-cu130
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
 
-docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:latest
-docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm
+## ROCm
+
+docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
+docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
+docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
 docker push vllm/vllm-openai-rocm:latest
-docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm
+docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
 
-Create multi-arch manifest:
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
 docker push vllm/vllm-openai-rocm:latest-base
 docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
 
-docker manifest rm vllm/vllm-openai:latest
-docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
-docker manifest push vllm/vllm-openai:latest
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
-
-docker manifest rm vllm/vllm-openai:latest-cu130
-docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
-docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
-docker manifest push vllm/vllm-openai:latest-cu130
-docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
-
-# CPU images (vllm/vllm-openai-cpu)
-docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
-docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
+## CPU
 
 docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
 docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
@@ -103,6 +94,20 @@ docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-a
 docker push vllm/vllm-openai-cpu:latest-arm64
 docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
 
+# Create multi-arch manifest:
+
+docker manifest rm vllm/vllm-openai:latest
+docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
+docker manifest push vllm/vllm-openai:latest
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
+
+docker manifest rm vllm/vllm-openai:latest-cu130
+docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
+docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
+docker manifest push vllm/vllm-openai:latest-cu130
+docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
+
 docker manifest rm vllm/vllm-openai-cpu:latest || true
 docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
 docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64

.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+set -euox pipefail
+
+echo "--- PP+TP"
+vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
+server_pid=$!
+timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
+vllm bench serve \
+    --backend vllm \
+    --dataset-name random \
+    --model meta-llama/Llama-3.2-3B-Instruct \
+    --num-prompts 20 \
+    --endpoint /v1/completions
+kill -s SIGTERM $server_pid &
+
+echo "--- DP+TP"
+vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
+server_pid=$!
+timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
+vllm bench serve \
+    --backend vllm \
+    --dataset-name random \
+    --model meta-llama/Llama-3.2-3B-Instruct \
+    --num-prompts 20 \
+    --endpoint /v1/completions
+kill -s SIGTERM $server_pid &

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

@@ -2,119 +2,19 @@
 
 # This script build the CPU docker image and run the offline inference inside the container.
 # It serves a sanity check for compilation and basic model usage.
-set -ex
+set -euox pipefail
 
 # allow to bind to different cores
 CORE_RANGE=${CORE_RANGE:-48-95}
-# used for TP/PP E2E test
-OMP_CORE_RANGE=${OMP_CORE_RANGE:-48-95}
 NUMA_NODE=${NUMA_NODE:-1}
+IMAGE_NAME="cpu-test-$NUMA_NODE"
+TIMEOUT_VAL=$1
+TEST_COMMAND=$2
 
-export CMAKE_BUILD_PARALLEL_LEVEL=32
-
-# Setup cleanup
-remove_docker_container() {
-    set -e;
-    docker rm -f cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"-avx2 || true;
-}
-trap remove_docker_container EXIT
-remove_docker_container
-
-# Try building the docker image
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --progress plain --build-arg VLLM_CPU_DISABLE_AVX512="true" --tag cpu-test-"$NUMA_NODE"-avx2 --target vllm-test -f docker/Dockerfile.cpu .
+# building the docker image
+echo "--- :docker: Building Docker image"
+docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
 
 # Run the image, setting --shm-size=4g for tensor parallel.
-docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
-docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE"-avx2 cpu-test-"$NUMA_NODE"-avx2
-
-function cpu_tests() {
-  set -e
-  export NUMA_NODE=$2
-
-  # list packages
-  docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c "
-    set -e
-    pip list"
-
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    pip list"
-
-  # offline inference
-  docker exec cpu-test-"$NUMA_NODE"-avx2 bash -c "
-    set -e
-    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
-
-  # Run kernel tests
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
-    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
-    pytest -x -v -s tests/kernels/test_onednn.py"
-
-  # Run basic model test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    # Note: disable until supports V1
-    # pytest -x -v -s tests/kernels/attention/test_cache.py -m cpu_model
-    # pytest -x -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
-
-    pytest -x -v -s tests/models/language/generation -m cpu_model
-    VLLM_CPU_SGL_KERNEL=1 pytest -x -v -s tests/models/language/generation -m cpu_model
-
-    pytest -x -v -s tests/models/language/pooling -m cpu_model
-    pytest -x -v -s tests/models/multimodal/generation \
-                --ignore=tests/models/multimodal/generation/test_pixtral.py \
-                -m cpu_model"
-
-  # Run compressed-tensor test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    pytest -x -s -v \
-    tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
-
-  # Run AWQ/GPTQ test
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    pytest -x -s -v \
-    tests/quantization/test_cpu_wna16.py"
-
-  # Run multi-lora tests
-  docker exec cpu-test-"$NUMA_NODE" bash -c "
-    set -e
-    pytest -x -s -v \
-    tests/lora/test_qwenvl.py"
-
-  # online serving: tp+pp
-  docker exec cpu-test-"$NUMA_NODE" bash -c '
-    set -e
-    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS VLLM_CPU_SGL_KERNEL=1 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 &
-    server_pid=$!
-    timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
-    vllm bench serve \
-      --backend vllm \
-      --dataset-name random \
-      --model meta-llama/Llama-3.2-3B-Instruct \
-      --num-prompts 20 \
-      --endpoint /v1/completions
-    kill -s SIGTERM $server_pid &'
-
-  # online serving: tp+dp
-  docker exec cpu-test-"$NUMA_NODE" bash -c '
-    set -e
-    VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS VLLM_CPU_SGL_KERNEL=1 vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 &
-    server_pid=$!
-    timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
-    vllm bench serve \
-      --backend vllm \
-      --dataset-name random \
-      --model meta-llama/Llama-3.2-3B-Instruct \
-      --num-prompts 20 \
-      --endpoint /v1/completions
-    kill -s SIGTERM $server_pid &'
-}
-
-# All of CPU tests are expected to be finished less than 40 mins.
-export -f cpu_tests
-timeout 2.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
+docker run --rm --cpuset-cpus=$CORE_RANGE --cpuset-mems=$NUMA_NODE -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g $IMAGE_NAME \
+        timeout $TIMEOUT_VAL bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"

.buildkite/test-amd.yaml

@@ -70,6 +70,7 @@ steps:
   - vllm/
   - tests/test_inputs.py
   - tests/test_outputs.py
+  - tests/test_pooling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -82,6 +83,7 @@ steps:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
+  - pytest -v -s test_pooling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_
@@ -231,6 +233,7 @@ steps:
   - tests/compile/fullgraph/test_basic_correctness.py
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
+  - examples/offline_inference/new_weight_syncing/
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
@@ -266,10 +269,16 @@ steps:
   - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
+  # OLD rlhf examples
   - pushd ../examples/offline_inference
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   - popd
+  # NEW rlhf examples
+  - pushd ../examples/offline_inference/new_weight_syncing
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
+  - popd
 
 - label: Distributed Tests (8 GPUs) # 4min
   timeout_in_minutes: 10
@@ -854,10 +863,11 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
+  - tests/models/test_terratorch.py
   - tests/models/test_transformers.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
 
 - label: Basic Models Test (Other CPU) # 5min
   mirror_hardwares: [amdexperimental, amdproduction]

.buildkite/test-pipeline.yaml

@@ -63,6 +63,7 @@ steps:
   - vllm/
   - tests/test_inputs.py
   - tests/test_outputs.py
+  - tests/test_pooling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -75,6 +76,7 @@ steps:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
+  - pytest -v -s test_pooling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_
@@ -204,6 +206,7 @@ steps:
   - tests/compile/fullgraph/test_basic_correctness.py
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
+  - examples/offline_inference/new_weight_syncing/
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
@@ -238,10 +241,16 @@ steps:
   - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
+  # OLD rlhf examples
   - pushd ../examples/offline_inference
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
   - popd
+  # NEW rlhf examples
+  - pushd ../examples/offline_inference/new_weight_syncing
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
+  - popd
 
 - label: Distributed Tests (8 GPUs) # 4min
   timeout_in_minutes: 10
@@ -795,10 +804,11 @@ steps:
   torch_nightly: true
   source_file_dependencies:
   - vllm/
+  - tests/models/test_terratorch.py
   - tests/models/test_transformers.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
 
 - label: Basic Models Test (Other CPU) # 5min
   timeout_in_minutes: 10
@@ -1144,6 +1154,8 @@ steps:
   - pytest -v -s distributed/test_shm_broadcast.py
   - pytest -v -s distributed/test_shm_buffer.py
   - pytest -v -s distributed/test_shm_storage.py
+  - pytest -v -s distributed/test_packed_tensor.py
+  - pytest -v -s distributed/test_weight_transfer.py
 
 - label: 2 Node Tests (4 GPUs in total) # 16min
   timeout_in_minutes: 30

.buildkite/test_areas/distributed.yaml

@@ -62,6 +62,7 @@ steps:
   - tests/compile/fullgraph/test_basic_correctness.py
   - examples/offline_inference/rlhf.py
   - examples/offline_inference/rlhf_colocate.py
+  - examples/offline_inference/new_weight_syncing/
   - tests/examples/offline_inference/data_parallel.py
   - tests/v1/distributed
   - tests/v1/engine/test_engine_core_client.py
@@ -96,9 +97,14 @@ steps:
   - pytest -v -s distributed/test_symm_mem_allreduce.py
   # TODO: create a dedicated test section for multi-GPU example tests
   # when we have multiple distributed example tests
+  # OLD rlhf examples
   - cd ../examples/offline_inference
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
   - VLLM_ALLOW_INSECURE_SERIALIZATION=1 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
+  # NEW rlhf examples
+  - cd new_weight_syncing
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 rlhf_async_new_apis.py
 
 - label: Distributed Tests (8 GPUs)(H100)
   timeout_in_minutes: 10

.buildkite/test_areas/misc.yaml

@@ -122,6 +122,7 @@ steps:
   - vllm/
   - tests/test_inputs.py
   - tests/test_outputs.py
+  - tests/test_pooling_params.py
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
@@ -134,6 +135,7 @@ steps:
   - python3 standalone_tests/lazy_imports.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
+  - pytest -v -s test_pooling_params.py
   - pytest -v -s -m 'cpu_test' multimodal
   - pytest -v -s renderers
   - pytest -v -s tokenizers_

.buildkite/test_areas/models_basic.yaml

@@ -33,10 +33,11 @@ steps:
   timeout_in_minutes: 45
   source_file_dependencies:
   - vllm/
+  - tests/models/test_terratorch.py
   - tests/models/test_transformers.py
   - tests/models/test_registry.py
   commands:
-    - pytest -v -s models/test_transformers.py models/test_registry.py
+    - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
 
 - label: Basic Models Test (Other CPU) # 5min
   depends_on: 

cmake/external_projects/vllm_flash_attn.cmake

@@ -38,7 +38,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 2adfc8c2177c5b0e8ddeedfd5a8990d80eb496ff
+          GIT_TAG 188be16520ceefdc625fdf71365585d2ee348fe2
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

csrc/rocm/skinny_gemms.cu

@@ -1899,8 +1899,9 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitKQ_hf_sml_(const int K, const int Kp, const int M, const int Bx,
-                      const int By, const fp8_t* B, const fp8_t* __restrict__ A,
+    wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, const int M,
+                      const int Bx, const int By, const fp8_t* B,
+                      const fp8_t* __restrict__ A,
                       const scalar_t* __restrict__ BIAS, scalar_t* C,
                       const float* __restrict__ s_A,
                       const float* __restrict__ s_B, const int _WvPrGrp,
@@ -1924,9 +1925,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   __shared__ fp8_t s[max_lds_len];
 
   for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
-       k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
+       k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
+  #if defined(__gfx950__)
+    __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
+  #else
     *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
+  #endif
   }
+  asm volatile("s_waitcnt vmcnt(0)");
   __syncthreads();
 
   if (threadIdx.y >= _WvPrGrp) return;
@@ -1934,37 +1940,24 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
 
   using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
-  floatx16 sum[N][YTILE];
   float sA = *s_A;
   float sB = *s_B;
 
   while (m < M) {
-    for (int i = 0; i < YTILE; i++)
-      for (int n = 0; n < N; n++) sum[n][i] = {0.f};
-
-    bigType bigA[N][UNRL];
-    bigType bigB[YTILE][UNRL];
-
+    floatx16 sum[N][YTILE] = {};
     for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
-  #pragma unroll
-      for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-  #pragma unroll
-        for (uint32_t n = 0; n < N; ++n) bigA[n][k2].h8 = {0.f};
-  #pragma unroll
-        for (uint32_t y = 0; y < YTILE; ++y) bigB[y][k2].h8 = {0.f};
-      }
+      bigType bigA[N][UNRL] = {};
+      bigType bigB[YTILE][UNRL];
 
       // Fetch the weight matrix from memory!
   #pragma unroll
       for (uint32_t k2 = 0; k2 < UNRL; k2++) {
         uint32_t k = k1 + k2 * THRDS * A_CHUNK;
         uint32_t k_ = k + threadIdx.x * A_CHUNK;
-        if (k_ >= K) break;
-
-        const fp8_t* B_ = &B[(m + 0) * Kp + k_];
+        const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)];
   #pragma unroll
         for (uint32_t y = 0; y < YTILE; ++y) {
-          bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp])));
+          bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
         }
       }
 
@@ -1975,16 +1968,13 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
         uint32_t k_ = k + threadIdx.x * A_CHUNK;
         if (k_ >= K) break;
         for (int n = 0; n < N; n++) {
-          bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
+          bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
         }
       }
 
   // Do the matrix multiplication in interleaved manner
   #pragma unroll
       for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-        uint32_t k = k1 + k2 * THRDS * A_CHUNK;
-        if (k >= K) break;
-
         for (uint32_t n = 0; n < N; n++) {
           for (int i = 0; i < A_CHUNK; i += 8) {
             for (int y = 0; y < YTILE; ++y) {
@@ -2002,48 +1992,27 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       for (int y = 0; y < YTILE; y++) {
         float accm0 = sum[n][y][0];
         float accm16 = sum[n][y][8];
-        asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16));
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
+                                          1);  // row_shl1
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
+                                          1);  // row_shl2
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
+                                          1);  // row_shl3
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
+                                          1);  // row_shl8
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
+                                          1);  // row_shl9
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
+                                          1);  // row_shl10
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
+                                          1);  // row_shl11
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
         accm0 += __shfl(accm0, 36);
         accm16 += __shfl(accm16, 52);
         sum[n][y][0] = accm0 + __shfl(accm16, 16);
@@ -2051,19 +2020,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     }
 
     if (threadIdx.x == 0) {
+      scalar_t biases[N][YTILE] = {};
+      if (BIAS)
+        for (int n = 0; n < N; n++) {
+          for (int y = 0; y < YTILE; y++) {
+            biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
+          }
+        }
       for (int n = 0; n < N; n++) {
         for (int y = 0; y < YTILE; y++) {
           if (y + m >= M) break;  // To avoid mem access fault.
           sum[n][y][0] *= sA * sB;
           if constexpr (std::is_same_v<scalar_t, half>) {
-            if (BIAS)
-              sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
+            sum[n][y][0] += __half2float(biases[n][y]);
           } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            if (BIAS)
-              sum[n][y][0] +=
-                  __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
+            sum[n][y][0] += __bfloat162float(biases[n][y]);
           }
-          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);  // * sA * sB);
+          C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
         }
       }
     }
@@ -2074,9 +2047,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
 #else   // !defined(__HIP__MI3XX__) TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
-__global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
-                                  const int Bx, const int By, const fp8_t* B,
-                                  const fp8_t* __restrict__ A,
+__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
+                                  const int M, const int Bx, const int By,
+                                  const fp8_t* B, const fp8_t* __restrict__ A,
                                   const scalar_t* __restrict__ BIAS,
                                   scalar_t* C, const float* __restrict__ s_A,
                                   const float* __restrict__ s_B,
@@ -2089,8 +2062,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kp, const int M,
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
 __global__ void __launch_bounds__(WvPrGrp* THRDS)
-    wvSplitKQ_hf_(const int K, const int Kp, const int M, const int Bx,
-                  const int By, const fp8_t* B, const fp8_t* __restrict__ A,
+    wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, const int M,
+                  const int Bx, const int By, const fp8_t* B,
+                  const fp8_t* __restrict__ A,
                   const scalar_t* __restrict__ BIAS, scalar_t* C,
                   const float* __restrict__ s_A, const float* __restrict__ s_B,
                   const int _WvPrGrp, const int CuCount) {
@@ -2113,9 +2087,14 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   __shared__ fp8_t s[max_lds_len];
 
   for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK;
-       k < min__(K * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
+       k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) {
+  #if defined(__gfx950__)
+    __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0);
+  #else
     *((bigType*)(&s[k])) = *((bigType*)(&A[k]));
+  #endif
   }
+  asm volatile("s_waitcnt vmcnt(0)");
   __syncthreads();
 
   if (threadIdx.y >= _WvPrGrp) return;
@@ -2123,29 +2102,23 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
   uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE;
 
   using floatx16 = __attribute__((__vector_size__(16 * sizeof(float)))) float;
-  floatx16 sum[N][YTILE];
   float sA = *s_A;
   float sB = *s_B;
 
   while (m < M) {
-    for (int i = 0; i < YTILE; i++)
-      for (int n = 0; n < N; n++) sum[n][i] = {0};
-
-    bigType bigA[N][UNRL];
-    bigType bigB[YTILE][UNRL];
-
+    floatx16 sum[N][YTILE] = {};
     for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) {
+      bigType bigA[N][UNRL] = {};
+      bigType bigB[YTILE][UNRL];
+
       // Fetch the weight matrix from memory!
   #pragma unroll
       for (uint32_t k2 = 0; k2 < UNRL; k2++) {
         uint32_t k = k1 + k2 * THRDS * A_CHUNK;
         uint32_t k_ = k + threadIdx.x * A_CHUNK;
-        if (k_ >= K) break;
-
-        const fp8_t* B_ = &B[(m + 0) * Kp + k_];
+        const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)];
         for (int y = 0; y < YTILE; ++y) {
-          if (y + m >= M) break;  // To avoid mem access fault.
-          bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[y * Kp])));
+          bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp])));
         }
       }
 
@@ -2156,20 +2129,16 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
         uint32_t k_ = k + threadIdx.x * A_CHUNK;
         if (k_ >= K) break;
         for (int n = 0; n < N; n++) {
-          if (k_ + K * n < max_lds_len)
-            bigA[n][k2] = *((const bigType*)(&(s[k_ + K * n])));
+          if (k_ + Kap * n < max_lds_len)
+            bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n])));
           else
-            bigA[n][k2] = *((const bigType*)(&(A[k_ + K * n])));
+            bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n])));
         }
       }
 
   // Do the matrix multiplication in interleaved manner
   #pragma unroll
       for (uint32_t k2 = 0; k2 < UNRL; k2++) {
-        uint32_t k = k1 + k2 * THRDS * A_CHUNK;
-        uint32_t k_ = k + threadIdx.x * A_CHUNK;
-        if (k_ >= K) break;
-
         for (uint32_t n = 0; n < N; n++) {
           for (int i = 0; i < A_CHUNK; i += 8) {
             for (int y = 0; y < YTILE; ++y) {
@@ -2187,48 +2156,27 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
       for (int y = 0; y < YTILE; y++) {
         float accm0 = sum[n][y][0];
         float accm16 = sum[n][y][8];
-        asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][1]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:1 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][9]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][2]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:2 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][10]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][3]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:3 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][11]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][4]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:8 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][12]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][5]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:9 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][13]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][6]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:10 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][14]), "v"(accm16));
-        asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
-            : "=v"(accm0)
-            : "0"(accm0), "v"(sum[n][y][7]), "v"(accm0));
-        asm("v_add_f32 %0, %2, %3 row_shl:11 bound_ctrl:0 "
-            : "=v"(accm16)
-            : "0"(accm16), "v"(sum[n][y][15]), "v"(accm16));
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf,
+                                          1);  // row_shl1
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][9], 0x101, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf,
+                                          1);  // row_shl2
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][10], 0x102, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf,
+                                          1);  // row_shl3
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][11], 0x103, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][4], 0x108, 0xf, 0xf,
+                                          1);  // row_shl8
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][12], 0x108, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][5], 0x109, 0xf, 0xf,
+                                          1);  // row_shl9
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][13], 0x109, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][6], 0x10a, 0xf, 0xf,
+                                          1);  // row_shl10
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][14], 0x10a, 0xf, 0xf, 1);
+        accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][7], 0x10b, 0xf, 0xf,
+                                          1);  // row_shl11
+        accm16 += __builtin_amdgcn_mov_dpp(sum[n][y][15], 0x10b, 0xf, 0xf, 1);
         accm0 += __shfl(accm0, 36);
         accm16 += __shfl(accm16, 52);
         sum[n][y][0] = accm0 + __shfl(accm16, 16);
@@ -2236,17 +2184,21 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
     }
 
     if (threadIdx.x == 0) {
+      scalar_t biases[N][YTILE] = {};
+      if (BIAS)
+        for (int n = 0; n < N; n++) {
+          for (int y = 0; y < YTILE; y++) {
+            biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx];
+          }
+        }
       for (int n = 0; n < N; n++) {
         for (int y = 0; y < YTILE; y++) {
           if (y + m >= M) break;  // To avoid mem access fault.
           sum[n][y][0] *= sA * sB;
           if constexpr (std::is_same_v<scalar_t, half>) {
-            if (BIAS)
-              sum[n][y][0] += __half2float(BIAS[(m + y) % Bx + (n % By) * M]);
+            sum[n][y][0] += __half2float(biases[n][y]);
           } else if constexpr (std::is_same_v<scalar_t, __hip_bfloat16>) {
-            if (BIAS)
-              sum[n][y][0] +=
-                  __bfloat162float(BIAS[(m + y) % Bx + (n % By) * M]);
+            sum[n][y][0] += __bfloat162float(biases[n][y]);
           }
           C[m + y + n * M] = __float2s<scalar_t>(sum[n][y][0]);
         }
@@ -2259,9 +2211,9 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
 #else   // !defined(__HIP__MI3XX__) TODO: Add NAVI support
 template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
           int A_CHUNK, int UNRL, int N>
-__global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
-                              const int Bx, const int By, const fp8_t* B,
-                              const fp8_t* __restrict__ A,
+__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
+                              const int M, const int Bx, const int By,
+                              const fp8_t* B, const fp8_t* __restrict__ A,
                               const scalar_t* __restrict__ BIAS, scalar_t* C,
                               const float* __restrict__ s_A,
                               const float* __restrict__ s_B, const int _WvPrGrp,
@@ -2270,17 +2222,18 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kp, const int M,
 }
 #endif  // defined(__HIP__MI3XX__) TODO: Add NAVI support
 
-void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
+void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
                const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
                const at::Tensor& scale_a, const at::Tensor& scale_b,
                const int64_t CuCount) {
   static c10::ScalarType kFp8Type = is_fp8_ocp()
                                         ? c10::ScalarType::Float8_e4m3fn
                                         : c10::ScalarType::Float8_e4m3fnuz;
-  auto M_in = in_a.size(0);
-  auto K_in = in_a.size(1);
-  auto N_in = in_b.size(0);
-  auto Kp_in = in_a.stride(0);
+  auto M_in = in_b.size(0);
+  auto K_in = in_b.size(1);
+  auto N_in = in_a.size(0);
+  auto Kap_in = in_a.stride(0);
+  auto Kbp_in = in_b.stride(0);
   auto Bx_in =
       (in_bias.has_value() && in_bias->numel() > 0)
           ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0)
@@ -2300,23 +2253,22 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   const int max_lds_len = get_lds_size();
 
-#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _YTILEb, _UNRLs, _UNRLm, _UNRLb, \
-                  _N)                                                          \
-  {                                                                            \
-    dim3 block(64, _WvPrGrp);                                                  \
-    if ((K_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) {               \
-      int __wvPrGrp = mindiv(M_in, CuCount * _YTILEs, _WvPrGrp);               \
-      wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N>  \
-          <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
-                                       b_ptr, bias_ptr, c_ptr, s_a, s_b,       \
-                                       __wvPrGrp, CuCount);                    \
-    } else {                                                                   \
-      int __wvPrGrp = mindiv(M_in, CuCount * _YTILEm, _WvPrGrp);               \
-      wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N>      \
-          <<<grid, block, 0, stream>>>(K_in, Kp_in, M_in, Bx_in, By_in, a_ptr, \
-                                       b_ptr, bias_ptr, c_ptr, s_a, s_b,       \
-                                       __wvPrGrp, CuCount);                    \
-    }                                                                          \
+#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N)             \
+  {                                                                           \
+    dim3 block(64, _WvPrGrp);                                                 \
+    if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) {            \
+      int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16));     \
+      wvSplitKQ_hf_sml_<fptype, fp8_t, 64, _YTILEs, _WvPrGrp, 16, _UNRLs, _N> \
+          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
+                                       By_in, b_ptr, a_ptr, bias_ptr, c_ptr,  \
+                                       s_a, s_b, __wvPrGrp, CuCount);         \
+    } else {                                                                  \
+      int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16));     \
+      wvSplitKQ_hf_<fptype, fp8_t, 64, _YTILEm, _WvPrGrp, 16, _UNRLm, _N>     \
+          <<<grid, block, 0, stream>>>(K_in, Kap_in, Kbp_in, M_in, Bx_in,     \
+                                       By_in, b_ptr, a_ptr, bias_ptr, c_ptr,  \
+                                       s_a, s_b, __wvPrGrp, CuCount);         \
+    }                                                                         \
   }
 
   AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] {
@@ -2332,16 +2284,16 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
                           : nullptr;
       switch (N_in) {
         case 1:
-          WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 1)
+          WVSPLITKQ(12, 2, 2, 2, 2, 1)
           break;
         case 2:
-          WVSPLITKQ(16, 2, 2, 2, 2, 2, 2, 2)
+          WVSPLITKQ(12, 2, 2, 2, 2, 2)
           break;
         case 3:
-          WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 3)
+          WVSPLITKQ(8, 2, 2, 1, 1, 3)
           break;
         case 4:
-          WVSPLITKQ(16, 4, 7, 7, 1, 1, 1, 4)
+          WVSPLITKQ(4, 2, 2, 1, 1, 4)
           break;
         default:
           throw std::runtime_error(

docs/contributing/model/basic.md

@@ -138,7 +138,7 @@ These models should follow the same instructions as case (1), but they should in
 
 For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
 Please follow the same guidelines as case (2) for implementing these models.
-We use "mamba-like" to refer to layers that posses a state that is updated in-place, rather than being appended-to (like KV cache for attention).
+We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
 For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
 It is also necessary to implement the "attention meta-data" class which handles the meta-data that is common across all layers.
 Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.

docs/contributing/model/multimodal.md

@@ -739,7 +739,7 @@ Each [PromptUpdate][vllm.multimodal.processing.PromptUpdate] instance specifies
         ```
 
     However, this is not entirely correct. After `FuyuImageProcessor.preprocess_with_tokenizer_info` is called,
-    a BOS token (`<s>`) is also added to the promopt:
+    a BOS token (`<s>`) is also added to the prompt:
 
     ??? code
 

docs/design/plugin_system.md

@@ -45,7 +45,7 @@ Every plugin has three parts:
 
 ## Types of supported plugins
 
-- **General plugins** (with group name `vllm.general_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree models into vLLM. This is done by calling `ModelRegistry.register_model` to register the model inside the plugin function.
+- **General plugins** (with group name `vllm.general_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree models into vLLM. This is done by calling `ModelRegistry.register_model` to register the model inside the plugin function. For an example of an official model plugin, see the [bart-plugin](https://github.com/vllm-project/bart-plugin) which adds support for `BartForConditionalGeneration`.
 
 - **Platform plugins** (with group name `vllm.platform_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree platforms into vLLM. The plugin function should return `None` when the platform is not supported in the current environment, or the platform class's fully qualified name when the platform is supported.
 

docs/features/nixl_connector_usage.md

@@ -213,6 +213,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
 ```
 
+### Cross layers blocks
+
+By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
+To enable this feature:
+
+```bash
+--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
+```
+
 ## Example Scripts/Code
 
 Refer to these example scripts in the vLLM repository:

docs/getting_started/quickstart.md

@@ -57,7 +57,7 @@ This guide will help you quickly get started with vLLM to perform:
         It currently supports Python 3.12, ROCm 7.0 and `glibc >= 2.35`.
 
     !!! note    
-        Note that, previously, docker images were published using AMD's docker release pipeline and were located `rocm/vlm-dev`. This is being deprecated by using vLLM's docker release pipeline.
+        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.
 
 === "Google TPU"
 

docs/models/supported_models.md

@@ -174,6 +174,16 @@ class MyConfig(PretrainedConfig):
     - The `list` in the first element of the `tuple` contains the names of the input arguments
     - The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code
 
+### Plugins
+
+Some model architectures are supported via vLLM plugins. These plugins extend vLLM's capabilities through the [plugin system](../design/plugin_system.md).
+
+| Architecture | Models | Plugin Repository |
+|--------------|--------|-------------------|
+| `BartForConditionalGeneration` | BART | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
+
+For other model architectures not natively supported, in particular for Encoder-Decoder models, we recommend following a similar pattern by implementing support through the plugin system.
+
 ## Loading a Model
 
 ### Hugging Face Hub

docs/usage/v1_guide.md

@@ -134,9 +134,12 @@ Please note that prefix caching is not yet supported for any of the above models
 
 #### Encoder-Decoder Models
 
-Whisper is supported. Other models requiring cross-attention between separate
-encoder and decoder (e.g., `BartForConditionalGeneration`,
-`MllamaForConditionalGeneration`) are no longer supported.
+Whisper is supported natively. Other encoder-decoder models are supported via the plugin system:
+
+- **BART**: `BartForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
+
+For other encoder-decoder models (e.g., `MllamaForConditionalGeneration`), we recommend
+following a similar pattern by implementing support through the [plugin system](../design/plugin_system.md).
 
 ### Features
 

examples/offline_inference/new_weight_syncing/rlhf.py

@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates reinforcement learning using vLLM and Ray,
+with native weight syncing APIs at engine instance.
+
+The script separates training and inference workloads onto distinct GPUs
+so that Ray can manage process placement and inter-process communication.
+A Hugging Face Transformer model occupies one GPU for training, whereas a
+2x tensor-parallel vLLM inference engine occupies two GPUs.
+
+The example performs the following steps:
+* Load the training model on one gpu (scheduled via ray)
+* Initialize the inference model with dummy weights across
+  two gpus using vLLM's tensor parallelism and Ray placement groups.
+* Generate gibberish from a list of prompts using the randomly initialized
+  inference engine.
+* Update the weights of the training model and broadcast the updated weights
+  to the inference engine by using a Ray collective RPC group.
+* Generating from the list of prompts after weight sync should result
+  in sensible outputs.
+
+This example assumes a single-node cluster with three GPUs, but Ray
+supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
+workloads. Residual GPU activity interferes with vLLM memory profiling and
+causes unexpected behavior.
+"""
+
+import os
+
+import ray
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from transformers import AutoModelForCausalLM
+
+from vllm import LLM, SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLWeightTransferEngine,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+
+MODEL_NAME = "facebook/opt-125m"
+# MODEL_NAME = "inference-optimization/Qwen3-0.6B-W4A16-G128"
+
+
+class MyLLM(LLM):
+    """Configure the vLLM worker for Ray placement group execution."""
+
+    def __init__(self, *args, **kwargs):
+        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
+        super().__init__(*args, **kwargs)
+
+
+@ray.remote(num_gpus=1)
+class TrainModel:
+    """Ray actor that wraps the training model on a dedicated GPU."""
+
+    def __init__(self, model_name: str):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+        ).to("cuda:0")
+
+        self.port = get_open_port()
+        self.master_address = get_ip()
+
+    def get_master_address_and_port(self):
+        return self.master_address, self.port
+
+    def get_weight_metadata(self):
+        """Return weight names, dtypes, and shapes for weight transfer."""
+        names = []
+        dtype_names = []
+        shapes = []
+        for name, p in self.model.named_parameters():
+            names.append(name)
+            dtype_names.append(str(p.dtype).split(".")[-1])
+            shapes.append(list(p.shape))
+        return names, dtype_names, shapes
+
+    def init_weight_transfer_group(self, world_size):
+        """Initialize the NCCL process group for weight transfer."""
+        self.model_update_group = NCCLWeightTransferEngine.trainer_init(
+            dict(
+                master_address=self.master_address,
+                master_port=self.port,
+                world_size=world_size,
+            ),
+        )
+
+    def broadcast_weights(self, packed: bool = True):
+        """Broadcast weights to the inference engine."""
+        NCCLWeightTransferEngine.trainer_send_weights(
+            iterator=self.model.named_parameters(),
+            group=self.model_update_group,
+            packed=packed,
+        )
+
+
+# Initialize Ray and set the visible devices. The vLLM engine will
+# be placed on GPUs 1 and 2.
+ray.init()
+
+# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
+# Learn more about Ray placement groups:
+# https://docs.ray.io/en/latest/placement-groups.html
+# Launch the training model actor. Ray's resource scheduler will allocate
+# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
+train_model = TrainModel.remote(MODEL_NAME)
+
+pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
+ray.get(pg_inference.ready())
+scheduling_inference = PlacementGroupSchedulingStrategy(
+    placement_group=pg_inference,
+    placement_group_capture_child_tasks=True,
+    placement_group_bundle_index=0,
+)
+
+# Launch the vLLM inference engine. The `enforce_eager` flag reduces
+# start-up latency.
+# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
+# are now native to vLLM workers.
+llm = ray.remote(
+    num_cpus=0,
+    num_gpus=0,
+    scheduling_strategy=scheduling_inference,
+)(MyLLM).remote(
+    model=MODEL_NAME,
+    enforce_eager=True,
+    tensor_parallel_size=2,
+    data_parallel_size=1,
+    distributed_executor_backend="ray",
+    weight_transfer_config=WeightTransferConfig(backend="nccl"),
+    load_format="dummy",
+    quantization="fp8",
+)
+
+# Generate text from the prompts.
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+sampling_params = SamplingParams(temperature=0)
+
+outputs = ray.get(llm.generate.remote(prompts, sampling_params))
+
+# Generate text with the initial model. The output is expected to be nonsense
+# because the weights are randomly initialized.
+print("-" * 50)
+for output in outputs:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+    print("-" * 50)
+
+# Set up the communication channel between the training process and the
+# inference engine.
+master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
+
+world_size = ray.get(llm.get_world_size.remote()) + 1  # +1 for the trainer
+inference_handle = llm.init_weight_transfer_engine.remote(
+    dict(
+        init_info=dict(
+            master_address=master_address,
+            master_port=master_port,
+            rank_offset=1,
+            world_size=world_size,
+        )
+    )
+)
+
+# Initialize weight transfer group on both the training actor and inference engine
+train_handle = train_model.init_weight_transfer_group.remote(world_size)
+ray.get([train_handle, inference_handle])
+
+# Synchronize the updated weights to the inference engine using batched API.
+# Collect all weight metadata from the training actor
+names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
+
+# Issue update_weights call with NCCL-specific update info
+# packed=True enables efficient batched tensor broadcasting
+inference_handle = llm.update_weights.remote(
+    dict(
+        update_info=dict(
+            names=names,
+            dtype_names=dtype_names,
+            shapes=shapes,
+            packed=True,
+        )
+    )
+)
+
+# Broadcast all weights from trainer using the weight transfer API
+train_handle = train_model.broadcast_weights.remote(packed=True)
+ray.get([train_handle, inference_handle])
+
+# Generate text with the updated model. The output is expected to be normal
+# because the weights are updated.
+outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
+print("-" * 50)
+for output in outputs_updated:
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+    print("-" * 50)

examples/offline_inference/new_weight_syncing/rlhf_async_new_apis.py

@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates async reinforcement learning using vLLM and Ray,
+with native weight syncing APIs at engine instance.
+
+The script separates training and inference workloads onto distinct GPUs
+so that Ray can manage process placement and inter-process communication.
+A Hugging Face Transformer model occupies one GPU for training, whereas a
+2x tensor-parallel vLLM inference engine occupies two GPUs.
+
+The example performs the following steps:
+* Load the training model on one gpu (scheduled via ray)
+* Initialize the inference model with dummy weights across
+  two gpus using vLLM's tensor parallelism and Ray placement groups.
+* Generate gibberish from a list of prompts using the randomly initialized
+  inference engine.
+* Pause generation once generation completes for one sequence
+* Update the weights of the training model and broadcast the updated weights
+  to the inference engine by using a Ray collective RPC group.
+* Resume generation and print out the results
+
+This example assumes a single-node cluster with three GPUs, but Ray
+supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
+workloads. Residual GPU activity interferes with vLLM memory profiling and
+causes unexpected behavior.
+"""
+
+import os
+import uuid
+from dataclasses import asdict
+
+import ray
+import torch
+from ray.util.placement_group import placement_group
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import vllm
+from vllm import SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLWeightTransferEngine,
+    NCCLWeightTransferInitInfo,
+    NCCLWeightTransferUpdateInfo,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+from vllm.v1.executor import Executor
+
+MODEL_NAME = "facebook/opt-125m"
+
+
+class MyLLM(vllm.AsyncLLMEngine):
+    """Configure the vLLM worker for Ray placement group execution."""
+
+    def __init__(self, **kwargs):
+        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0,1"
+        engine_args = vllm.AsyncEngineArgs(**kwargs)
+        vllm_config = engine_args.create_engine_config()
+        executor_class = Executor.get_class(vllm_config)
+        super().__init__(
+            vllm_config=vllm_config,
+            executor_class=executor_class,
+            log_requests=engine_args.enable_log_requests,
+            log_stats=not engine_args.disable_log_stats,
+        )
+
+    async def generate_with_retry(
+        self, prompt_token_ids: list[int], sampling_params: vllm.SamplingParams
+    ) -> vllm.RequestOutput:
+        finish_reason = "abort"
+        while finish_reason == "abort":
+            async for request_output in self.generate(
+                {"prompt_token_ids": prompt_token_ids},
+                sampling_params,
+                request_id=str(uuid.uuid4()),
+            ):
+                output = request_output
+            finish_reason = output.outputs[0].finish_reason
+            if finish_reason == "abort":
+                print(
+                    f"ABORT, prompt_token_ids: {prompt_token_ids}, "
+                    f"generated token_ids: {list(output.outputs[0].token_ids)}"
+                )
+            prompt_token_ids = prompt_token_ids + list(output.outputs[0].token_ids)
+        return output
+
+
+@ray.remote(num_gpus=1)
+class TrainModel:
+    """Ray actor that wraps the training model on a dedicated GPU."""
+
+    def __init__(self, model_name: str):
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, dtype=torch.bfloat16
+        ).to("cuda:0")
+        self.port = get_open_port()
+        self.master_address = get_ip()
+
+    def get_master_address_and_port(self):
+        return self.master_address, self.port
+
+    def get_weight_metadata(self):
+        """Return weight names, dtypes, and shapes for weight transfer."""
+        names = []
+        dtype_names = []
+        shapes = []
+        for name, p in self.model.named_parameters():
+            names.append(name)
+            dtype_names.append(str(p.dtype).split(".")[-1])
+            shapes.append(list(p.shape))
+        return names, dtype_names, shapes
+
+    def init_weight_transfer_group(self, world_size):
+        """Initialize the NCCL process group for weight transfer."""
+        self.model_update_group = NCCLWeightTransferEngine.trainer_init(
+            dict(
+                master_address=self.master_address,
+                master_port=self.port,
+                world_size=world_size,
+            ),
+        )
+
+    def broadcast_weights(self, packed: bool = True):
+        """Broadcast weights to the inference engine."""
+        NCCLWeightTransferEngine.trainer_send_weights(
+            iterator=self.model.named_parameters(),
+            group=self.model_update_group,
+            packed=packed,
+        )
+
+
+# Initialize Ray and set the visible devices. The vLLM engine will
+# be placed on GPUs 1 and 2.
+ray.init()
+
+# Launch the training model actor. Ray's resource scheduler will allocate
+# 1 GPU (via num_gpus=1 in the decorator), ensuring pg_inference gets different GPUs.
+train_model = TrainModel.remote(MODEL_NAME)
+
+# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
+# Learn more about Ray placement groups:
+# https://docs.ray.io/en/latest/placement-groups.html
+
+pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
+ray.get(pg_inference.ready())
+scheduling_inference = PlacementGroupSchedulingStrategy(
+    placement_group=pg_inference,
+    placement_group_capture_child_tasks=True,
+    placement_group_bundle_index=0,
+)
+
+# Launch the vLLM inference engine. The `enforce_eager` flag reduces
+# start-up latency.
+# Note: Weight transfer APIs (init_weight_transfer_engine, update_weights)
+# are now native to vLLM workers.
+llm = ray.remote(
+    num_cpus=0,
+    num_gpus=0,
+    scheduling_strategy=scheduling_inference,
+)(MyLLM).remote(
+    model=MODEL_NAME,
+    enforce_eager=True,
+    tensor_parallel_size=2,
+    distributed_executor_backend="ray",
+    load_format="dummy",
+    weight_transfer_config=WeightTransferConfig(backend="nccl"),
+)
+
+# Generate text from the prompts.
+prompts = [
+    "My name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+# Tokenize prompts to token IDs
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+prompt_token_ids_list = [
+    tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts
+]
+
+sampling_params = [
+    SamplingParams(temperature=0, max_tokens=2),
+    SamplingParams(temperature=0, max_tokens=32),
+    SamplingParams(temperature=0, max_tokens=32),
+    SamplingParams(temperature=0, max_tokens=32),
+]
+
+# Set up the communication channel between the training process and the
+# inference engine.
+master_address, master_port = ray.get(train_model.get_master_address_and_port.remote())
+
+world_size = 3  # 1 trainer + 2 inference workers (tensor_parallel_size=2)
+inference_handle = llm.init_weight_transfer_engine.remote(
+    WeightTransferInitRequest(
+        init_info=asdict(
+            NCCLWeightTransferInitInfo(
+                master_address=master_address,
+                master_port=master_port,
+                rank_offset=1,
+                world_size=world_size,
+            )
+        )
+    )
+)
+
+# Initialize weight transfer group on both the training actor and inference engine
+train_handle = train_model.init_weight_transfer_group.remote(world_size)
+ray.get([train_handle, inference_handle])
+
+
+generation_futures = [
+    llm.generate_with_retry.remote(prompt_token_ids, params)
+    for prompt_token_ids, params in zip(prompt_token_ids_list, sampling_params)
+]
+
+finished, pending = ray.wait(generation_futures, num_returns=1)
+
+# Pause generation in preparation for weight sync
+ray.get(llm.pause_generation.remote(wait_for_inflight_requests=False))
+
+# Synchronize the updated weights to the inference engine using batched API.
+# Collect all weight metadata from the training actor
+names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
+
+# Issue update_weights call with NCCL-specific update info
+# packed=True enables efficient batched tensor broadcasting
+inference_handle = llm.update_weights.remote(
+    WeightTransferUpdateRequest(
+        update_info=asdict(
+            NCCLWeightTransferUpdateInfo(
+                names=names,
+                dtype_names=dtype_names,
+                shapes=shapes,
+                packed=True,
+            )
+        )
+    )
+)
+
+# Broadcast all weights from trainer using the weight transfer API
+train_handle = train_model.broadcast_weights.remote(packed=True)
+ray.get([train_handle, inference_handle])
+
+# Resume generation since weight sync is complete
+ray.get(llm.resume_generation.remote())
+
+# Get outputs separately - finished completed before pause, pending were paused/resumed
+finished_outputs = ray.get(finished)
+pending_outputs = ray.get(pending)
+
+# Requests that finished before the pause: all generation used original weights
+print("-" * 50)
+print("Requests that completed BEFORE weight change:")
+print("-" * 50)
+for output in finished_outputs:
+    prompt_text = tokenizer.decode(output.prompt_token_ids)
+    print(f"Prompt: {prompt_text!r}")
+    print(f"Generated (with original weights): {output.outputs[0].text!r}")
+    print("-" * 50)
+
+# Requests that were paused mid-generation: some text before, some after weight change
+print("Requests that were PAUSED and RESUMED after weight change:")
+print("-" * 50)
+for output in pending_outputs:
+    # Decode the full prompt token IDs (original + generated before pause)
+    full_prompt_text = tokenizer.decode(output.prompt_token_ids)
+    # Find the original prompt by checking which one this output started with
+    original_prompt = next(p for p in prompts if full_prompt_text.startswith(p))
+    # output.prompt_token_ids contains original prompt + tokens generated before pause
+    # output.outputs[0].text is what was generated after resuming with new weights
+    text_before_pause = full_prompt_text[len(original_prompt) :]
+    text_after_pause = output.outputs[0].text
+    print(f"Original prompt: {original_prompt!r}")
+    print(f"Generated before weight change: {text_before_pause!r}")
+    print(f"Generated after weight change: {text_after_pause!r}")
+    print("-" * 50)

examples/offline_inference/spec_decode.py

@@ -75,6 +75,7 @@ def parse_args():
     parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
     parser.add_argument("--disable-padded-drafter-batch", action="store_true")
     parser.add_argument("--max-num-seqs", type=int, default=None)
+    parser.add_argument("--parallel-drafting", action="store_true")
     parser.add_argument("--allowed-local-media-path", type=str, default="")
     return parser.parse_args()
 
@@ -121,6 +122,7 @@ def main(args):
             "model": eagle_dir,
             "num_speculative_tokens": args.num_spec_tokens,
             "disable_padded_drafter_batch": args.disable_padded_drafter_batch,
+            "parallel_drafting": args.parallel_drafting,
         }
     elif args.method == "ngram":
         speculative_config = {
@@ -137,6 +139,7 @@ def main(args):
             "num_speculative_tokens": args.num_spec_tokens,
             "enforce_eager": args.enforce_eager,
             "max_model_len": args.max_model_len,
+            "parallel_drafting": args.parallel_drafting,
         }
     elif args.method == "mtp":
         speculative_config = {

examples/online_serving/rlhf_http.py

@@ -0,0 +1,241 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
+via HTTP API, with native weight syncing APIs.
+
+Unlike rlhf.py which creates a vLLM instance programmatically, this script
+assumes you have already started a vLLM server using `vllm serve`. It uses:
+- OpenAI-compatible API for inference requests
+- HTTP endpoints for weight transfer control plane
+- NCCL for actual weight data transfer
+
+Prerequisites:
+    Start a vLLM server with weight transfer enabled:
+
+    $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
+        --enforce-eager \
+        --weight-transfer-config '{"backend": "nccl"}' \
+        --load-format dummy
+
+    Then run this script:
+
+    $ python rlhf_http.py
+
+The example performs the following steps:
+
+* Load the training model on GPU 0.
+* Generate text using the vLLM server via OpenAI-compatible API. The output
+  is expected to be nonsense because the server is initialized with dummy weights.
+* Initialize weight transfer via HTTP endpoint.
+* Broadcast the real weights from the training model to the vLLM server
+  using NCCL.
+* Generate text again to show normal output after the weight update.
+"""
+
+import requests
+import torch
+from openai import OpenAI
+from transformers import AutoModelForCausalLM
+
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLWeightTransferEngine,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+
+BASE_URL = "http://localhost:8000"
+MODEL_NAME = "facebook/opt-125m"
+
+
+def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
+    """Generate completions using the OpenAI-compatible API."""
+    results = []
+    for prompt in prompts:
+        response = client.completions.create(
+            model=model,
+            prompt=prompt,
+            max_tokens=32,
+            temperature=0,
+        )
+        results.append(response.choices[0].text)
+    return results
+
+
+def init_weight_transfer_engine(
+    base_url: str,
+    master_address: str,
+    master_port: int,
+    rank_offset: int,
+    world_size: int,
+) -> None:
+    """Initialize weight transfer via HTTP endpoint."""
+    url = f"{base_url}/init_weight_transfer_engine"
+    payload = {
+        "init_info": dict(
+            master_address=master_address,
+            master_port=master_port,
+            rank_offset=rank_offset,
+            world_size=world_size,
+        )
+    }
+    response = requests.post(url, json=payload, timeout=60)
+    response.raise_for_status()
+
+
+def update_weights(
+    base_url: str,
+    names: list[str],
+    dtype_names: list[str],
+    shapes: list[list[int]],
+    packed: bool = False,
+) -> None:
+    """Update weights via HTTP endpoint."""
+    url = f"{base_url}/update_weights"
+    payload = {
+        "update_info": dict(
+            names=names,
+            dtype_names=dtype_names,
+            shapes=shapes,
+            packed=packed,
+        )
+    }
+    response = requests.post(url, json=payload, timeout=300)
+    response.raise_for_status()
+
+
+def pause_generation(base_url: str) -> None:
+    """Pause generation via HTTP endpoint."""
+    url = f"{base_url}/pause"
+    response = requests.post(url, timeout=60)
+    response.raise_for_status()
+
+
+def resume_generation(base_url: str) -> None:
+    """Resume generation via HTTP endpoint."""
+    url = f"{base_url}/resume"
+    response = requests.post(url, timeout=60)
+    response.raise_for_status()
+
+
+def get_world_size(base_url: str) -> int:
+    """Get world size from the vLLM server."""
+    url = f"{base_url}/get_world_size"
+    response = requests.get(url, timeout=10)
+    response.raise_for_status()
+    return response.json()["world_size"]
+
+
+def main():
+    # Get the inference world size from the vLLM server
+    inference_world_size = get_world_size(BASE_URL)
+    world_size = inference_world_size + 1  # +1 for the trainer
+    device = f"cuda:{inference_world_size}"
+    torch.cuda.set_device(device)
+
+    # Load the training model
+    print(f"Loading training model: {MODEL_NAME}")
+    train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
+    train_model.to(device)
+
+    # Create OpenAI client pointing to the vLLM server
+    client = OpenAI(
+        base_url=f"{BASE_URL}/v1",
+        api_key="EMPTY",  # vLLM doesn't require an API key by default
+    )
+
+    # Test prompts
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    # Generate text before weight update. The output is expected to be nonsense
+    # because the server is initialized with dummy weights.
+    print("-" * 50)
+    print("Generating text BEFORE weight update (expect nonsense):")
+    print("-" * 50)
+    outputs = generate_completions(client, MODEL_NAME, prompts)
+    for prompt, generated_text in zip(prompts, outputs):
+        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+    # Set up the communication channel between the training process and the
+    # vLLM server. The trainer is rank 0, vLLM worker(s) start at rank_offset.
+    master_address = get_ip()
+    master_port = get_open_port()
+    rank_offset = 1
+
+    print(f"Initializing weight transfer: master={master_address}:{master_port}")
+
+    # Initialize weight transfer on vLLM server (this is async, server will
+    # wait for NCCL connection)
+    import threading
+
+    init_thread = threading.Thread(
+        target=init_weight_transfer_engine,
+        args=(BASE_URL, master_address, master_port, rank_offset, world_size),
+    )
+    init_thread.start()
+
+    # Initialize NCCL process group on trainer side
+    model_update_group = NCCLWeightTransferEngine.trainer_init(
+        dict(
+            master_address=master_address,
+            master_port=master_port,
+            world_size=world_size,
+        ),
+    )
+
+    # Wait for init_weight_transfer_engine to complete
+    init_thread.join()
+
+    # Pause generation before weight sync
+    pause_generation(BASE_URL)
+
+    # Collect weight metadata for the update request
+    names = []
+    dtype_names = []
+    shapes = []
+    for name, p in train_model.named_parameters():
+        names.append(name)
+        dtype_names.append(str(p.dtype).split(".")[-1])
+        shapes.append(list(p.shape))
+
+    # Start the update_weights call in a separate thread since it will block
+    # waiting for NCCL broadcasts
+    # packed=True enables efficient batched tensor broadcasting
+    update_thread = threading.Thread(
+        target=update_weights,
+        args=(BASE_URL, names, dtype_names, shapes, True),  # packed=True
+    )
+    update_thread.start()
+
+    # Broadcast all weights from trainer to vLLM workers
+    print("Broadcasting weights via NCCL...")
+    NCCLWeightTransferEngine.trainer_send_weights(
+        iterator=train_model.named_parameters(),
+        group=model_update_group,
+        packed=True,
+    )
+
+    # Wait for update_weights to complete
+    update_thread.join()
+
+    # Resume generation after weight sync
+    resume_generation(BASE_URL)
+
+    # Generate text after weight update. The output is expected to be normal
+    # because the real weights are now loaded.
+    print("-" * 50)
+    print("Generating text AFTER weight update:")
+    print("-" * 50)
+    outputs_updated = generate_completions(client, MODEL_NAME, prompts)
+    for prompt, generated_text in zip(prompts, outputs_updated):
+        print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
+        print("-" * 50)
+
+
+if __name__ == "__main__":
+    main()

tests/distributed/test_packed_tensor.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for packed tensor broadcasting functionality.
+
+Unit tests for packed_broadcast_producer and packed_broadcast_consumer.
+These utilities enable efficient batched tensor transfer over NCCL.
+"""
+
+import pytest
+import torch
+
+from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo
+from vllm.distributed.weight_transfer.packed_tensor import (
+    packed_broadcast_consumer,
+    packed_broadcast_producer,
+)
+
+
+class MockCommunicationGroup:
+    """Mock communication group for testing producer broadcast operations."""
+
+    def __init__(self):
+        self.broadcasted_tensors: list[torch.Tensor] = []
+        self.broadcast_count = 0
+        self.device = torch.device("cuda:0")
+
+    def broadcast(self, tensor, src):
+        """Mock broadcast that stores the tensor for later verification."""
+        self.broadcasted_tensors.append(tensor.clone())
+        self.broadcast_count += 1
+
+
+class MockConsumerCommunicationGroup:
+    """Mock communication group for consumer that returns pre-stored tensors."""
+
+    def __init__(self, tensors_to_return: list[torch.Tensor]):
+        self.tensors_to_return = tensors_to_return
+        self.current_index = 0
+        self.device = torch.device("cuda:0")
+
+    def broadcast(self, tensor, src):
+        """Mock broadcast that fills the tensor with pre-stored data."""
+        if self.current_index < len(self.tensors_to_return):
+            tensor.copy_(self.tensors_to_return[self.current_index])
+            self.current_index += 1
+
+
+def create_mock_model_params(
+    num_layers: int = 3,
+    dtype: torch.dtype = torch.float32,
+) -> list[tuple[str, torch.Tensor]]:
+    """Create mock model parameters for testing."""
+    params = []
+    for i in range(num_layers):
+        params.append((f"layer{i}.weight", torch.randn(10, 20, dtype=dtype)))
+        params.append((f"layer{i}.bias", torch.randn(10, dtype=dtype)))
+    return params
+
+
+def create_state_dict_info(
+    params: list[tuple[str, torch.Tensor]],
+) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
+    """Create state dict info (name -> (shape, dtype)) from params."""
+    return {name: (tuple(tensor.shape), tensor.dtype) for name, tensor in params}
+
+
+# --- Unit Tests: NCCLWeightTransferUpdateInfo packed field ---
+
+
+class TestNCCLWeightTransferUpdateInfoPacked:
+    """Test NCCLWeightTransferUpdateInfo dataclass packed field."""
+
+    def test_packed_default_false(self):
+        """Test that packed defaults to False."""
+        info = NCCLWeightTransferUpdateInfo(
+            names=["layer.weight"],
+            dtype_names=["float32"],
+            shapes=[[10, 10]],
+        )
+        assert info.packed is False
+
+    def test_packed_can_be_set_true(self):
+        """Test that packed can be set to True."""
+        info = NCCLWeightTransferUpdateInfo(
+            names=["layer.weight"],
+            dtype_names=["float32"],
+            shapes=[[10, 10]],
+            packed=True,
+        )
+        assert info.packed is True
+
+
+# --- Unit Tests: packed_broadcast_producer ---
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestPackedBroadcastProducer:
+    """Test packed_broadcast_producer function."""
+
+    def test_producer_broadcasts_tensors(self):
+        """Test that producer broadcasts all tensors."""
+        params = create_mock_model_params()
+        params_cuda = [(name, tensor.cuda()) for name, tensor in params]
+
+        mock_group = MockCommunicationGroup()
+
+        # Use a small target size to force multiple batches
+        packed_broadcast_producer(
+            iterator=iter(params_cuda),
+            group=mock_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=500,
+        )
+
+        # Should have broadcasted some tensors
+        assert mock_group.broadcast_count > 0
+        assert len(mock_group.broadcasted_tensors) > 0
+
+    def test_producer_single_large_tensor(self):
+        """Test with a single tensor larger than target size."""
+        # Create a large tensor
+        large_tensor = torch.randn(1000, 1000, dtype=torch.float32).cuda()
+        params = [("large_weight", large_tensor)]
+
+        mock_group = MockCommunicationGroup()
+
+        # Small target size to force the tensor to exceed it
+        packed_broadcast_producer(
+            iterator=iter(params),
+            group=mock_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=100,
+        )
+
+        # Should still broadcast the tensor (at least 1 broadcast)
+        assert mock_group.broadcast_count >= 1
+        assert len(mock_group.broadcasted_tensors) >= 1
+
+        # Verify the total broadcasted size matches the tensor
+        expected_size = large_tensor.numel() * large_tensor.element_size()
+        actual_size = sum(t.numel() for t in mock_group.broadcasted_tensors)
+        assert actual_size == expected_size
+
+    def test_producer_multiple_batches(self):
+        """Test that tensors are properly batched when exceeding target size."""
+        # Create many small tensors
+        params = [
+            (f"weight_{i}", torch.randn(10, 10, dtype=torch.float32).cuda())
+            for i in range(20)
+        ]
+
+        mock_group = MockCommunicationGroup()
+
+        # Small target size to force multiple batches
+        packed_broadcast_producer(
+            iterator=iter(params),
+            group=mock_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=2000,
+        )
+
+        # Should have multiple broadcasts
+        assert mock_group.broadcast_count > 1
+
+        # Total size should match sum of all tensors
+        expected_total = sum(t.numel() * t.element_size() for _, t in params)
+        actual_total = sum(t.numel() for t in mock_group.broadcasted_tensors)
+        assert actual_total == expected_total
+
+    def test_producer_empty_iterator(self):
+        """Test producer handles empty iterator gracefully."""
+        mock_group = MockCommunicationGroup()
+
+        packed_broadcast_producer(
+            iterator=iter([]),
+            group=mock_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=1000,
+        )
+
+        # No broadcasts for empty iterator
+        assert mock_group.broadcast_count == 0
+
+
+# --- Unit Tests: packed_broadcast_consumer ---
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestPackedBroadcastConsumer:
+    """Test packed_broadcast_consumer function."""
+
+    def test_consumer_receives_tensors(self):
+        """Test that consumer receives and unpacks tensors."""
+        params = create_mock_model_params()
+        params_cuda = [(name, tensor.cuda()) for name, tensor in params]
+
+        buffer_size = 2000
+
+        # First, run producer to get the broadcasted tensors
+        producer_group = MockCommunicationGroup()
+
+        packed_broadcast_producer(
+            iterator=iter(params_cuda),
+            group=producer_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=buffer_size,
+        )
+
+        # Now run consumer with the broadcasted tensors
+        consumer_group = MockConsumerCommunicationGroup(
+            producer_group.broadcasted_tensors
+        )
+
+        state_dict_info = create_state_dict_info(params_cuda)
+
+        unpacked_tensors = {}
+
+        def post_unpack_func(tensor_list):
+            for name, tensor in tensor_list:
+                unpacked_tensors[name] = tensor.clone()
+
+        packed_broadcast_consumer(
+            iterator=iter(state_dict_info.items()),
+            group=consumer_group,
+            src=0,
+            post_unpack_func=post_unpack_func,
+            buffer_size_bytes=buffer_size,
+        )
+
+        # Verify all parameters were unpacked
+        assert len(unpacked_tensors) == len(params)
+
+        # Verify each tensor matches the original
+        for name, original_tensor in params_cuda:
+            assert name in unpacked_tensors
+            unpacked = unpacked_tensors[name]
+            assert unpacked.shape == original_tensor.shape
+            assert unpacked.dtype == original_tensor.dtype
+            assert torch.allclose(unpacked, original_tensor, rtol=1e-5, atol=1e-7)
+
+
+# --- Integration Tests: Producer-Consumer Roundtrip ---
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestPackedBroadcastRoundtrip:
+    """Test producer-consumer roundtrip behavior."""
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_roundtrip_different_dtypes(self, dtype):
+        """Test roundtrip with different data types."""
+        params = create_mock_model_params(num_layers=2, dtype=dtype)
+        params_cuda = [(name, tensor.cuda()) for name, tensor in params]
+
+        buffer_size = 1000
+        producer_group = MockCommunicationGroup()
+
+        packed_broadcast_producer(
+            iterator=iter(params_cuda),
+            group=producer_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=buffer_size,
+        )
+
+        consumer_group = MockConsumerCommunicationGroup(
+            producer_group.broadcasted_tensors
+        )
+
+        state_dict_info = create_state_dict_info(params_cuda)
+        unpacked_tensors = {}
+
+        def post_unpack_func(tensor_list):
+            for name, tensor in tensor_list:
+                unpacked_tensors[name] = tensor.clone()
+
+        packed_broadcast_consumer(
+            iterator=iter(state_dict_info.items()),
+            group=consumer_group,
+            src=0,
+            post_unpack_func=post_unpack_func,
+            buffer_size_bytes=buffer_size,
+        )
+
+        # Verify roundtrip preserves data
+        for name, original_tensor in params_cuda:
+            assert name in unpacked_tensors
+            unpacked = unpacked_tensors[name]
+            assert unpacked.dtype == dtype
+            assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
+
+    def test_roundtrip_mixed_dtypes(self):
+        """Test roundtrip with mixed data types."""
+        # Create params with mixed dtypes
+        params = [
+            ("layer1.weight", torch.randn(10, 20, dtype=torch.float32).cuda()),
+            ("layer1.bias", torch.randn(10, dtype=torch.float16).cuda()),
+            ("layer2.weight", torch.randn(20, 30, dtype=torch.bfloat16).cuda()),
+        ]
+
+        buffer_size = 500
+        producer_group = MockCommunicationGroup()
+
+        packed_broadcast_producer(
+            iterator=iter(params),
+            group=producer_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=buffer_size,
+        )
+
+        consumer_group = MockConsumerCommunicationGroup(
+            producer_group.broadcasted_tensors
+        )
+
+        state_dict_info = create_state_dict_info(params)
+        unpacked_tensors = {}
+
+        def post_unpack_func(tensor_list):
+            for name, tensor in tensor_list:
+                unpacked_tensors[name] = tensor.clone()
+
+        packed_broadcast_consumer(
+            iterator=iter(state_dict_info.items()),
+            group=consumer_group,
+            src=0,
+            post_unpack_func=post_unpack_func,
+            buffer_size_bytes=buffer_size,
+        )
+
+        # Verify all params roundtrip correctly with correct dtypes
+        for name, original_tensor in params:
+            assert name in unpacked_tensors
+            unpacked = unpacked_tensors[name]
+            assert unpacked.shape == original_tensor.shape
+            assert unpacked.dtype == original_tensor.dtype
+            assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
+
+    @pytest.mark.parametrize("target_size", [100, 1000, 10000, 100000])
+    def test_roundtrip_different_batch_sizes(self, target_size):
+        """Test roundtrip with different target batch sizes."""
+        params = create_mock_model_params(num_layers=5)
+        params_cuda = [(name, tensor.cuda()) for name, tensor in params]
+
+        producer_group = MockCommunicationGroup()
+
+        packed_broadcast_producer(
+            iterator=iter(params_cuda),
+            group=producer_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=target_size,
+        )
+
+        consumer_group = MockConsumerCommunicationGroup(
+            producer_group.broadcasted_tensors
+        )
+
+        state_dict_info = create_state_dict_info(params_cuda)
+        unpacked_tensors = {}
+
+        def post_unpack_func(tensor_list):
+            for name, tensor in tensor_list:
+                unpacked_tensors[name] = tensor.clone()
+
+        packed_broadcast_consumer(
+            iterator=iter(state_dict_info.items()),
+            group=consumer_group,
+            src=0,
+            post_unpack_func=post_unpack_func,
+            buffer_size_bytes=target_size,
+        )
+
+        # Verify all params roundtrip correctly
+        assert len(unpacked_tensors) == len(params)
+        for name, original_tensor in params_cuda:
+            assert name in unpacked_tensors
+            assert torch.allclose(
+                unpacked_tensors[name], original_tensor, rtol=1e-5, atol=1e-7
+            )
+
+    def test_roundtrip_non_contiguous_tensors(self):
+        """Test roundtrip with non-contiguous tensors from the trainer."""
+        # Create non-contiguous tensors (simulating trainer outputs)
+        # Transposed tensors are non-contiguous
+        weight1 = torch.randn(20, 10, dtype=torch.float32).cuda().T
+        # Sliced tensors with step are non-contiguous
+        weight2 = torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2]
+        # Permuted tensors are non-contiguous
+        weight3 = torch.randn(5, 10, 15, dtype=torch.bfloat16).cuda().permute(2, 0, 1)
+
+        params = [
+            ("layer1.weight", weight1),
+            ("layer2.weight", weight2),
+            ("layer3.weight", weight3),
+        ]
+
+        # Verify tensors are indeed non-contiguous
+        for name, tensor in params:
+            assert not tensor.is_contiguous(), f"{name} should be non-contiguous"
+
+        buffer_size = 500
+        producer_group = MockCommunicationGroup()
+
+        packed_broadcast_producer(
+            iterator=iter(params),
+            group=producer_group,
+            src=0,
+            post_iter_func=lambda x: x[1],
+            buffer_size_bytes=buffer_size,
+        )
+
+        consumer_group = MockConsumerCommunicationGroup(
+            producer_group.broadcasted_tensors
+        )
+
+        state_dict_info = create_state_dict_info(params)
+        unpacked_tensors = {}
+
+        def post_unpack_func(tensor_list):
+            for name, tensor in tensor_list:
+                unpacked_tensors[name] = tensor.clone()
+
+        packed_broadcast_consumer(
+            iterator=iter(state_dict_info.items()),
+            group=consumer_group,
+            src=0,
+            post_unpack_func=post_unpack_func,
+            buffer_size_bytes=buffer_size,
+        )
+
+        # Verify all non-contiguous params roundtrip correctly
+        for name, original_tensor in params:
+            assert name in unpacked_tensors
+            unpacked = unpacked_tensors[name]
+            assert unpacked.shape == original_tensor.shape
+            assert unpacked.dtype == original_tensor.dtype
+            assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)

tests/distributed/test_weight_transfer.py

@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for weight transfer engine backends.
+
+Unit tests for engine classes (parsing, validation, registry).
+Integration test for NCCL weight transfer between processes using Ray.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import ray
+import torch
+
+from vllm.config.parallel import ParallelConfig
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer import WeightTransferEngineFactory
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLWeightTransferEngine,
+    NCCLWeightTransferInitInfo,
+    NCCLWeightTransferUpdateInfo,
+)
+from vllm.utils.network_utils import get_open_port
+
+
+def create_mock_parallel_config(
+    rank: int = 0,
+    world_size: int = 1,
+    dp_rank: int = 0,
+) -> ParallelConfig:
+    """Create a mock ParallelConfig for testing."""
+    config = MagicMock(spec=ParallelConfig)
+    config.rank = rank
+    config.world_size = world_size
+    config.data_parallel_rank = dp_rank
+    return config
+
+
+# --- Unit Tests: NCCLWeightTransferUpdateInfo Validation ---
+
+
+class TestNCCLWeightTransferUpdateInfoValidation:
+    """Test NCCLWeightTransferUpdateInfo dataclass validation."""
+
+    def test_valid_update_info(self):
+        """Test creating valid NCCLWeightTransferUpdateInfo."""
+        info = NCCLWeightTransferUpdateInfo(
+            names=["layer.weight", "layer.bias"],
+            dtype_names=["float32", "float32"],
+            shapes=[[10, 10], [10]],
+        )
+        assert info.names == ["layer.weight", "layer.bias"]
+        assert info.dtype_names == ["float32", "float32"]
+        assert info.shapes == [[10, 10], [10]]
+
+    def test_mismatched_dtype_names_raises(self):
+        """Test that mismatched dtype_names length raises ValueError."""
+        with pytest.raises(ValueError, match="dtype_names"):
+            NCCLWeightTransferUpdateInfo(
+                names=["layer.weight", "layer.bias"],
+                dtype_names=["float32"],  # Only one dtype
+                shapes=[[10, 10], [10]],
+            )
+
+    def test_mismatched_shapes_raises(self):
+        """Test that mismatched shapes length raises ValueError."""
+        with pytest.raises(ValueError, match="shapes"):
+            NCCLWeightTransferUpdateInfo(
+                names=["layer.weight", "layer.bias"],
+                dtype_names=["float32", "float32"],
+                shapes=[[10, 10]],  # Only one shape
+            )
+
+    def test_empty_lists_valid(self):
+        """Test that empty lists are valid."""
+        info = NCCLWeightTransferUpdateInfo(
+            names=[],
+            dtype_names=[],
+            shapes=[],
+        )
+        assert len(info.names) == 0
+
+
+# --- Unit Tests: Engine Parsing ---
+
+
+class TestNCCLEngineParsing:
+    """Test NCCLWeightTransferEngine parsing methods."""
+
+    def test_parse_init_info_valid(self):
+        """Test parsing valid init info dict."""
+        config = WeightTransferConfig(backend="nccl")
+        parallel_config = create_mock_parallel_config()
+        engine = NCCLWeightTransferEngine(config, parallel_config)
+
+        init_info = engine.parse_init_info(
+            {
+                "master_address": "127.0.0.1",
+                "master_port": 12345,
+                "rank_offset": 1,
+                "world_size": 3,
+            }
+        )
+
+        assert isinstance(init_info, NCCLWeightTransferInitInfo)
+        assert init_info.master_address == "127.0.0.1"
+        assert init_info.master_port == 12345
+        assert init_info.rank_offset == 1
+        assert init_info.world_size == 3
+
+    def test_parse_init_info_missing_field_raises(self):
+        """Test parsing init info with missing required field."""
+        config = WeightTransferConfig(backend="nccl")
+        parallel_config = create_mock_parallel_config()
+        engine = NCCLWeightTransferEngine(config, parallel_config)
+
+        with pytest.raises(ValueError, match="Invalid init_info"):
+            engine.parse_init_info(
+                {
+                    "master_address": "127.0.0.1",
+                    # Missing master_port, rank_offset, world_size
+                }
+            )
+
+    def test_parse_update_info_valid(self):
+        """Test parsing valid update info dict."""
+        config = WeightTransferConfig(backend="nccl")
+        parallel_config = create_mock_parallel_config()
+        engine = NCCLWeightTransferEngine(config, parallel_config)
+
+        update_info = engine.parse_update_info(
+            {
+                "names": ["w1", "w2"],
+                "dtype_names": ["float32", "bfloat16"],
+                "shapes": [[100, 100], [50]],
+            }
+        )
+
+        assert isinstance(update_info, NCCLWeightTransferUpdateInfo)
+        assert update_info.names == ["w1", "w2"]
+        assert update_info.dtype_names == ["float32", "bfloat16"]
+        assert update_info.shapes == [[100, 100], [50]]
+
+
+# --- Unit Tests: Engine Registry ---
+
+
+class TestEngineRegistry:
+    """Test weight transfer engine registry."""
+
+    def test_create_engine_nccl(self):
+        """Test factory creates NCCL engine."""
+        config = WeightTransferConfig(backend="nccl")
+        parallel_config = create_mock_parallel_config()
+        engine = WeightTransferEngineFactory.create_engine(config, parallel_config)
+        assert isinstance(engine, NCCLWeightTransferEngine)
+
+    def test_create_engine_invalid_backend(self):
+        """Test factory raises for invalid backend."""
+        config = WeightTransferConfig(backend="invalid")
+        parallel_config = create_mock_parallel_config()
+        with pytest.raises(ValueError, match="Invalid weight transfer backend"):
+            WeightTransferEngineFactory.create_engine(config, parallel_config)
+
+    def test_register_duplicate_raises(self):
+        """Test registering duplicate engine name raises."""
+        with pytest.raises(ValueError, match="already registered"):
+            WeightTransferEngineFactory.register_engine(
+                "nccl", NCCLWeightTransferEngine
+            )
+
+
+# --- Test receive_weights without init raises ---
+
+
+def test_nccl_receive_weights_without_init_raises():
+    """Test that receive_weights raises if init_transfer_engine wasn't called."""
+    if torch.cuda.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    config = WeightTransferConfig(backend="nccl")
+    parallel_config = create_mock_parallel_config()
+    engine = NCCLWeightTransferEngine(config, parallel_config)
+
+    update_info = NCCLWeightTransferUpdateInfo(
+        names=["w"],
+        dtype_names=["float32"],
+        shapes=[[10]],
+    )
+
+    with pytest.raises(RuntimeError, match="not initialized"):
+        engine.receive_weights(update_info, lambda x: None)
+
+
+# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
+
+
+@ray.remote(num_gpus=1)
+def trainer_broadcast_tensor(
+    master_address: str,
+    master_port: int,
+    world_size: int,
+    tensor_shape: list[int],
+    tensor_dtype: str,
+) -> bool:
+    """Trainer task that broadcasts a tensor via NCCL."""
+    import torch
+
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.distributed.utils import StatelessProcessGroup
+
+    # Create process group as rank 0 (trainer)
+    pg = StatelessProcessGroup.create(
+        host=master_address,
+        port=master_port,
+        rank=0,
+        world_size=world_size,
+    )
+    # Ray sets CUDA_VISIBLE_DEVICES, so device 0 is the assigned GPU
+    comm = PyNcclCommunicator(pg, device=0)
+
+    # Create and broadcast the tensor
+    dtype = getattr(torch, tensor_dtype)
+    tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device="cuda:0")
+    comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream())
+    torch.cuda.synchronize()
+
+    return True
+
+
+@ray.remote(num_gpus=1)
+def inference_receive_tensor(
+    master_address: str,
+    master_port: int,
+    world_size: int,
+    tensor_shape: list[int],
+    tensor_dtype: str,
+) -> dict:
+    """Inference task that receives tensor via NCCLWeightTransferEngine."""
+    from unittest.mock import MagicMock
+
+    import torch
+
+    from vllm.config.parallel import ParallelConfig
+    from vllm.config.weight_transfer import WeightTransferConfig
+    from vllm.distributed.weight_transfer.nccl_engine import (
+        NCCLWeightTransferEngine,
+        NCCLWeightTransferInitInfo,
+        NCCLWeightTransferUpdateInfo,
+    )
+
+    # Create engine with mock parallel config
+    config = WeightTransferConfig(backend="nccl")
+    parallel_config = MagicMock(spec=ParallelConfig)
+    parallel_config.rank = 0
+    parallel_config.world_size = 1
+    parallel_config.data_parallel_rank = 0
+
+    engine = NCCLWeightTransferEngine(config, parallel_config)
+
+    # Initialize the engine (joins as rank 1)
+    init_info = NCCLWeightTransferInitInfo(
+        master_address=master_address,
+        master_port=master_port,
+        rank_offset=1,  # Trainer is rank 0, we become rank 1
+        world_size=world_size,
+    )
+    engine.init_transfer_engine(init_info)
+
+    # Receive weights with a no-op load_weights that captures the tensor
+    received_tensors = []
+
+    def noop_load_weights(weights: list[tuple[str, torch.Tensor]]):
+        for name, tensor in weights:
+            # Clone tensor to keep it after engine cleans up
+            received_tensors.append((name, tensor.clone()))
+
+    update_info = NCCLWeightTransferUpdateInfo(
+        names=["test.weight"],
+        dtype_names=[tensor_dtype],
+        shapes=[tensor_shape],
+    )
+    engine.receive_weights(update_info, noop_load_weights)
+    torch.cuda.synchronize()
+
+    # Verify we received the tensor
+    success = False
+    received_shape = None
+    received_sum = None
+
+    if len(received_tensors) == 1:
+        name, tensor = received_tensors[0]
+        received_shape = list(tensor.shape)
+        received_sum = tensor.sum().item()
+        # Check shape matches and values are all 1s (trainer sends ones)
+        if received_shape == tensor_shape:
+            expected_sum = 1.0 * torch.tensor(tensor_shape).prod().item()
+            if abs(received_sum - expected_sum) < 0.01:
+                success = True
+
+    engine.shutdown()
+
+    return {
+        "success": success,
+        "received_shape": received_shape,
+        "received_sum": received_sum,
+    }
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2,
+    reason="Need at least 2 GPUs to run NCCL weight transfer test.",
+)
+def test_nccl_weight_transfer_between_processes():
+    """Test NCCL weight transfer from trainer to inference process using Ray.
+
+    This test verifies that the NCCLWeightTransferEngine can receive
+    tensors broadcast by a trainer process via NCCL.
+    """
+    ray.init(ignore_reinit_error=True)
+
+    master_address = "127.0.0.1"
+    master_port = get_open_port()
+    world_size = 2  # 1 trainer + 1 inference worker
+
+    # Tensor to transfer: 100x100 ones
+    tensor_shape = [100, 100]
+    tensor_dtype = "float32"
+
+    # Start both tasks concurrently - Ray assigns GPUs automatically
+    inference_future = inference_receive_tensor.remote(
+        master_address, master_port, world_size, tensor_shape, tensor_dtype
+    )
+    trainer_future = trainer_broadcast_tensor.remote(
+        master_address, master_port, world_size, tensor_shape, tensor_dtype
+    )
+
+    # Wait for both to complete
+    trainer_result, result = ray.get([trainer_future, inference_future])
+
+    assert trainer_result, "Trainer should complete successfully"
+    assert result["success"], (
+        f"Weight transfer failed. "
+        f"Received shape: {result['received_shape']}, "
+        f"Received sum: {result['received_sum']}"
+    )

tests/entrypoints/openai/responses/test_harmony.py

@@ -62,7 +62,7 @@ async def client(server):
 async def test_basic(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24?",
+        input="What is 123 * 456?",
     )
     assert response is not None
     print("response: ", response)
@@ -74,7 +74,7 @@ async def test_basic(client: OpenAI, model_name: str):
 async def test_basic_with_instructions(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24?",
+        input="What is 123 * 456?",
         instructions="Respond in Korean.",
     )
     assert response is not None
@@ -116,7 +116,7 @@ async def test_chat(client: OpenAI, model_name: str):
             {"role": "system", "content": "Respond in Korean."},
             {"role": "user", "content": "Hello!"},
             {"role": "assistant", "content": "Hello! How can I help you today?"},
-            {"role": "user", "content": "What is 13 * 24? Explain your answer."},
+            {"role": "user", "content": "What is 123 * 456? Explain your answer."},
         ],
     )
     assert response is not None
@@ -131,7 +131,7 @@ async def test_chat_with_input_type(client: OpenAI, model_name: str):
         input=[
             {
                 "role": "user",
-                "content": [{"type": "input_text", "text": "What is 13*24?"}],
+                "content": [{"type": "input_text", "text": "What is 123 * 456?"}],
             },
         ],
     )
@@ -200,7 +200,7 @@ async def test_store(client: OpenAI, model_name: str):
     for store in [True, False]:
         response = await client.responses.create(
             model=model_name,
-            input="What is 13 * 24?",
+            input="What is 123 * 456?",
             store=store,
         )
         assert response is not None
@@ -219,7 +219,7 @@ async def test_store(client: OpenAI, model_name: str):
 async def test_background(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24?",
+        input="What is 123 * 456?",
         background=True,
     )
     assert response is not None
@@ -256,7 +256,7 @@ async def test_background_cancel(client: OpenAI, model_name: str):
 async def test_stateful_multi_turn(client: OpenAI, model_name: str):
     response1 = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24?",
+        input="What is 123 * 456?",
     )
     assert response1 is not None
     assert response1.status == "completed"
@@ -361,7 +361,7 @@ async def test_streaming(client: OpenAI, model_name: str, background: bool):
     # TODO: Add back when web search and code interpreter are available in CI
     prompts = [
         "tell me a story about a cat in 20 words",
-        "What is 13 * 24? Use python to calculate the result.",
+        "What is 123 * 456? Use python to calculate the result.",
         # "When did Jensen found NVIDIA? Search it and answer the year only.",
     ]
 
@@ -976,6 +976,9 @@ async def test_mcp_code_interpreter_streaming(client: OpenAI, model_name: str, s
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.dependency(
+    depends=["test_mcp_code_interpreter_streaming[openai/gpt-oss-20b]"]
+)
 async def test_mcp_tool_multi_turn(client: OpenAI, model_name: str, server):
     """Test MCP tool calling across multiple turns.
 
@@ -1117,8 +1120,10 @@ async def test_function_call_with_previous_input_messages(
         model=model_name,
         input="What is the horoscope for Aquarius today?",
         tools=tools,
+        temperature=0.0,
         extra_body={"enable_response_messages": True},
         stream=True,
+        max_output_tokens=1000,
     )
 
     response = None
@@ -1170,6 +1175,7 @@ async def test_function_call_with_previous_input_messages(
     stream_response_2 = await client.responses.create(
         model=model_name,
         tools=tools,
+        temperature=0.0,
         input="",
         extra_body={
             "previous_input_messages": previous_messages,

tests/entrypoints/openai/responses/test_mcp_tools.py

@@ -160,6 +160,7 @@ class TestMCPEnabled:
                 "No developer messages should be present with valid mcp tool"
             )
 
+    @pytest.mark.flaky(reruns=3)
     @pytest.mark.asyncio
     @pytest.mark.parametrize("model_name", [MODEL_NAME])
     async def test_mcp_tool_with_allowed_tools_star(

tests/entrypoints/openai/responses/test_parsable_context.py

@@ -53,7 +53,7 @@ async def client(server):
 async def test_basic(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24?",
+        input="What is 123 * 456?",
     )
     assert response is not None
     print("response: ", response)
@@ -164,7 +164,7 @@ async def test_function_call_first_turn(client: OpenAI, model_name: str):
 async def test_mcp_tool_call(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24? Use python to calculate the result.",
+        input="What is 123 * 456? Use python to calculate the result.",
         tools=[{"type": "code_interpreter", "container": {"type": "auto"}}],
         extra_body={"enable_response_messages": True},
         temperature=0.0,
@@ -179,12 +179,12 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
     assert response.output[2].type == "reasoning"
     # make sure the correct math is in the final output
     assert response.output[3].type == "message"
-    assert "312" in response.output[3].content[0].text
+    assert "56088" in response.output[3].content[0].text
 
     # test raw input_messages / output_messages
     assert len(response.input_messages) == 1
     assert len(response.output_messages) == 3
-    assert "312" in response.output_messages[2]["message"]
+    assert "56088" in response.output_messages[2]["message"]
 
 
 @pytest.mark.asyncio

tests/entrypoints/openai/responses/test_simple.py

@@ -34,7 +34,7 @@ async def client(server):
 async def test_basic(client: OpenAI, model_name: str):
     response = await client.responses.create(
         model=model_name,
-        input="What is 13 * 24?",
+        input="What is 123 * 456?",
     )
     assert response is not None
     print("response: ", response)

tests/entrypoints/openai/test_openai_schema.py

@@ -139,6 +139,14 @@ def test_openapi_stateless(case: schemathesis.Case):
         # Skip responses API as it is meant to be stateful.
         return
 
+    # Skip weight transfer endpoints as they require special setup
+    # (weight_transfer_config) and are meant to be stateful.
+    if case.operation.path in (
+        "/init_weight_transfer_engine",
+        "/update_weights",
+    ):
+        return
+
     timeout = {
         # requires a longer timeout
         ("POST", "/v1/chat/completions"): LONG_TIMEOUT_SECONDS,

tests/entrypoints/openai/test_serving_tokens.py

@@ -7,6 +7,7 @@ import pytest_asyncio
 from transformers import AutoTokenizer
 
 from vllm.config import ModelConfig
+from vllm.config.utils import getattr_iter
 from vllm.v1.engine.detokenizer import check_stop_strings
 
 from ...utils import RemoteOpenAIServer
@@ -131,7 +132,14 @@ async def test_same_response_as_chat_completions(client, tokenizer, messages):
             # Post-EOS generation is undefined and may differ
             eos_tokens = {
                 tokenizer.eos_token_id,
-                *tokenizer.additional_special_tokens_ids,
+                *getattr_iter(
+                    tokenizer,
+                    [
+                        "extra_special_tokens_ids",  # Transformers v5
+                        "additional_special_tokens_ids",  # Transformers v4
+                    ],
+                    [],
+                ),
             }
             # Find first EOS in generated tokens
             eos_pos = None

tests/entrypoints/pooling/classify/test_online.py

@@ -469,6 +469,4 @@ async def test_pooling_not_supported(
         },
     )
     assert response.json()["error"]["type"] == "BadRequestError"
-    assert response.json()["error"]["message"].startswith(
-        f"Task {task} is not supported"
-    )
+    assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

tests/entrypoints/pooling/embed/test_online.py

@@ -757,6 +757,4 @@ async def test_pooling_not_supported(
         },
     )
     assert response.json()["error"]["type"] == "BadRequestError"
-    assert response.json()["error"]["message"].startswith(
-        f"Task {task} is not supported"
-    )
+    assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

tests/entrypoints/pooling/score/test_online_colbert.py

@@ -138,17 +138,17 @@ def test_colbert_token_embed(server: RemoteOpenAIServer, model_name: str):
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 def test_colbert_embed_not_supported(server: RemoteOpenAIServer, model_name: str):
     """Test that ColBERT model does not support 'embed' task."""
+    task = "embed"
     text = "What is the capital of France?"
 
-    pooling_response = requests.post(
+    response = requests.post(
         server.url_for("pooling"),
         json={
             "model": model_name,
             "input": text,
-            "task": "embed",
+            "task": task,
         },
     )
 
-    # Should return error
-    assert pooling_response.status_code == 400
-    assert "Task embed is not supported" in pooling_response.text
+    assert response.json()["error"]["type"] == "BadRequestError"
+    assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

tests/entrypoints/pooling/score/test_online_rerank.py

@@ -232,6 +232,4 @@ async def test_pooling_not_supported(
         },
     )
     assert response.json()["error"]["type"] == "BadRequestError"
-    assert response.json()["error"]["message"].startswith(
-        f"Task {task} is not supported"
-    )
+    assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

tests/entrypoints/weight_transfer/__init__.py

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

tests/entrypoints/weight_transfer/test_weight_transfer_llm.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for weight transfer APIs via LLM class.
+
+These tests use a mock weight transfer engine to verify that the API
+calls the correct methods with the right arguments, without requiring
+actual NCCL communication.
+"""
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm import LLM
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferEngine,
+    WeightTransferInitInfo,
+    WeightTransferInitRequest,
+    WeightTransferUpdateInfo,
+    WeightTransferUpdateRequest,
+)
+
+from ...utils import create_new_process_for_each_test
+
+# Use a tiny model for fast testing
+MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
+
+
+# --- Mock Weight Transfer Engine ---
+
+
+@dataclass
+class MockInitInfo(WeightTransferInitInfo):
+    """Mock initialization info."""
+
+    test_param: str = "test"
+
+
+@dataclass
+class MockUpdateInfo(WeightTransferUpdateInfo):
+    """Mock update info."""
+
+    names: list[str] | None = None
+    dtype_names: list[str] | None = None
+    shapes: list[list[int]] | None = None
+
+
+class MockWeightTransferEngine(WeightTransferEngine[MockInitInfo, MockUpdateInfo]):
+    """Mock weight transfer engine that tracks method calls."""
+
+    init_info_cls = MockInitInfo
+    update_info_cls = MockUpdateInfo
+
+    # Class-level tracking for verification across processes
+    init_transfer_engine_called: bool = False
+    receive_weights_called: bool = False
+    shutdown_called: bool = False
+    last_init_info: MockInitInfo | None = None
+    last_update_info: MockUpdateInfo | None = None
+
+    def __init__(self, config, parallel_config):
+        super().__init__(config, parallel_config)
+        # Reset tracking on init
+        MockWeightTransferEngine.init_transfer_engine_called = False
+        MockWeightTransferEngine.receive_weights_called = False
+        MockWeightTransferEngine.shutdown_called = False
+        MockWeightTransferEngine.last_init_info = None
+        MockWeightTransferEngine.last_update_info = None
+
+    def init_transfer_engine(self, init_info: MockInitInfo) -> None:
+        MockWeightTransferEngine.init_transfer_engine_called = True
+        MockWeightTransferEngine.last_init_info = init_info
+
+    def receive_weights(
+        self,
+        update_info: MockUpdateInfo,
+        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
+    ) -> None:
+        MockWeightTransferEngine.receive_weights_called = True
+        MockWeightTransferEngine.last_update_info = update_info
+        # Simulate loading weights by calling load_weights with empty list
+        # (In real implementation, this would receive and load actual weights)
+        load_weights([])
+
+    def shutdown(self) -> None:
+        MockWeightTransferEngine.shutdown_called = True
+
+
+def mock_create_engine(config, parallel_config):
+    """Mock factory function that returns our mock engine."""
+    return MockWeightTransferEngine(config, parallel_config)
+
+
+# --- Tests ---
+
+
+@create_new_process_for_each_test()
+def test_get_world_size_tp1():
+    """Test world_size is correctly configured for TP=1."""
+    if torch.cuda.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    llm = LLM(
+        model=MODEL_NAME,
+        enforce_eager=True,
+        load_format="dummy",
+        tensor_parallel_size=1,
+        weight_transfer_config=WeightTransferConfig(backend="nccl"),
+    )
+
+    world_size = llm.llm_engine.vllm_config.parallel_config.world_size
+    assert world_size == 1
+
+
+@create_new_process_for_each_test()
+def test_init_weight_transfer_engine_calls_engine():
+    """Test that init_weight_transfer_engine calls the engine's
+    init_transfer_engine method."""
+    if torch.cuda.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    # Enable insecure serialization to allow pickling functions for collective_rpc
+    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
+    with patch(
+        "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
+        mock_create_engine,
+    ):
+        llm = LLM(
+            model=MODEL_NAME,
+            enforce_eager=True,
+            load_format="dummy",
+            tensor_parallel_size=1,
+            weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        )
+
+        # Verify engine was created
+        def check_engine_exists(self):
+            return self.weight_transfer_engine is not None
+
+        results = llm.collective_rpc(check_engine_exists)
+        assert all(results), "Weight transfer engine should be initialized"
+
+        # Call init_weight_transfer_engine
+        llm.init_weight_transfer_engine(
+            WeightTransferInitRequest(init_info={"test_param": "hello"})
+        )
+
+        # Verify init_transfer_engine was called on the engine
+        def check_init_called(self):
+            engine = self.weight_transfer_engine
+            return (
+                engine.init_transfer_engine_called,
+                engine.last_init_info.test_param if engine.last_init_info else None,
+            )
+
+        results = llm.collective_rpc(check_init_called)
+        for called, param in results:
+            assert called, "init_transfer_engine should have been called"
+            assert param == "hello", f"Expected 'hello', got {param}"
+
+
+@create_new_process_for_each_test()
+def test_update_weights_calls_engine():
+    """Test that update_weights calls the engine's receive_weights method."""
+    if torch.cuda.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    # Enable insecure serialization to allow pickling functions for collective_rpc
+    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
+    with patch(
+        "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
+        mock_create_engine,
+    ):
+        llm = LLM(
+            model=MODEL_NAME,
+            enforce_eager=True,
+            load_format="dummy",
+            tensor_parallel_size=1,
+            weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        )
+
+        # First init the weight transfer
+        llm.init_weight_transfer_engine(
+            WeightTransferInitRequest(init_info={"test_param": "init"})
+        )
+
+        # Call update_weights
+        test_names = ["layer.weight", "layer.bias"]
+        test_dtypes = ["float32", "float32"]
+        test_shapes = [[10, 10], [10]]
+
+        llm.update_weights(
+            WeightTransferUpdateRequest(
+                update_info={
+                    "names": test_names,
+                    "dtype_names": test_dtypes,
+                    "shapes": test_shapes,
+                }
+            )
+        )
+
+        # Verify receive_weights was called with correct info
+        def check_update_called(self):
+            engine = self.weight_transfer_engine
+            if not engine.receive_weights_called:
+                return False, None, None, None
+            info = engine.last_update_info
+            return (True, info.names, info.dtype_names, info.shapes)
+
+        results = llm.collective_rpc(check_update_called)
+        for called, names, dtypes, shapes in results:
+            assert called, "receive_weights should have been called"
+            assert names == test_names
+            assert dtypes == test_dtypes
+            assert shapes == test_shapes
+
+
+@create_new_process_for_each_test()
+def test_full_weight_transfer_flow():
+    """Test the complete weight transfer flow: init -> update."""
+    if torch.cuda.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    # Enable insecure serialization to allow pickling functions for collective_rpc
+    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
+    with patch(
+        "vllm.v1.worker.gpu_worker.WeightTransferEngineFactory.create_engine",
+        mock_create_engine,
+    ):
+        llm = LLM(
+            model=MODEL_NAME,
+            enforce_eager=True,
+            load_format="dummy",
+            tensor_parallel_size=1,
+            weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        )
+
+        # Step 1: Initialize
+        llm.init_weight_transfer_engine(
+            WeightTransferInitRequest(init_info={"test_param": "flow_test"})
+        )
+
+        # Step 2: Update weights
+        llm.update_weights(
+            WeightTransferUpdateRequest(
+                update_info={
+                    "names": ["test.weight"],
+                    "dtype_names": ["bfloat16"],
+                    "shapes": [[100, 100]],
+                }
+            )
+        )
+
+        # Verify the full flow completed
+        def check_flow(self):
+            engine = self.weight_transfer_engine
+            return {
+                "init_called": engine.init_transfer_engine_called,
+                "update_called": engine.receive_weights_called,
+                "init_param": (
+                    engine.last_init_info.test_param if engine.last_init_info else None
+                ),
+                "update_names": (
+                    engine.last_update_info.names if engine.last_update_info else None
+                ),
+            }
+
+        results = llm.collective_rpc(check_flow)
+        for result in results:
+            assert result["init_called"], "init_transfer_engine should be called"
+            assert result["update_called"], "receive_weights should be called"
+            assert result["init_param"] == "flow_test"
+            assert result["update_names"] == ["test.weight"]
+
+
+@create_new_process_for_each_test()
+def test_weight_transfer_config_backend():
+    """Test that WeightTransferConfig backend is properly configured."""
+    if torch.cuda.device_count() < 1:
+        pytest.skip("Need at least 1 GPU for this test")
+
+    # Test with nccl backend
+    llm = LLM(
+        model=MODEL_NAME,
+        enforce_eager=True,
+        load_format="dummy",
+        tensor_parallel_size=1,
+        weight_transfer_config=WeightTransferConfig(backend="nccl"),
+    )
+
+    config = llm.llm_engine.vllm_config.weight_transfer_config
+    assert config.backend == "nccl"

tests/kernels/moe/modular_kernel_tools/common.py

@@ -620,6 +620,7 @@ def make_modular_kernel(
     modular_kernel = mk.FusedMoEModularKernel(
         prepare_finalize=prepare_finalize,
         fused_experts=fused_experts,
+        inplace=False,
     )
 
     return modular_kernel

tests/kernels/moe/test_batched_deepgemm.py

@@ -74,7 +74,11 @@ def test_batched_deepgemm_vs_triton(
         quant_config=quant_config,
         moe_config=make_dummy_moe_config(),
     )
-    mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts)
+    mk_triton = FusedMoEModularKernel(
+        prep_finalize,
+        triton_experts,
+        inplace=False,
+    )
 
     out_triton = mk_triton(
         hidden_states=a,
@@ -82,7 +86,6 @@ def test_batched_deepgemm_vs_triton(
         w2=w2,
         topk_weights=topk_weights,
         topk_ids=topk_ids,
-        inplace=False,
         global_num_experts=E,
     )
 
@@ -93,7 +96,11 @@ def test_batched_deepgemm_vs_triton(
         quant_config=quant_config,
         moe_config=make_dummy_moe_config(),
     )
-    mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts)
+    mk_deepgemm = FusedMoEModularKernel(
+        prep_finalize,
+        deepgemm_experts,
+        inplace=False,
+    )
 
     out_deepgemm = mk_deepgemm(
         hidden_states=a,
@@ -101,7 +108,6 @@ def test_batched_deepgemm_vs_triton(
         w2=w2,
         topk_weights=topk_weights,
         topk_ids=topk_ids,
-        inplace=False,
         global_num_experts=E,
     )
 

tests/kernels/moe/test_block_fp8.py

@@ -9,6 +9,7 @@ from tests.kernels.moe.utils import (
     make_dummy_moe_config,
     make_test_quant_config,
     make_test_weights,
+    modular_triton_fused_moe,
 )
 from tests.kernels.quant_utils import (
     native_per_token_group_quant_fp8,
@@ -26,9 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import (
 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
     _valid_deep_gemm_shape,
 )
-from vllm.model_executor.layers.fused_moe.fused_moe import (
-    modular_triton_fused_moe,
-)
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
     MoEPrepareAndFinalizeNoEP,
 )
@@ -261,6 +259,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)
             moe_config=make_dummy_moe_config(),
             quant_config=quant_config,
         ),
+        inplace=False,
     )
 
     def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids):

tests/kernels/moe/test_cutlass_moe.py

@@ -207,6 +207,7 @@ def run_with_expert_maps(
                 ),
                 quant_config=new_quant_config,
             ),
+            inplace=False,
         )
         out_tensor = out_tensor + kernel(**kwargs)
 
@@ -266,6 +267,7 @@ def run_8_bit(
                 ),
                 quant_config=quant_config,
             ),
+            inplace=False,
         )
         return kernel(**kwargs)
 

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -194,8 +194,11 @@ def make_ll_modular_kernel(
         quant_config=quant_config,
         moe_config=make_dummy_moe_config(),
     )
-    mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)
-    return mk
+    return FusedMoEModularKernel(
+        prepare_finalize=a2a,
+        fused_experts=fused_experts,
+        inplace=False,
+    )
 
 
 def make_ht_modular_kernel(
@@ -224,8 +227,11 @@ def make_ht_modular_kernel(
         moe_config=make_dummy_moe_config(),
         quant_config=quant_config,
     )
-    mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)
-    return mk
+    return FusedMoEModularKernel(
+        prepare_finalize=a2a,
+        fused_experts=fused_experts,
+        inplace=False,
+    )
 
 
 def make_modular_kernel(
@@ -318,7 +324,6 @@ def deepep_deepgemm_moe_impl(
             w2=w2,
             topk_weights=test_tensors.topk_weights,
             topk_ids=test_tensors.topk,
-            inplace=False,
             activation="silu",
             global_num_experts=num_experts,
             expert_map=build_expert_map(),

tests/kernels/moe/test_deepep_moe.py

@@ -179,7 +179,11 @@ def make_modular_kernel(
             quant_config=quant_config,
         )
 
-    mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)
+    mk = FusedMoEModularKernel(
+        prepare_finalize=a2a,
+        fused_experts=fused_experts,
+        inplace=False,
+    )
     return mk
 
 
@@ -256,7 +260,6 @@ def deep_ep_moe_impl(
             w2=w2,
             topk_weights=topk_weights_chunk,
             topk_ids=topk_chunk,
-            inplace=False,
             activation="silu",
             global_num_experts=num_experts,
             expert_map=build_expert_map(),

tests/kernels/moe/test_deepgemm.py

@@ -115,6 +115,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
             moe_config=make_dummy_moe_config(),
             quant_config=quant_config,
         ),
+        inplace=False,
     )
 
     # triton reference
@@ -135,7 +136,6 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
         w2=w2,
         topk_weights=topk_weights,
         topk_ids=topk_ids,
-        inplace=False,
     )
     diff = calc_diff(out_deepgemm, out_triton)
     assert diff < 0.001, f"Diff exceeded 1%: {diff}"

tests/kernels/moe/test_flashinfer.py

@@ -301,6 +301,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
                 moe_config=moe_config,
                 quant_config=quant_config,
             ),
+            inplace=False,
         )
 
         flashinfer_cutlass_output = kernel(
@@ -309,7 +310,6 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
             td.layer.w2_weight,
             topk_weights,
             topk_ids,
-            inplace=False,
             activation=activation,
             global_num_experts=e,
             expert_map=None,

tests/kernels/moe/test_flashinfer_moe.py

@@ -108,6 +108,7 @@ def test_flashinfer_fp4_moe_no_graph(
         flashinfer_experts = FusedMoEModularKernel(
             MoEPrepareAndFinalizeNoEP(),
             FlashInferExperts(moe_config=moe_config, quant_config=quant_config),
+            inplace=False,
         )
 
         fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]

tests/kernels/moe/test_modular_oai_triton_moe.py

@@ -180,7 +180,11 @@ def oai_triton_moe_impl(
     else:
         fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config)
 
-    mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts)
+    mk = FusedMoEModularKernel(
+        MoEPrepareAndFinalizeNoEP(),
+        fused_experts,
+        inplace=False,
+    )
 
     return mk.forward(
         hidden_states=x,
@@ -188,7 +192,6 @@ def oai_triton_moe_impl(
         w2=w2,
         topk_weights=topk_weights,
         topk_ids=topk_ids,
-        inplace=True,
         activation="swigluoai",
         global_num_experts=num_experts,
         expert_map=None,

tests/kernels/moe/test_moe.py

@@ -18,7 +18,11 @@ from transformers import MixtralConfig
 from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
 
 import vllm.model_executor.layers.fused_moe  # noqa
-from tests.kernels.moe.utils import fused_moe, make_dummy_moe_config
+from tests.kernels.moe.utils import (
+    fused_moe,
+    make_dummy_moe_config,
+    modular_triton_fused_moe,
+)
 from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig, set_current_vllm_config
@@ -36,9 +40,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
     batched_fused_marlin_moe,
     fused_marlin_moe,
 )
-from vllm.model_executor.layers.fused_moe.fused_moe import (
-    modular_triton_fused_moe,
-)
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     marlin_permute_bias,
 )

tests/kernels/moe/test_nvfp4_moe.py

@@ -95,6 +95,7 @@ def test_cutlass_fp4_moe_no_graph(
                 moe_config=make_dummy_moe_config(),
                 quant_config=quant_config,
             ),
+            inplace=False,
         )
 
         cutlass_output = kernel(

tests/kernels/moe/test_pplx_cutlass_moe.py

@@ -172,6 +172,7 @@ def pplx_cutlass_moe(
     fused_cutlass_experts = FusedMoEModularKernel(
         prepare_finalize,
         experts,
+        inplace=False,
     )
 
     a_chunk = chunk_by_rank(a, rank, world_size).to(device)

tests/kernels/moe/test_pplx_moe.py

@@ -592,6 +592,7 @@ def pplx_moe(
         prepare_finalize,
         experts,
         shared_experts,
+        inplace=False,
     )
 
     # Note: for now use_compile will error out if the problem size is

tests/kernels/moe/test_rocm_aiter_topk.py

@@ -10,22 +10,28 @@
 # and the platform is not ROCm.
 
 import importlib.util
+import os
 
 import pytest
 import torch
 
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip("This test can only run on ROCm.", allow_module_level=True)
+
+# This environment variable must be set so ops will be registered.
+os.environ["VLLM_ROCM_USE_AITER"] = "1"
+
 # this import statement is needed to ensure the ops are registered
 import vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe  # noqa: F401
-from vllm.platforms import current_platform
 
 # need to import once to ensure the ops are registered
 # Check if aiter package is installed
 aiter_available = importlib.util.find_spec("aiter") is not None
 
-pytestmark = pytest.mark.skipif(
-    not (current_platform.is_rocm() and aiter_available),
-    reason="AITER ops are only available on ROCm with aiter package installed",
-)
+if not aiter_available:
+    pytest.skip("These tests require AITER to run.", allow_module_level=True)
 
 
 def test_rocm_aiter_biased_grouped_topk_custom_op_registration():

tests/kernels/moe/utils.py

@@ -7,7 +7,11 @@ import vllm._custom_ops as ops
 from tests.kernels.quant_utils import per_block_cast_to_int8
 from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
+from vllm.model_executor.layers.fused_moe import (
+    TritonExperts,
+    fused_experts,
+    fused_topk,
+)
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEParallelConfig,
@@ -20,6 +24,9 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
     NaiveBatchedExperts,
 )
 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNoEP,
+)
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 from vllm.utils.deep_gemm import per_block_cast_to_fp8
 from vllm.utils.math_utils import round_up
@@ -116,6 +123,7 @@ def batched_moe(
             quant_config=quant_config,
             moe_config=make_dummy_moe_config(),
         ),
+        inplace=False,
     )
 
     return fused_experts(a, w1, w2, topk_weight, topk_ids)
@@ -157,6 +165,7 @@ def naive_batched_moe(
             quant_config=quant_config,
             moe_config=make_dummy_moe_config(),
         ),
+        inplace=False,
     )
 
     return fused_experts(a, w1, w2, topk_weight, topk_ids)
@@ -554,3 +563,16 @@ def make_shared_experts(
         return RealMLP(K, N, w1, w2, "silu", quant_config, w1_s=w1_s, w2_s=w2_s)
     finally:
         torch.set_default_dtype(old_dtype)
+
+
+def modular_triton_fused_moe(
+    moe_config: FusedMoEConfig,
+    quant_config: FusedMoEQuantConfig,
+    shared_experts: torch.nn.Module | None = None,
+) -> FusedMoEModularKernel:
+    return FusedMoEModularKernel(
+        MoEPrepareAndFinalizeNoEP(),
+        TritonExperts(moe_config, quant_config),
+        shared_experts,
+        inplace=False,
+    )

tests/kernels/quantization/test_rocm_skinny_gemms.py

@@ -73,21 +73,40 @@ NKM_FACTORS_WVSPLITKRC = [
 NKM_FACTORS_WVSPLITK_FP8 = [
     # FP8-specific cases with K % 16 == 0
     (1, 16, 16),
+    (1, 32, 16 + 16),
     (1, 64, 64),
+    (1, 64, 64 + 16),
+    (1, 64 + 16, 64),
+    (1, 64 + 16, 64 + 16),
+    (4, 64, 64),
+    (4, 64, 64 + 16),
+    (4, 64 + 16, 64),
+    (4, 64 + 16, 64 + 16),
     (2, 512, 512),
+    (3, 512, 512),
+    (3, 512, 512 + 16),
+    (4, 512, 512),
     (3, 2048, 2048),
+    (3, 2048, 2048 + 16),
+    (4, 2048 + 16, 2048),
+    (4, 2048 + 16, 2048 + 16),
     (4, 4096, 4096),
     (4, 16400, 2048),
+    (4, 16400, 2048 + 16),
     # Extended FP8 dimensions not covered by WVSPLITK
     (1, 14336, 1024),
     (2, 24576, 2048),
     (4, 32768, 28672),
+    (4, 32768 * 2, 28672),
+    (4, 32768 * 2, 28672 + 16),
+    (4, 32768 * 2 + 16, 28672),
+    (4, 32768 * 2 + 16, 28672 + 16),
 ]
 
 SEEDS = [0]
 
 
-def pad_weights_fp8(weight):
+def pad_fp8(weight):
     num_pad = 256 // weight.element_size()
     import torch.nn.functional as F
 
@@ -195,72 +214,41 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
     assert torch.allclose(out, ref_out, rtol=0.01)
 
 
+@pytest.mark.parametrize("xnorm", [False, True])
 @pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("padded", [False, True])
+@pytest.mark.parametrize("padded_a", [False, True])
+@pytest.mark.parametrize("padded_b", [False, True])
+@pytest.mark.parametrize("biased", [False, True])
 @pytest.mark.skipif(
     not (current_platform.is_rocm() and current_platform.supports_fp8()),
     reason="only test for rocm fp8",
 )
-def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
+def test_rocm_wvsplitk_fp8_kernel(
+    xnorm, n, k, m, dtype, seed, padded_a, padded_b, biased
+):
     torch.manual_seed(seed)
 
-    A = torch.rand(n, k, device="cuda") - 0.5
-    B = torch.rand(m, k, device="cuda") - 0.5
+    xavier = math.sqrt(2 / k) if xnorm else 1  # normalize to avoid large deltas
+    A = (torch.rand(n, k, device="cuda") * 2 - 1) * xavier
+    B = (torch.rand(m, k, device="cuda") * 2 - 1) * xavier
 
     A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
     B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
-    if padded:
-        B = pad_weights_fp8(B)
+    if padded_b:
+        B = pad_fp8(B)
+    if padded_a:
+        A = pad_fp8(A)
 
-    ref_out = torch._scaled_mm(
-        A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b
-    )
-    out = ops.wvSplitKQ(
-        B,
-        A,
-        dtype,
-        scale_a,
-        scale_b,
-        get_cu_count(),
-    )
-
-    assert torch.allclose(out, ref_out, rtol=0.01)
-
-
-@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("seed", SEEDS)
-@pytest.mark.parametrize("padded", [False, True])
-@pytest.mark.skipif(
-    not (current_platform.is_rocm() and current_platform.supports_fp8()),
-    reason="only test for rocm fp8",
-)
-def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded):
-    torch.manual_seed(seed)
-
-    xavier = math.sqrt(2 / k)  # normalize to avoid large output-bias deltas
-    A = (torch.rand(n, k, device="cuda") - 0.5) * xavier
-    B = (torch.rand(m, k, device="cuda") - 0.5) * xavier
-    BIAS = torch.rand(m, dtype=dtype, device="cuda") - 0.5
-
-    A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
-    B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
-    if padded:
-        B = pad_weights_fp8(B)
+    BIAS = None if (not biased) else (torch.rand(m, dtype=dtype, device="cuda") * 2 - 1)
 
     ref_out = torch._scaled_mm(
         A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
     )
-    out = ops.wvSplitKQ(
-        B,
-        A,
-        dtype,
-        scale_a,
-        scale_b,
-        get_cu_count(),
-        BIAS,
-    )
+    out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, get_cu_count(), BIAS)
 
-    assert torch.allclose(out, ref_out, rtol=0.01)
+    if xnorm:
+        assert torch.allclose(out, ref_out, atol=1e-3, rtol=1e-8)
+    else:
+        assert torch.allclose(out, ref_out, 0.01)

tests/models/language/generation/test_hybrid.py

@@ -8,6 +8,7 @@ import pytest
 from tests.models.registry import HF_EXAMPLE_MODELS
 from tests.utils import multi_gpu_test
 from vllm.engine.arg_utils import EngineArgs
+from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 
@@ -577,6 +578,10 @@ def test_apc_multiple_prompts_all_cached_outputs(
         model, max_model_len, tensor_parallel_size=tensor_parallel_size
     )
     vllm_runner_kwargs["mamba_ssm_cache_dtype"] = "float32"
+    # Reduce the effects of batch variance on ROCm since batch invariance is not
+    # yet supported. See: https://github.com/vllm-project/vllm/issues/27433
+    if current_platform.is_rocm():
+        vllm_runner_kwargs["max_num_seqs"] = 4
 
     vllm_outputs_no_cache, _ = _get_vLLM_output(
         vllm_runner, vllm_runner_kwargs, generated_prompts, max_tokens, num_logprobs

tests/models/language/pooling/test_bge_m3.py

@@ -136,6 +136,16 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
     )
 
 
+@pytest.mark.asyncio
+async def test_bge_m3_api_server_sparse_embedding_corner_case(
+    client: openai.AsyncOpenAI,
+):
+    embeddings = await sparse_embeddings(client, ["Hi"])
+    assert len(embeddings) == 1
+    assert 2673 in embeddings[0]
+    assert embeddings[0][2673] == pytest.approx(0.26710861921310425, rel=0.01)
+
+
 # https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163
 def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
     token_scores = torch.einsum("in,jn->ij", q_reps, p_reps)

tests/models/multimodal/pooling/test_jinavl_reranker.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import cast
 
 import pytest
 import transformers
@@ -117,7 +116,7 @@ def _normalize_image(image_val: str) -> str:
 
 def create_score_multimodal_param(
     content_parts: list[dict],
-) -> ScoreMultiModalParam:
+) -> list[ScoreMultiModalParam]:
     """
     Create a ScoreMultiModalParam from a list of content dictionaries.
 
@@ -152,7 +151,7 @@ def create_score_multimodal_param(
                     )
                 )
 
-    return ScoreMultiModalParam(content=formatted_content)
+    return [ScoreMultiModalParam(content=[content]) for content in formatted_content]
 
 
 def _run_vllm(
@@ -198,23 +197,7 @@ def _run_hf(
     else:
         raise ValueError("Unsupported query format")
 
-    # Separate documents by type
-    text_docs: list[str] = []
-    image_docs: list[str] = []
-    text_indices: list[int] = []
-    image_indices: list[int] = []
-
-    for idx, doc in enumerate(document_strs):
-        if "text" in doc:
-            text_docs.append(doc["text"])
-            text_indices.append(idx)
-        elif "image" in doc:
-            image_docs.append(_normalize_image(doc["image"]))
-            image_indices.append(idx)
-        else:
-            raise ValueError(f"Unsupported document format at index {idx}")
-
-    scores: list[None | float] = [None] * len(document_strs)
+    scores: list[float] = []
 
     with hf_runner(
         model,
@@ -223,30 +206,24 @@ def _run_hf(
         auto_cls=AutoModel,
         model_kwargs={"key_mapping": CHECKPOINT_TO_HF_MAPPER},
     ) as hf_model:
-        # Score text documents
-        if text_docs:
-            text_scores = hf_model.model.compute_score(
-                [[query_data, d] for d in text_docs],
-                max_length=2048,
-                query_type=query_type,
-                doc_type="text",
-            )
-            for i, s in zip(text_indices, text_scores):
-                scores[i] = s
-
-        # Score image documents
-        if image_docs:
-            image_scores = hf_model.model.compute_score(
-                [[query_data, d] for d in image_docs],
-                max_length=2048,
-                query_type=query_type,
-                doc_type="image",
-            )
-            for i, s in zip(image_indices, image_scores):
-                scores[i] = s
-
-    assert all(s is not None for s in scores)
-    return cast(list[float], scores)
+        for doc in document_strs:
+            if "text" in doc:
+                score = hf_model.model.compute_score(
+                    [[query_data, doc["text"]]],
+                    max_length=2048,
+                    query_type=query_type,
+                    doc_type="text",
+                )
+                scores.append(score)
+            elif "image" in doc:
+                score = hf_model.model.compute_score(
+                    [[query_data, doc["image"]]],
+                    max_length=2048,
+                    query_type=query_type,
+                    doc_type="image",
+                )
+                scores.append(score)
+    return scores
 
 
 def _run_test(

tests/models/test_terratorch.py

@@ -5,8 +5,10 @@ import pytest
 import torch
 
 from tests.conftest import VllmRunner
+from tests.utils import create_new_process_for_each_test
 
 
+@create_new_process_for_each_test()  # Memory is not cleaned up properly otherwise
 @pytest.mark.parametrize(
     "model",
     [
@@ -22,10 +24,14 @@ def test_inference(
     location_coords = torch.full((1, 2), 1.0, dtype=torch.float16)
     prompt = dict(
         prompt_token_ids=[1],
-        multi_modal_data=dict(
-            pixel_values=pixel_values, location_coords=location_coords
-        ),
+        multi_modal_data={
+            "image": {
+                "pixel_values": pixel_values,
+                "location_coords": location_coords,
+            }
+        },
     )
+
     with vllm_runner(
         model,
         runner="pooling",

tests/multimodal/media/test_connector.py

@@ -0,0 +1,320 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import base64
+import mimetypes
+import os
+from tempfile import NamedTemporaryFile, TemporaryDirectory
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image, ImageChops
+
+from vllm.multimodal.image import convert_image_mode
+from vllm.multimodal.inputs import PlaceholderRange
+from vllm.multimodal.media import MediaConnector
+
+# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
+TEST_IMAGE_ASSETS = [
+    "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
+    "Grayscale_8bits_palette_sample_image.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
+    "1280px-Venn_diagram_rgb.svg.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
+    "RGBA_comp.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
+]
+
+TEST_VIDEO_URLS = [
+    "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4",
+    "https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi",
+]
+
+
+@pytest.fixture(scope="module")
+def url_images(local_asset_server) -> dict[str, Image.Image]:
+    return {
+        image_url: local_asset_server.get_image_asset(image_url)
+        for image_url in TEST_IMAGE_ASSETS
+    }
+
+
+def get_supported_suffixes() -> tuple[str, ...]:
+    # We should at least test the file types mentioned in GPT-4 with Vision
+    OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif")
+
+    # Additional file types that are supported by us
+    EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff")
+
+    return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES
+
+
+def _image_equals(a: Image.Image, b: Image.Image) -> bool:
+    return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
+async def test_fetch_image_http(image_url: str):
+    connector = MediaConnector()
+
+    image_sync = connector.fetch_image(image_url)
+    image_async = await connector.fetch_image_async(image_url)
+    assert _image_equals(image_sync, image_async)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS)
+@pytest.mark.parametrize("suffix", get_supported_suffixes())
+async def test_fetch_image_base64(
+    url_images: dict[str, Image.Image], raw_image_url: str, suffix: str
+):
+    connector = MediaConnector(
+        # Domain restriction should not apply to data URLs.
+        allowed_media_domains=[
+            "www.bogotobogo.com",
+            "github.com",
+        ]
+    )
+    url_image = url_images[raw_image_url]
+
+    try:
+        mime_type = Image.MIME[Image.registered_extensions()[suffix]]
+    except KeyError:
+        try:
+            mime_type = mimetypes.types_map[suffix]
+        except KeyError:
+            pytest.skip("No MIME type")
+
+    with NamedTemporaryFile(suffix=suffix) as f:
+        try:
+            url_image.save(f.name)
+        except Exception as e:
+            if e.args[0] == "cannot write mode RGBA as JPEG":
+                pytest.skip("Conversion not supported")
+
+            raise
+
+        base64_image = base64.b64encode(f.read()).decode("utf-8")
+        data_url = f"data:{mime_type};base64,{base64_image}"
+
+        data_image_sync = connector.fetch_image(data_url)
+        if _image_equals(url_image, Image.open(f)):
+            assert _image_equals(url_image, data_image_sync)
+        else:
+            pass  # Lossy format; only check that image can be opened
+
+        data_image_async = await connector.fetch_image_async(data_url)
+        assert _image_equals(data_image_sync, data_image_async)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
+async def test_fetch_image_local_files(image_url: str):
+    connector = MediaConnector()
+
+    with TemporaryDirectory() as temp_dir:
+        local_connector = MediaConnector(allowed_local_media_path=temp_dir)
+
+        origin_image = connector.fetch_image(image_url)
+        origin_image.save(
+            os.path.join(temp_dir, os.path.basename(image_url)),
+            quality=100,
+            icc_profile=origin_image.info.get("icc_profile"),
+        )
+
+        image_async = await local_connector.fetch_image_async(
+            f"file://{temp_dir}/{os.path.basename(image_url)}"
+        )
+        image_sync = local_connector.fetch_image(
+            f"file://{temp_dir}/{os.path.basename(image_url)}"
+        )
+        # Check that the images are equal
+        assert not ImageChops.difference(image_sync, image_async).getbbox()
+
+        with pytest.raises(ValueError, match="must be a subpath"):
+            await local_connector.fetch_image_async(
+                f"file://{temp_dir}/../{os.path.basename(image_url)}"
+            )
+        with pytest.raises(RuntimeError, match="Cannot load local files"):
+            await connector.fetch_image_async(
+                f"file://{temp_dir}/../{os.path.basename(image_url)}"
+            )
+
+        with pytest.raises(ValueError, match="must be a subpath"):
+            local_connector.fetch_image(
+                f"file://{temp_dir}/../{os.path.basename(image_url)}"
+            )
+        with pytest.raises(RuntimeError, match="Cannot load local files"):
+            connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True)
+async def test_fetch_image_local_files_with_space_in_name(image_url: str):
+    connector = MediaConnector()
+
+    with TemporaryDirectory() as temp_dir:
+        local_connector = MediaConnector(allowed_local_media_path=temp_dir)
+
+        origin_image = connector.fetch_image(image_url)
+        filename = "file name with space.jpg"
+        origin_image.save(
+            os.path.join(temp_dir, filename),
+            quality=100,
+            icc_profile=origin_image.info.get("icc_profile"),
+        )
+
+        try:
+            image_async = await local_connector.fetch_image_async(
+                f"file://{temp_dir}/{filename}"
+            )
+            image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}")
+        except FileNotFoundError as e:
+            pytest.fail("Failed to fetch image with space in name: {}".format(e))
+        # Check that the images are equal
+        assert not ImageChops.difference(image_sync, image_async).getbbox()
+
+
+@pytest.mark.asyncio
+async def test_fetch_image_error_conversion():
+    connector = MediaConnector()
+    broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK"
+
+    # PIL.UnidentifiedImageError should be converted to ValueError
+    with pytest.raises(ValueError):
+        await connector.fetch_image_async(broken_img)
+
+    with pytest.raises(ValueError):
+        connector.fetch_image(broken_img)
+
+
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
+@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
+async def test_fetch_video_http(video_url: str, num_frames: int):
+    connector = MediaConnector(
+        media_io_kwargs={
+            "video": {
+                "num_frames": num_frames,
+            }
+        }
+    )
+
+    try:
+        video_sync, metadata_sync = connector.fetch_video(video_url)
+        video_async, metadata_async = await connector.fetch_video_async(video_url)
+    except (TimeoutError, asyncio.TimeoutError) as e:
+        pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
+
+    assert np.array_equal(video_sync, video_async)
+    assert metadata_sync == metadata_async
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
+@pytest.mark.parametrize("max_duration", [1, 60, 1800])
+@pytest.mark.parametrize("requested_fps", [2, 24])
+async def test_fetch_video_http_with_dynamic_loader(
+    video_url: str,
+    max_duration: int,
+    requested_fps: int,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic")
+        connector = MediaConnector(
+            media_io_kwargs={
+                "video": {
+                    "max_duration": max_duration,
+                    "requested_fps": requested_fps,
+                }
+            }
+        )
+
+        video_sync, metadata_sync = connector.fetch_video(video_url)
+        video_async, metadata_async = await connector.fetch_video_async(video_url)
+
+        assert np.array_equal(video_sync, video_async)
+        assert metadata_sync == metadata_async
+        assert metadata_sync["video_backend"] == "opencv_dynamic"
+
+
+@pytest.mark.parametrize(
+    "is_embed,start_idx,end_idx,expected",
+    [
+        (None, 2, 4, (2, 4)),
+        (
+            torch.tensor([False, True, False, True, True]),
+            3,
+            5,
+            (1, 3),
+        ),
+        (
+            torch.tensor([False, True, False, True, True]),
+            0,
+            2,
+            (0, 1),
+        ),
+        (
+            torch.tensor([True, False, True, False]),
+            2,
+            2,
+            (1, 1),
+        ),
+    ],
+)
+def test_placeholder_range_get_embeds_indices_in_range(
+    is_embed, start_idx, end_idx, expected
+):
+    length = len(is_embed) if is_embed is not None else 5
+    pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
+    assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected
+
+
+@pytest.mark.parametrize(
+    "offset,is_embed,expected",
+    [
+        (0, None, [(0, 4)]),
+        (
+            2,
+            torch.tensor([False, True, False, True, True]),
+            [(3, 3), (5, 6)],
+        ),
+        (0, torch.tensor([True, True, True, True]), [(0, 3)]),
+        (0, torch.tensor([False, False, False, False]), []),
+    ],
+)
+def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
+    length = len(is_embed) if is_embed is not None else 5
+    pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed)
+    assert pr.extract_embeds_range() == expected
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
+@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
+async def test_allowed_media_domains(video_url: str, num_frames: int):
+    connector = MediaConnector(
+        media_io_kwargs={
+            "video": {
+                "num_frames": num_frames,
+            }
+        },
+        allowed_media_domains=[
+            "www.bogotobogo.com",
+            "github.com",
+        ],
+    )
+
+    video_sync, metadata_sync = connector.fetch_video(video_url)
+    video_async, metadata_async = await connector.fetch_video_async(video_url)
+    assert np.array_equal(video_sync, video_async)
+    assert metadata_sync == metadata_async
+
+    disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png"
+    with pytest.raises(ValueError):
+        _, _ = connector.fetch_video(disallowed_url)
+
+    with pytest.raises(ValueError):
+        _, _ = await connector.fetch_video_async(disallowed_url)

tests/multimodal/test_hasher.py

@@ -16,6 +16,22 @@ ASSETS_DIR = Path(__file__).parent / "assets"
 assert ASSETS_DIR.exists()
 
 
+def test_hash_single_item_different_shape():
+    x1 = torch.zeros(())
+    x2 = torch.zeros((1,))
+
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs(x=x1) != hasher.hash_kwargs(x=x2)
+
+
+def test_hash_key_order_invariant():
+    x = torch.zeros((5, 10))
+    y = torch.ones((5, 10))
+
+    hasher = MultiModalHasher
+    assert hasher.hash_kwargs(x=x, y=y) == hasher.hash_kwargs(y=y, x=x)
+
+
 # NOTE: Images that are the same visually are allowed to have the same hash
 @pytest.mark.parametrize("mode_pair", [("1", "L"), ("RGBA", "CMYK")])
 def test_hash_collision_image_mode(mode_pair):

tests/multimodal/test_inputs.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+
+from vllm.multimodal.inputs import PlaceholderRange
+
+
+@pytest.mark.parametrize(
+    "is_embed,expected",
+    [
+        (None, 5),
+        (torch.tensor([True, True, True, True, True]), 5),
+        (torch.tensor([False, False, False, False, False]), 0),
+        (torch.tensor([True, False, True, False, True]), 3),
+        (torch.tensor([True]), 1),
+    ],
+)
+def test_placeholder_range_get_num_embeds(is_embed, expected):
+    length = len(is_embed) if is_embed is not None else 5
+    pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
+    assert pr.get_num_embeds == expected
+
+
+@pytest.mark.parametrize(
+    "is_embed,expected",
+    [
+        (None, None),
+        (
+            torch.tensor([False, True, False, True, True]),
+            torch.tensor([0, 1, 1, 2, 3]),
+        ),
+        (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
+    ],
+)
+def test_placeholder_range_embeds_cumsum(is_embed, expected):
+    length = len(is_embed) if is_embed is not None else 5
+    pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
+
+    if expected is None:
+        assert pr.embeds_cumsum is None
+        return
+
+    assert torch.equal(pr.embeds_cumsum, expected)
+    # cached_property should return the same object on repeated access
+    assert pr.embeds_cumsum is pr.embeds_cumsum

tests/multimodal/test_utils.py

@@ -1,244 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import asyncio
-import base64
-import mimetypes
-import os
-from tempfile import NamedTemporaryFile, TemporaryDirectory
-
-import numpy as np
 import pytest
 import torch
-from PIL import Image, ImageChops
 
-from vllm.multimodal.image import convert_image_mode
-from vllm.multimodal.inputs import PlaceholderRange
-from vllm.multimodal.media import MediaConnector
-from vllm.multimodal.utils import argsort_mm_positions
-
-# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
-TEST_IMAGE_ASSETS = [
-    "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
-    "Grayscale_8bits_palette_sample_image.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
-    "1280px-Venn_diagram_rgb.svg.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
-    "RGBA_comp.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
-]
-
-TEST_VIDEO_URLS = [
-    "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4",
-    "https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi",
-]
-
-
-@pytest.fixture(scope="module")
-def url_images(local_asset_server) -> dict[str, Image.Image]:
-    return {
-        image_url: local_asset_server.get_image_asset(image_url)
-        for image_url in TEST_IMAGE_ASSETS
-    }
-
-
-def get_supported_suffixes() -> tuple[str, ...]:
-    # We should at least test the file types mentioned in GPT-4 with Vision
-    OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif")
-
-    # Additional file types that are supported by us
-    EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff")
-
-    return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES
-
-
-def _image_equals(a: Image.Image, b: Image.Image) -> bool:
-    return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all()
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
-async def test_fetch_image_http(image_url: str):
-    connector = MediaConnector()
-
-    image_sync = connector.fetch_image(image_url)
-    image_async = await connector.fetch_image_async(image_url)
-    assert _image_equals(image_sync, image_async)
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS)
-@pytest.mark.parametrize("suffix", get_supported_suffixes())
-async def test_fetch_image_base64(
-    url_images: dict[str, Image.Image], raw_image_url: str, suffix: str
-):
-    connector = MediaConnector(
-        # Domain restriction should not apply to data URLs.
-        allowed_media_domains=[
-            "www.bogotobogo.com",
-            "github.com",
-        ]
-    )
-    url_image = url_images[raw_image_url]
-
-    try:
-        mime_type = Image.MIME[Image.registered_extensions()[suffix]]
-    except KeyError:
-        try:
-            mime_type = mimetypes.types_map[suffix]
-        except KeyError:
-            pytest.skip("No MIME type")
-
-    with NamedTemporaryFile(suffix=suffix) as f:
-        try:
-            url_image.save(f.name)
-        except Exception as e:
-            if e.args[0] == "cannot write mode RGBA as JPEG":
-                pytest.skip("Conversion not supported")
-
-            raise
-
-        base64_image = base64.b64encode(f.read()).decode("utf-8")
-        data_url = f"data:{mime_type};base64,{base64_image}"
-
-        data_image_sync = connector.fetch_image(data_url)
-        if _image_equals(url_image, Image.open(f)):
-            assert _image_equals(url_image, data_image_sync)
-        else:
-            pass  # Lossy format; only check that image can be opened
-
-        data_image_async = await connector.fetch_image_async(data_url)
-        assert _image_equals(data_image_sync, data_image_async)
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
-async def test_fetch_image_local_files(image_url: str):
-    connector = MediaConnector()
-
-    with TemporaryDirectory() as temp_dir:
-        local_connector = MediaConnector(allowed_local_media_path=temp_dir)
-
-        origin_image = connector.fetch_image(image_url)
-        origin_image.save(
-            os.path.join(temp_dir, os.path.basename(image_url)),
-            quality=100,
-            icc_profile=origin_image.info.get("icc_profile"),
-        )
-
-        image_async = await local_connector.fetch_image_async(
-            f"file://{temp_dir}/{os.path.basename(image_url)}"
-        )
-        image_sync = local_connector.fetch_image(
-            f"file://{temp_dir}/{os.path.basename(image_url)}"
-        )
-        # Check that the images are equal
-        assert not ImageChops.difference(image_sync, image_async).getbbox()
-
-        with pytest.raises(ValueError, match="must be a subpath"):
-            await local_connector.fetch_image_async(
-                f"file://{temp_dir}/../{os.path.basename(image_url)}"
-            )
-        with pytest.raises(RuntimeError, match="Cannot load local files"):
-            await connector.fetch_image_async(
-                f"file://{temp_dir}/../{os.path.basename(image_url)}"
-            )
-
-        with pytest.raises(ValueError, match="must be a subpath"):
-            local_connector.fetch_image(
-                f"file://{temp_dir}/../{os.path.basename(image_url)}"
-            )
-        with pytest.raises(RuntimeError, match="Cannot load local files"):
-            connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}")
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True)
-async def test_fetch_image_local_files_with_space_in_name(image_url: str):
-    connector = MediaConnector()
-
-    with TemporaryDirectory() as temp_dir:
-        local_connector = MediaConnector(allowed_local_media_path=temp_dir)
-
-        origin_image = connector.fetch_image(image_url)
-        filename = "file name with space.jpg"
-        origin_image.save(
-            os.path.join(temp_dir, filename),
-            quality=100,
-            icc_profile=origin_image.info.get("icc_profile"),
-        )
-
-        try:
-            image_async = await local_connector.fetch_image_async(
-                f"file://{temp_dir}/{filename}"
-            )
-            image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}")
-        except FileNotFoundError as e:
-            pytest.fail("Failed to fetch image with space in name: {}".format(e))
-        # Check that the images are equal
-        assert not ImageChops.difference(image_sync, image_async).getbbox()
-
-
-@pytest.mark.asyncio
-async def test_fetch_image_error_conversion():
-    connector = MediaConnector()
-    broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK"
-
-    # PIL.UnidentifiedImageError should be converted to ValueError
-    with pytest.raises(ValueError):
-        await connector.fetch_image_async(broken_img)
-
-    with pytest.raises(ValueError):
-        connector.fetch_image(broken_img)
-
-
-@pytest.mark.flaky(reruns=3, reruns_delay=5)
-@pytest.mark.asyncio
-@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
-@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
-async def test_fetch_video_http(video_url: str, num_frames: int):
-    connector = MediaConnector(
-        media_io_kwargs={
-            "video": {
-                "num_frames": num_frames,
-            }
-        }
-    )
-
-    try:
-        video_sync, metadata_sync = connector.fetch_video(video_url)
-        video_async, metadata_async = await connector.fetch_video_async(video_url)
-    except (TimeoutError, asyncio.TimeoutError) as e:
-        pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
-
-    assert np.array_equal(video_sync, video_async)
-    assert metadata_sync == metadata_async
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
-@pytest.mark.parametrize("max_duration", [1, 60, 1800])
-@pytest.mark.parametrize("requested_fps", [2, 24])
-async def test_fetch_video_http_with_dynamic_loader(
-    video_url: str,
-    max_duration: int,
-    requested_fps: int,
-    monkeypatch: pytest.MonkeyPatch,
-):
-    with monkeypatch.context() as m:
-        m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic")
-        connector = MediaConnector(
-            media_io_kwargs={
-                "video": {
-                    "max_duration": max_duration,
-                    "requested_fps": requested_fps,
-                }
-            }
-        )
-
-        video_sync, metadata_sync = connector.fetch_video(video_url)
-        video_async, metadata_async = await connector.fetch_video_async(video_url)
-
-        assert np.array_equal(video_sync, video_async)
-        assert metadata_sync == metadata_async
-        assert metadata_sync["video_backend"] == "opencv_dynamic"
+from vllm.multimodal.inputs import (
+    MultiModalBatchedField,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    MultiModalSharedField,
+    PlaceholderRange,
+)
+from vllm.multimodal.utils import argsort_mm_positions, group_and_batch_mm_items
 
 
 @pytest.mark.parametrize(
@@ -412,121 +184,35 @@ def test_argsort_mm_positions(case):
     assert modality_idxs == expected_modality_idxs
 
 
-@pytest.mark.parametrize(
-    "is_embed,expected",
-    [
-        (None, 5),
-        (torch.tensor([True, True, True, True, True]), 5),
-        (torch.tensor([False, False, False, False, False]), 0),
-        (torch.tensor([True, False, True, False, True]), 3),
-        (torch.tensor([True]), 1),
-    ],
-)
-def test_placeholder_range_get_num_embeds(is_embed, expected):
-    length = len(is_embed) if is_embed is not None else 5
-    pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
-    assert pr.get_num_embeds == expected
-
-
-@pytest.mark.parametrize(
-    "is_embed,expected",
-    [
-        (None, None),
-        (
-            torch.tensor([False, True, False, True, True]),
-            torch.tensor([0, 1, 1, 2, 3]),
-        ),
-        (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
-    ],
-)
-def test_placeholder_range_embeds_cumsum(is_embed, expected):
-    length = len(is_embed) if is_embed is not None else 5
-    pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
-
-    if expected is None:
-        assert pr.embeds_cumsum is None
-        return
-
-    assert torch.equal(pr.embeds_cumsum, expected)
-    # cached_property should return the same object on repeated access
-    assert pr.embeds_cumsum is pr.embeds_cumsum
-
-
-@pytest.mark.parametrize(
-    "is_embed,start_idx,end_idx,expected",
-    [
-        (None, 2, 4, (2, 4)),
-        (
-            torch.tensor([False, True, False, True, True]),
-            3,
-            5,
-            (1, 3),
-        ),
-        (
-            torch.tensor([False, True, False, True, True]),
-            0,
-            2,
-            (0, 1),
-        ),
-        (
-            torch.tensor([True, False, True, False]),
-            2,
-            2,
-            (1, 1),
-        ),
-    ],
-)
-def test_placeholder_range_get_embeds_indices_in_range(
-    is_embed, start_idx, end_idx, expected
-):
-    length = len(is_embed) if is_embed is not None else 5
-    pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
-    assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected
-
-
-@pytest.mark.parametrize(
-    "offset,is_embed,expected",
-    [
-        (0, None, [(0, 4)]),
-        (
-            2,
-            torch.tensor([False, True, False, True, True]),
-            [(3, 3), (5, 6)],
-        ),
-        (0, torch.tensor([True, True, True, True]), [(0, 3)]),
-        (0, torch.tensor([False, False, False, False]), []),
-    ],
-)
-def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
-    length = len(is_embed) if is_embed is not None else 5
-    pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed)
-    assert pr.extract_embeds_range() == expected
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
-@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
-async def test_allowed_media_domains(video_url: str, num_frames: int):
-    connector = MediaConnector(
-        media_io_kwargs={
-            "video": {
-                "num_frames": num_frames,
-            }
-        },
-        allowed_media_domains=[
-            "www.bogotobogo.com",
-            "github.com",
-        ],
+def test_group_and_batch_mm_items_split_by_fieldset():
+    elem = MultiModalFieldElem(
+        data=torch.empty(1, dtype=torch.uint8),
+        field=MultiModalBatchedField(),
     )
+    item1 = MultiModalKwargsItem({"x": elem, "y": elem})
+    item2 = MultiModalKwargsItem({"y": elem, "x": elem})
+    item3 = MultiModalKwargsItem({"x": elem, "y": elem, "z": elem})
+    item4 = MultiModalKwargsItem({"x": elem})
+    item5 = MultiModalKwargsItem({"x": elem, "y": elem})
 
-    video_sync, metadata_sync = connector.fetch_video(video_url)
-    video_async, metadata_async = await connector.fetch_video_async(video_url)
-    assert np.array_equal(video_sync, video_async)
-    assert metadata_sync == metadata_async
+    res = group_and_batch_mm_items([item1, item2, item3, item4, item5])
+    assert [num_items for num_items, _ in res] == [2, 1, 1, 1]
 
-    disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png"
-    with pytest.raises(ValueError):
-        _, _ = connector.fetch_video(disallowed_url)
 
-    with pytest.raises(ValueError):
-        _, _ = await connector.fetch_video_async(disallowed_url)
+def test_group_and_batch_mm_items_split_by_shared_data():
+    elem1 = MultiModalFieldElem(
+        data=torch.zeros(1, dtype=torch.uint8),
+        field=MultiModalSharedField(batch_size=1),
+    )
+    elem2 = MultiModalFieldElem(
+        data=torch.zeros(2, dtype=torch.uint8),
+        field=MultiModalSharedField(batch_size=1),
+    )
+    item1 = MultiModalKwargsItem({"x": elem1})
+    item2 = MultiModalKwargsItem({"x": elem1})
+    item3 = MultiModalKwargsItem({"x": elem2})
+    item4 = MultiModalKwargsItem({"x": elem1})
+    item5 = MultiModalKwargsItem({"x": elem2})
+
+    res = group_and_batch_mm_items([item1, item2, item3, item4, item5])
+    assert [num_items for num_items, _ in res] == [2, 1, 1, 1]

tests/test_pooling_params.py

@@ -27,35 +27,24 @@ class MockModelConfig:
     pooler_config: PoolerConfig
 
 
-def test_task():
-    pooling_params = PoolingParams()
-    pooling_params.verify(task="score")
-
-    pooling_params = PoolingParams(task="score")
-    pooling_params.verify(task="score")
-
-    with pytest.raises(ValueError):
-        pooling_params.verify(task="classify")
-
-
 def test_embed():
     task = "embed"
     model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
 
-    pooling_params = PoolingParams(use_activation=None)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=None)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=True)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=True)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=False)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=False)
+    pooling_params.verify(model_config)
 
     invalid_parameters = classify_parameters + step_pooling_parameters
     for p in set(invalid_parameters) - set(embed_parameters):
         with pytest.raises(ValueError):
-            pooling_params = PoolingParams(**{p: True})
-            pooling_params.verify(task=task, model_config=model_config)
+            pooling_params = PoolingParams(task=task, **{p: True})
+            pooling_params.verify(model_config)
 
 
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
@@ -63,7 +52,6 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
     task = "embed"
     model_config = ModelConfig(
         model_info.name,
-        task="auto",
         tokenizer=model_info.name,
         tokenizer_mode="auto",
         trust_remote_code=False,
@@ -71,37 +59,39 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
         dtype="float16",
     )
 
-    pooling_params = PoolingParams(dimensions=None)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, dimensions=None)
+    pooling_params.verify(model_config)
 
     with pytest.raises(ValueError):
-        pooling_params = PoolingParams(dimensions=1)
-        pooling_params.verify(task=task, model_config=model_config)
+        pooling_params = PoolingParams(task=task, dimensions=1)
+        pooling_params.verify(model_config)
 
     if model_info.is_matryoshka:
         assert model_info.matryoshka_dimensions is not None
-        pooling_params = PoolingParams(dimensions=model_info.matryoshka_dimensions[0])
-        pooling_params.verify(task=task, model_config=model_config)
+        pooling_params = PoolingParams(
+            task=task, dimensions=model_info.matryoshka_dimensions[0]
+        )
+        pooling_params.verify(model_config)
 
 
 @pytest.mark.parametrize("task", ["score", "classify"])
 def test_classify(task):
     model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
 
-    pooling_params = PoolingParams(use_activation=None)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=None)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=True)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=True)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=False)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=False)
+    pooling_params.verify(model_config)
 
     invalid_parameters = embed_parameters + step_pooling_parameters
     for p in set(invalid_parameters) - set(classify_parameters):
         with pytest.raises(ValueError):
-            pooling_params = PoolingParams(**{p: True})
-            pooling_params.verify(task=task, model_config=model_config)
+            pooling_params = PoolingParams(task=task, **{p: True})
+            pooling_params.verify(model_config)
 
 
 @pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
@@ -111,14 +101,14 @@ def test_token_embed(pooling_type: str):
         pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
     )
 
-    pooling_params = PoolingParams(use_activation=None)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=None)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=True)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=True)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=False)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=False)
+    pooling_params.verify(model_config)
 
     invalid_parameters = classify_parameters
     if pooling_type != "STEP":
@@ -126,8 +116,8 @@ def test_token_embed(pooling_type: str):
 
     for p in set(invalid_parameters) - set(embed_parameters):
         with pytest.raises(ValueError):
-            pooling_params = PoolingParams(**{p: True})
-            pooling_params.verify(task=task, model_config=model_config)
+            pooling_params = PoolingParams(task=task, **{p: True})
+            pooling_params.verify(model_config)
 
 
 @pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
@@ -137,14 +127,14 @@ def test_token_classify(pooling_type: str):
         pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
     )
 
-    pooling_params = PoolingParams(use_activation=None)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=None)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=True)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=True)
+    pooling_params.verify(model_config)
 
-    pooling_params = PoolingParams(use_activation=False)
-    pooling_params.verify(task=task, model_config=model_config)
+    pooling_params = PoolingParams(task=task, use_activation=False)
+    pooling_params.verify(model_config)
 
     invalid_parameters = embed_parameters
     if pooling_type != "STEP":
@@ -152,5 +142,5 @@ def test_token_classify(pooling_type: str):
 
     for p in set(invalid_parameters) - set(classify_parameters):
         with pytest.raises(ValueError):
-            pooling_params = PoolingParams(**{p: True})
-            pooling_params.verify(task=task, model_config=model_config)
+            pooling_params = PoolingParams(task=task, **{p: True})
+            pooling_params.verify(model_config)

tests/utils.py

@@ -197,12 +197,86 @@ class RemoteOpenAIServer:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        pid = self.proc.pid
+        # Graceful shutdown
         self.proc.terminate()
         try:
-            self.proc.wait(8)
+            self.proc.wait(timeout=15)
+            print(f"[RemoteOpenAIServer] Server {pid} terminated gracefully")
         except subprocess.TimeoutExpired:
-            # force kill if needed
+            print(
+                f"[RemoteOpenAIServer] Server {pid} did not respond "
+                "to SIGTERM, sending SIGKILL"
+            )
             self.proc.kill()
+            try:
+                self.proc.wait(timeout=5)
+                print(f"[RemoteOpenAIServer] Server {pid} killed")
+            except subprocess.TimeoutExpired as err:
+                raise RuntimeError(
+                    f"[RemoteOpenAIServer] Failed to kill server process {pid}"
+                ) from err
+        # Wait for GPU memory to be released
+        self._wait_for_gpu_memory_release()
+
+    def _get_gpu_memory_used(self) -> float | None:
+        """Get total GPU memory used across all visible devices in bytes."""
+        try:
+            if current_platform.is_rocm():
+                with _nvml():
+                    handles = amdsmi_get_processor_handles()
+                    total_used = 0
+                    for handle in handles:
+                        vram_info = amdsmi_get_gpu_vram_usage(handle)
+                        total_used += vram_info["vram_used"]
+                    return total_used
+            elif current_platform.is_cuda():
+                with _nvml():
+                    total_used = 0
+                    device_count = cuda_device_count_stateless()
+                    for i in range(device_count):
+                        handle = nvmlDeviceGetHandleByIndex(i)
+                        mem_info = nvmlDeviceGetMemoryInfo(handle)
+                        total_used += mem_info.used
+                    return total_used
+        except Exception as e:
+            print(f"[RemoteOpenAIServer] Could not query GPU memory: {e}")
+            return None
+        return None
+
+    def _wait_for_gpu_memory_release(self, timeout: float = 30.0):
+        """Poll GPU memory until it stabilizes, indicating cleanup is complete."""
+        start = time.time()
+        prev_used: float | None = None
+        stable_count = 0
+
+        while time.time() - start < timeout:
+            used = self._get_gpu_memory_used()
+
+            if used is None:
+                return  # Can't query, assume ok
+
+            if prev_used is not None and abs(used - prev_used) < 100 * 1024 * 1024:
+                stable_count += 1
+                if stable_count >= 3:
+                    used_gb = used / 1e9
+                    print(
+                        f"[RemoteOpenAIServer] GPU memory stabilized "
+                        f"at {used_gb:.2f} GB"
+                    )
+                    return
+            else:
+                stable_count = 0
+
+            prev_used = used
+            time.sleep(0.1)
+
+        last_reading = prev_used / 1e9 if prev_used is not None else 0.0
+        raise RuntimeError(
+            f"[RemoteOpenAIServer] GPU memory did not stabilize within {timeout}s. "
+            f"Last reading: {last_reading:.2f} GB. "
+            "Child processes may still be holding GPU memory."
+        )
 
     def _poll(self) -> int | None:
         """Subclasses override this method to customize process polling"""

tests/v1/core/test_scheduler.py

@@ -1136,7 +1136,7 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
         prompt_logprobs_dict={},
         pooler_output=[],
     )
-    scheduler.update_from_output(output, EMPTY_OUTPUT)
+    initial_ecos = scheduler.update_from_output(output, EMPTY_OUTPUT)
 
     # Simulate KV transfer completion using KVConnectorOutput.finished_recving
     output = scheduler.schedule()
@@ -1156,6 +1156,8 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]):
     for req_id in req_ids:
         assert req_id in scheduler.finished_recving_kv_req_ids
 
+    return initial_ecos
+
 
 @pytest.mark.parametrize("is_async", [False, True])
 def test_kv_connector_basic(is_async: bool):
@@ -1286,29 +1288,72 @@ def test_kv_connector_basic(is_async: bool):
 
 
 @pytest.mark.parametrize("is_async", [False, True])
-def test_external_prefix_cache_metrics(is_async: bool):
+@pytest.mark.parametrize("local_cache_hits", [False, True])
+def test_external_prefix_cache_metrics(is_async: bool, local_cache_hits: bool):
     """
     Verify connector prefix cache metrics are updated
     correctly when the scheduler processes requests with KV connector hits.
     """
 
+    BLOCK_SIZE = 16
+    if local_cache_hits:
+        NUM_MATCHED_NEW_TOKENS = BLOCK_SIZE * 2  # 32 tokens
+        NUM_LOCAL_HITS = NUM_MATCHED_NEW_TOKENS * 2  # 64 tokens
+        NUM_REQUESTS = 1
+        NUM_TOKENS = NUM_LOCAL_HITS * 2  # 128 tokens
+    else:
+        NUM_MATCHED_NEW_TOKENS = 4
+        NUM_LOCAL_HITS = 0
+        NUM_REQUESTS = 2
+        NUM_TOKENS = 8  # 8 tokens
+
     # Setup Scheduler.
-    NUM_MATCHED_NEW_TOKENS = 4
     scheduler = create_scheduler(
-        enable_prefix_caching=False,
+        enable_prefix_caching=local_cache_hits,
         use_kv_connector=mock_kv(
             matched_tokens=NUM_MATCHED_NEW_TOKENS, is_async=is_async
         ),
+        block_size=BLOCK_SIZE,
     )
 
-    # --- Prepare simple requests ---
-    NUM_REQUESTS = 2
-    NUM_TOKENS = 8
+    if local_cache_hits:
+        # First, establish local cache by running a request to completion
+        requests = create_requests(
+            num_requests=1,
+            num_tokens=NUM_LOCAL_HITS,
+            max_tokens=2,
+            block_size=BLOCK_SIZE,
+        )
+        req_ids = []
+        req_to_index = {}
+        for i, request in enumerate(requests):
+            scheduler.add_request(request)
+            req_ids.append(request.request_id)
+            req_to_index[request.request_id] = i
+
+        if is_async:
+            _step_until_kv_transfer_finished(scheduler, req_ids)
+
+        # Run first request to completion to establish local cache
+        output = scheduler.schedule()
+        MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
+            req_ids=req_ids,
+            req_id_to_index=req_to_index,
+            sampled_token_ids=[[1000]] * len(req_ids),
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        )
+        _step_until_done(scheduler, output, MODEL_RUNNER_OUTPUT)
+        _ = scheduler.schedule()
+
+    # --- Prepare test requests ---
     MAX_TOKENS = 2
     requests = create_requests(
         num_requests=NUM_REQUESTS,
         num_tokens=NUM_TOKENS,
         max_tokens=MAX_TOKENS,
+        block_size=BLOCK_SIZE,
     )
     req_ids = []
     req_to_index = {}
@@ -1317,8 +1362,9 @@ def test_external_prefix_cache_metrics(is_async: bool):
         req_ids.append(request.request_id)
         req_to_index[request.request_id] = i
 
+    initial_ecos = None
     if is_async:
-        _step_until_kv_transfer_finished(scheduler, req_ids)
+        initial_ecos = _step_until_kv_transfer_finished(scheduler, req_ids)
 
     # --- Trigger scheduling and simulate model output ---
     output = scheduler.schedule()
@@ -1338,10 +1384,23 @@ def test_external_prefix_cache_metrics(is_async: bool):
     assert ecos is not None and len(ecos) > 0
     assert ecos[0].scheduler_stats is not None
 
-    external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats
+    if local_cache_hits:
+        # For async, local cache stats come from the first step
+        if initial_ecos:
+            local_stats = initial_ecos[0].scheduler_stats.prefix_cache_stats
+        else:
+            local_stats = ecos[0].scheduler_stats.prefix_cache_stats
+        assert local_stats is not None
+        assert local_stats.queries == NUM_TOKENS * NUM_REQUESTS
+        assert local_stats.hits == NUM_LOCAL_HITS * NUM_REQUESTS
+
+    if initial_ecos:
+        external_stats = initial_ecos[0].scheduler_stats.connector_prefix_cache_stats
+    else:
+        external_stats = ecos[0].scheduler_stats.connector_prefix_cache_stats
     assert external_stats is not None
 
-    assert external_stats.queries == NUM_TOKENS * NUM_REQUESTS
+    assert external_stats.queries == (NUM_TOKENS - NUM_LOCAL_HITS) * NUM_REQUESTS
     assert external_stats.hits == NUM_MATCHED_NEW_TOKENS * NUM_REQUESTS
     assert external_stats.requests == NUM_REQUESTS
     assert external_stats.preempted_requests == 0

tests/v1/e2e/test_spec_decode.py

@@ -13,15 +13,12 @@ from vllm import LLM, SamplingParams
 from vllm.assets.base import VLLM_S3_BUCKET_URL
 from vllm.assets.image import VLM_IMAGES_DIR
 from vllm.benchmarks.datasets import InstructCoderDataset
-from vllm.config.vllm import VllmConfig
+from vllm.config import VllmConfig
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.v1.metrics.reader import Metric
-from vllm.v1.spec_decode.draft_model import (
-    create_vllm_config_for_draft_model,
-    merge_toks_kernel,
-)
+from vllm.v1.spec_decode.utils import create_vllm_config_for_draft_model
 
 MTP_SIMILARITY_RATE = 0.8
 
@@ -625,6 +622,8 @@ class ArgsTest:
     expected_acceptance_rate: float
     expected_acceptance_len: float
     # Defaults
+    enforce_eager: bool = True
+    parallel_drafting: bool = False
     target_tensor_parallel_size: int = 1
     draft_tensor_parallel_size: int = 1
     max_model_len: int = 1024
@@ -658,7 +657,8 @@ cases = [
 @pytest.mark.parametrize("args", cases)
 @pytest.mark.parametrize("enforce_eager", [True, False])
 def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
-    assert_draft_model_correctness(args, enforce_eager)
+    args.enforce_eager = enforce_eager
+    assert_draft_model_correctness(args)
 
 
 def test_draft_model_realistic_example():
@@ -668,11 +668,28 @@ def test_draft_model_realistic_example():
         dataset="likaixin/InstructCoder",
         num_speculative_tokens=3,
         sampling_config=greedy_sampling(),
+        enforce_eager=False,
         # values below are not derived, but just prevent a regression
         expected_acceptance_len=2.8,
         expected_acceptance_rate=0.55,
     )
-    assert_draft_model_correctness(args, enforce_eager=False)
+    assert_draft_model_correctness(args)
+
+
+def test_draft_model_parallel_drafting():
+    args = ArgsTest(
+        target_model="Qwen/Qwen3-1.7B",
+        draft_model="amd/PARD-Qwen3-0.6B",
+        dataset="likaixin/InstructCoder",
+        num_speculative_tokens=3,
+        sampling_config=greedy_sampling(),
+        parallel_drafting=True,
+        enforce_eager=False,
+        # values below are collected from a stable run, with ~5% tolerance
+        expected_acceptance_len=2.375,
+        expected_acceptance_rate=0.45,
+    )
+    assert_draft_model_correctness(args)
 
 
 @pytest.mark.parametrize(
@@ -691,8 +708,9 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
         target_model=tgt_model,
         draft_model=draft_model,
         **some_high_acceptance_metrics(),
+        enforce_eager=enforce_eager,
     )
-    assert_draft_model_correctness(sd_case, enforce_eager)
+    assert_draft_model_correctness(sd_case)
 
 
 def test_draft_model_tensor_parallelism():
@@ -704,8 +722,9 @@ def test_draft_model_tensor_parallelism():
         draft_model="Qwen/Qwen3-0.6B",
         draft_tensor_parallel_size=2,
         **some_high_acceptance_metrics(),
+        enforce_eager=False,
     )
-    assert_draft_model_correctness(sd_case, enforce_eager=False)
+    assert_draft_model_correctness(sd_case)
 
 
 def test_draft_model_engine_args_tensor_parallelism():
@@ -750,7 +769,7 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname():
         engine_args.create_engine_config()
 
 
-def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
+def assert_draft_model_correctness(args: ArgsTest):
     """Compare the outputs using and not using speculative decoding.
     In the greedy decoding case, the outputs must match EXACTLY."""
     test_prompts: list[Messages] = get_messages(
@@ -764,14 +783,15 @@ def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
             "method": "draft_model",
             "num_speculative_tokens": args.num_speculative_tokens,
             "max_model_len": args.max_model_len,
-            "enforce_eager": enforce_eager,
+            "enforce_eager": args.enforce_eager,
             "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
+            "parallel_drafting": args.parallel_drafting,
         },
         max_num_seqs=100,  # limit cudagraph capture runtime
         max_model_len=args.max_model_len,
         gpu_memory_utilization=args.gpu_memory_utilization,
         tensor_parallel_size=args.target_tensor_parallel_size,
-        enforce_eager=enforce_eager,
+        enforce_eager=args.enforce_eager,
         disable_log_stats=False,  # enables get_metrics()
     )
     # we don't check the outputs, only check the metrics
@@ -813,57 +833,6 @@ def some_high_acceptance_metrics() -> dict:
     }
 
 
-def test_merge_toks_kernel():
-    device = "cuda"
-    merged_len = 5 + 2  # len(target_toks) = 5, batch_size = 2
-    merged = torch.full((merged_len,), -100, device=device)  # -100 is arbitrary
-    is_rejected_tok = torch.full((merged_len,), True, device=device)
-    grid = (2,)
-    merge_toks_kernel[grid](
-        target_toks_ptr=torch.tensor([0, 1, 2, 0, 1], device=device),
-        next_toks_ptr=torch.tensor([3, 2], device=device),
-        query_start_locs_ptr=torch.tensor([0, 3], device=device),
-        query_end_locs_ptr=torch.tensor([2, 4], device=device),
-        out_ptr_merged_toks=merged,
-        out_ptr_is_rejected_tok=is_rejected_tok,
-        target_toks_size=5,
-        rejected_tok_fill=-1,
-    )
-    expected_merged = torch.tensor([0, 1, 2, 3, 0, 1, 2], device=device)
-    assert torch.allclose(merged, expected_merged)
-
-    expected_rejected_toks = torch.tensor([False] * merged_len, device=device)
-    assert torch.allclose(is_rejected_tok, expected_rejected_toks)
-
-
-def test_merge_toks_kernel_with_rejected_tokens():
-    device = "cuda"
-    merged_size = 9 + 2  # len(target_toks) = 9, batch_size = 2
-    merged = torch.full((merged_size,), -100, device=device)
-    is_rejected_tok = torch.full((merged_size,), True, device=device)
-    grid = (2,)
-    merge_toks_kernel[grid](
-        #                                       rejected tokens
-        #                                       ↓   ↓   ↓         ↓
-        target_toks_ptr=torch.tensor([0, 1, 2, 13, 14, 15, 0, 1, 22], device=device),
-        next_toks_ptr=torch.tensor([3, 2], device=device),
-        query_start_locs_ptr=torch.tensor([0, 6], device=device),
-        query_end_locs_ptr=torch.tensor([2, 7], device=device),
-        out_ptr_merged_toks=merged,
-        out_ptr_is_rejected_tok=is_rejected_tok,
-        target_toks_size=9,
-        rejected_tok_fill=-1,
-    )
-    expected_merged = torch.tensor([0, 1, 2, 3, -1, -1, -1, 0, 1, 2, -1], device=device)
-    assert torch.allclose(merged, expected_merged)
-
-    expected_rejected_toks = torch.tensor(
-        [False, False, False, False, True, True, True, False, False, False, True],
-        device=device,
-    )
-    assert torch.allclose(is_rejected_tok, expected_rejected_toks)
-
-
 def compute_acceptance_rate(metrics: list[Metric]) -> float:
     name2metric = {metric.name: metric for metric in metrics}
     n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value  # type: ignore

tests/v1/e2e/test_streaming_input.py

@@ -19,7 +19,7 @@ import pytest
 import pytest_asyncio
 
 from vllm import SamplingParams
-from vllm.inputs.data import StreamingInput
+from vllm.inputs import StreamingInput
 from vllm.outputs import RequestOutput
 from vllm.platforms import current_platform
 from vllm.sampling_params import RequestOutputKind

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -14,8 +14,8 @@ tp_configs=(
   "GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"
 )
 dp_ep_configs=(
-"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1) 
-"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1) 
+"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
+"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
 )
 
 # Select config array based on DP_EP env var
@@ -57,3 +57,9 @@ if [[ -n "${FLASHINFER:-}" ]]; then
 else
   echo "FLASHINFER not set, skipping FLASHINFER runs."
 fi
+
+# Check if cross-layers is enabled (non-empty)
+if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then
+  echo "CROSS_LAYERS_BLOCKS is set, rerunning with --enable-cross-layers"
+  run_tests "default backend" "--enable-cross-layers"
+fi

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -4,6 +4,7 @@ set -xe
 # Parse command line arguments
 KV_BUFFER_DEVICE="cuda"  # Default to cuda
 ATTENTION_BACKEND=""  # Default to empty (use vllm default)
+CROSS_LAYERS_BLOCKS="False"
 while [[ $# -gt 0 ]]; do
   case $1 in
     --kv_buffer_device)
@@ -14,6 +15,10 @@ while [[ $# -gt 0 ]]; do
       ATTENTION_BACKEND="$2"
       shift 2
       ;;
+    --enable-cross-layers)
+      CROSS_LAYERS_BLOCKS="True"
+      shift 1
+      ;;
     *)
       echo "Unknown option $1"
       echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]"
@@ -34,11 +39,17 @@ else
   KV_CONFIG_HETERO_LAYOUT=''
 fi
 
+if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
+  KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
+else
+  KV_EXTRA_CONFIG=''
+fi
+
 # Build the kv-transfer-config once
 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
-  KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
+  KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
 else
-  KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
+  KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
 fi
 
 # Models to run

tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py

@@ -281,6 +281,17 @@ def test_sync_fail_invalid_blocks_evicted(fail_scheduler: Scheduler):
         f"(hash should be None), but hash is still {block.block_hash}"
     )
 
+    # Verify connector prefix cache stats:
+    # - queries = num_prompt_tokens (total tokens not in local cache)
+    # - hits = num_external_computed_tokens (tokens loaded externally)
+    assert engine_outputs.scheduler_stats is not None
+    stats = engine_outputs.scheduler_stats
+    assert stats.connector_prefix_cache_stats is not None
+    conn_stats = stats.connector_prefix_cache_stats
+    assert conn_stats.requests == 1
+    assert conn_stats.queries == num_prompt_tokens
+    assert conn_stats.hits == num_external_computed_tokens
+
 
 def test_async_recompute_blocks_not_cached_when_invalid(
     recompute_scheduler: Scheduler,
@@ -364,7 +375,9 @@ def test_async_recompute_blocks_not_cached_when_invalid(
     with patch.object(
         recompute_scheduler.kv_cache_manager, "evict_blocks", evict_blocks_spy
     ):
-        recompute_scheduler.update_from_output(scheduler_output, model_runner_output)
+        outputs = recompute_scheduler.update_from_output(
+            scheduler_output, model_runner_output
+        )
 
     # verify evict_blocks was NOT called (async blocks excluded from eviction)
     assert len(evict_blocks_calls) == 0, (
@@ -386,6 +399,19 @@ def test_async_recompute_blocks_not_cached_when_invalid(
         f"Block {invalid_block_id} hash should be None but is {block.block_hash}"
     )
 
+    # Verify connector prefix cache stats:
+    # - queries = num_prompt_tokens (total tokens not in local cache)
+    # - hits = num_external_computed_tokens (tokens loaded externally)
+    assert len(outputs) == 1
+    engine_outputs = next(iter(outputs.values()))
+    assert engine_outputs.scheduler_stats is not None
+    stats = engine_outputs.scheduler_stats
+    assert stats.connector_prefix_cache_stats is not None
+    conn_stats = stats.connector_prefix_cache_stats
+    assert conn_stats.requests == 1
+    assert conn_stats.queries == num_prompt_tokens
+    assert conn_stats.hits == num_external_computed_tokens
+
     # now simulate async transfer completing
     model_runner_output_2 = create_model_runner_output(
         reqs=[],

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -18,8 +18,12 @@ import ray
 import torch
 
 from vllm import LLM
-from vllm.config import KVTransferConfig
-from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
+from vllm.config import KVTransferConfig, set_current_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.utils import (
+    KVOutputAggregator,
+    TpKVTopology,
+    get_current_attn_backend,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
@@ -46,10 +50,14 @@ from vllm.platforms import current_platform
 from vllm.platforms.interface import Platform
 from vllm.sampling_params import SamplingParams
 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.output_processor import OutputProcessor
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import RequestStatus
+from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
+from vllm.v1.worker.utils import AttentionGroup
 
 from .utils import create_request, create_scheduler, create_vllm_config
 
@@ -366,6 +374,7 @@ def test_kv_transfer_handshake(dist_init):
 
         # Decode connector will be able to create handshake with the prefill connector.
         decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
+        decode_connector.register_kv_caches(kv_caches)
 
         # Here we are testing the retrieval of NIXLAgentMetadata.
         # Knowing the implementation detail, we override the add_remote_agent
@@ -402,6 +411,23 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self.kv_cache_layout = kv_cache_layout
         # Mock register_kv_caches attribute needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
+        test_shape = self.attn_backend.get_kv_cache_shape(
+            num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
+        )
+        self.kv_topo = TpKVTopology(
+            tp_rank=self.tp_rank,
+            engine_id=self.engine_id,
+            remote_tp_size=self._tp_size,  # shared state
+            remote_block_size=self._block_size,  # shared state
+            is_mla=self.use_mla,
+            total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
+            attn_backend=self.attn_backend,
+            tensor_shape=test_shape,
+        )
+
+        self.compat_hash = compute_nixl_compatibility_hash(
+            self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
+        )
 
     def _nixl_handshake(
         self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
@@ -1352,6 +1378,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
     llm.llm_engine.engine_core.shutdown()
 
 
+@pytest.mark.parametrize("enable_cross_layers", ["False", "True"])
 @pytest.mark.parametrize(
     "attn_backend",
     [
@@ -1372,7 +1399,9 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
         "TRITON_ATTN",
     ],
 )
-def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
+def test_register_kv_caches(
+    default_vllm_config, dist_init, attn_backend, enable_cross_layers
+):
     """
     Test that register_kv_caches() properly calls nixl_wrapper methods with
     correct data.
@@ -1386,6 +1415,12 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
 
     vllm_config = create_vllm_config(attention_backend=attn_backend)
 
+    # Enable cross layers blocks
+    vllm_config.kv_transfer_config.kv_connector_extra_config[
+        "enable_cross_layers_blocks"
+    ] = enable_cross_layers
+    set_kv_cache_layout("HND")
+
     # Import the appropriate backend based on the parameter
     if attn_backend == "FLASH_ATTN":
         from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
@@ -1400,44 +1435,6 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
 
         backend_cls = TritonAttentionBackend
 
-    # Create test kv cache tensors using proper backend shape
-    kv_cache_shape = backend_cls.get_kv_cache_shape(
-        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
-    )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-    kv_caches = {
-        "layer0": shared_tensor,
-        "layer1": unique_tensor,
-        "layer2": shared_tensor,
-    }
-
-    # Store tensor info for validation
-
-    test_shape = backend_cls.get_kv_cache_shape(
-        num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
-    )
-    is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
-
-    if is_blocks_first:
-        expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
-        expected_base_addrs = [
-            shared_tensor.data_ptr(),
-            unique_tensor.data_ptr(),
-        ]
-        expected_num_entries = 2
-    else:
-        expected_tensor_size = (
-            shared_tensor[0].element_size() * shared_tensor[0].numel()
-        )
-        expected_base_addrs = [
-            shared_tensor[0].data_ptr(),
-            shared_tensor[1].data_ptr(),
-            unique_tensor[0].data_ptr(),
-            unique_tensor[1].data_ptr(),
-        ]
-        expected_num_entries = 4
-
     nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
     with (
         patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
@@ -1466,6 +1463,111 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
         # Reassure the shutdown() check that the thread is terminated
         mock_thread.return_value.is_alive.return_value = False
 
+        expected_tensor_size: int
+        expected_base_addrs: list[int]
+        expected_num_entries: int
+        kv_caches: dict[str, torch.Tensor]
+        assert str(enable_cross_layers).lower() != "true" or (
+            (attn_backend not in ("FLASH_ATTN", "FLASHINFER"))
+            or connector.prefer_cross_layer_blocks
+        )
+        if connector.prefer_cross_layer_blocks:
+            num_layers = 32
+            block_size = 16
+            num_blocks = 8
+            kv_cache_spec = AttentionSpec(
+                block_size=block_size,
+                num_kv_heads=4,
+                head_size=64,
+                dtype=torch.bfloat16,
+            )
+            kv_cache_config = KVCacheConfig(
+                num_blocks=num_blocks,
+                kv_cache_tensors=[
+                    KVCacheTensor(
+                        size=kv_cache_spec.page_size_bytes * num_blocks,
+                        shared_by=["dummy-layer"],
+                    )
+                    for i in range(num_layers)
+                ],
+                # allocate_uniform_kv_caches does not use this
+                kv_cache_groups=[],
+            )
+
+            with set_current_vllm_config(vllm_config):
+                _, cross_layers_kv_cache, _ = (
+                    KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
+                        kv_cache_config=kv_cache_config,
+                        attn_groups=[
+                            [
+                                AttentionGroup(
+                                    backend=backend_cls,
+                                    layer_names=[],
+                                    kv_cache_spec=kv_cache_spec,
+                                    kv_cache_group_id=0,
+                                )
+                            ]
+                        ],
+                        cache_dtype=torch.bfloat16,
+                        device=torch.cuda.current_device(),
+                        kernel_block_sizes=[block_size],
+                    )
+                )
+            # Store tensor info for validation
+            expected_tensor_size = (
+                cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
+            )
+            expected_base_addrs = [
+                cross_layers_kv_cache.data_ptr(),
+            ]
+            expected_num_entries = 1
+
+            expected_blocks_count = 8
+
+            kv_caches = {"all-layers": cross_layers_kv_cache}
+
+        else:
+            # Create test kv cache tensors using proper backend shape
+            kv_cache_shape = backend_cls.get_kv_cache_shape(
+                num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+            )
+            shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+            unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+            kv_caches = {
+                "layer0": shared_tensor,
+                "layer1": unique_tensor,
+                "layer2": shared_tensor,
+            }
+
+            # Store tensor info for validation
+
+            test_shape = backend_cls.get_kv_cache_shape(
+                num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
+            )
+            is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
+
+            if is_blocks_first:
+                expected_tensor_size = (
+                    shared_tensor.element_size() * shared_tensor.numel()
+                )
+                expected_base_addrs = [
+                    shared_tensor.data_ptr(),
+                    unique_tensor.data_ptr(),
+                ]
+                expected_num_entries = 2
+            else:
+                expected_tensor_size = (
+                    shared_tensor[0].element_size() * shared_tensor[0].numel()
+                )
+                expected_base_addrs = [
+                    shared_tensor[0].data_ptr(),
+                    shared_tensor[1].data_ptr(),
+                    unique_tensor[0].data_ptr(),
+                    unique_tensor[1].data_ptr(),
+                ]
+                expected_num_entries = 4
+            expected_blocks_count = 8
+
         # Execute register_kv_caches
         connector.register_kv_caches(kv_caches)
 
@@ -1489,16 +1591,19 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
         blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
 
         # Validate blocks_data structure and size
-        expected_blocks_count = 8
         assert len(blocks_data) == expected_blocks_count, (
             f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
         )
 
-        num_blocks = 2
-        if is_blocks_first:
-            expected_block_len = expected_tensor_size // num_blocks // 2
-        else:
+        if connector.prefer_cross_layer_blocks:
+            num_blocks = 8
             expected_block_len = expected_tensor_size // num_blocks
+        else:
+            num_blocks = 2
+            if is_blocks_first:
+                expected_block_len = expected_tensor_size // num_blocks // 2
+            else:
+                expected_block_len = expected_tensor_size // num_blocks
 
         for i, block_entry in enumerate(blocks_data):
             block_start_addr, block_len, tp_rank = block_entry
@@ -1507,6 +1612,8 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
                 f"got {block_len}"
             )
 
+        assert connector.connector_worker.block_size == 16
+
 
 class FakePlatform(Platform):
     device_type: str = "oot"
@@ -2049,6 +2156,17 @@ def test_compatibility_hash_validation(
     )
     decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
     decode_worker = decode_connector.connector_worker
+    kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
+        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    )
+    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    kv_caches = {
+        "layer0": shared_tensor,
+        "layer1": unique_tensor,
+        "layer2": shared_tensor,
+    }
+    decode_connector.register_kv_caches(kv_caches)
 
     remote_config_params: dict[str, Any] = {
         "model": "facebook/opt-125m",
@@ -2071,7 +2189,9 @@ def test_compatibility_hash_validation(
                 )
             )
         remote_hash = compute_nixl_compatibility_hash(
-            remote_vllm_config, decode_worker.backend_name
+            remote_vllm_config,
+            decode_worker.backend_name,
+            decode_worker.kv_topo.cross_layers_blocks,
         )
 
     prefill_block_size = config_overrides.get("block_size", 16)
@@ -2150,6 +2270,27 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
     decode_worker = decode_connector.connector_worker
 
+    backend = get_current_attn_backend(local_vllm_config)
+    test_shape = backend.get_kv_cache_shape(
+        num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
+    )
+    decode_worker.kv_topo = TpKVTopology(
+        tp_rank=decode_worker.tp_rank,
+        engine_id=decode_worker.engine_id,
+        remote_tp_size=decode_worker._tp_size,  # shared state
+        remote_block_size=decode_worker._block_size,  # shared state
+        is_mla=decode_worker.use_mla,
+        total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
+        attn_backend=backend,
+        tensor_shape=test_shape,
+    )
+
+    decode_worker.compat_hash = compute_nixl_compatibility_hash(
+        decode_worker.vllm_config,
+        decode_worker.backend_name,
+        decode_worker.kv_topo.cross_layers_blocks,
+    )
+
     if error_scenario == "handshake_decode_error":
         msg_bytes = b"this is not valid msgpack data"
     elif error_scenario == "handshake_validation_error":

tests/v1/spec_decode/test_eagle.py

@@ -27,6 +27,7 @@ from vllm.config.load import LoadConfig
 from vllm.model_executor.models.llama import LlamaForCausalLM
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.spec_decode.draft_model import DraftModelProposer
 from vllm.v1.spec_decode.eagle import EagleProposer
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -34,6 +35,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 model_dir = "meta-llama/Llama-3.1-8B-Instruct"
 eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
 eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
+ar_draft_model_dir = "amd/PARD-Llama-3.2-1B"  # Compatible with parallel and AR drafting
 
 
 def _create_proposer(
@@ -41,11 +43,19 @@ def _create_proposer(
     num_speculative_tokens: int,
     attention_backend: str | None = None,
     speculative_token_tree: list[tuple[int, ...]] | None = None,
+    parallel_drafting: bool = False,
 ) -> EagleProposer:
     model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
 
-    # Choose model directory based on method
-    draft_model_dir = eagle_dir if method == "eagle" else eagle3_dir
+    # Method-dependent setup
+    if method == "eagle":
+        draft_model_dir = eagle_dir
+    elif method == "eagle3":
+        draft_model_dir = eagle3_dir
+    elif method == "draft_model":
+        draft_model_dir = ar_draft_model_dir
+    else:
+        raise ValueError(f"Unknown method: {method}")
 
     spec_token_tree_str = None
     if speculative_token_tree is not None:
@@ -59,13 +69,18 @@ def _create_proposer(
         method=method,
         num_speculative_tokens=num_speculative_tokens,
         speculative_token_tree=spec_token_tree_str,
+        parallel_drafting=parallel_drafting,
     )
+    if parallel_drafting:
+        # Overwrite pard_token to avoid crash during init
+        speculative_config.draft_model_config.hf_config.pard_token = 0
 
+    device = current_platform.device_type
     vllm_config = VllmConfig(
         model_config=model_config,
         cache_config=CacheConfig(),
         speculative_config=speculative_config,
-        device_config=DeviceConfig(device=current_platform.device_type),
+        device_config=DeviceConfig(device=device),
         parallel_config=ParallelConfig(),
         load_config=LoadConfig(),
         scheduler_config=SchedulerConfig(
@@ -75,7 +90,10 @@ def _create_proposer(
         attention_config=AttentionConfig(backend=attention_backend),
     )
 
-    return EagleProposer(vllm_config=vllm_config, device=current_platform.device_type)
+    if "eagle" in method:
+        return EagleProposer(vllm_config=vllm_config, device=device)
+    else:
+        return DraftModelProposer(vllm_config=vllm_config, device=device)
 
 
 def test_prepare_next_token_ids():
@@ -321,6 +339,390 @@ def test_prepare_inputs_padded():
     assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)
 
 
+def test_set_inputs_first_pass_default_eagle():
+    """
+    Test for set_inputs_first_pass without extra input slots (default EAGLE).
+
+    This tests the path where needs_extra_input_slots=False, which is the
+    default EAGLE pathway. In this case:
+    - Input IDs are rotated (shifted by one)
+    - The next_token_ids are inserted at the last position of each request
+    - Positions are copied as-is
+    - Hidden states are copied as-is
+    - The CommonAttentionMetadata is returned unchanged
+
+    Setup:
+    - 3 requests with query_lens [3, 2, 4]
+    - Tokens: [a1, a2, a3, b1, b2, c1, c2, c3, c4]
+    - After rotation: [a2, a3, -, b2, -, c2, c3, c4, -]
+    - After inserting next_tokens [100, 200, 300]:
+        [a2, a3, 100, b2, 200, c2, c3, c4, 300]
+    """
+    device = torch.device(current_platform.device_type)
+
+    num_speculative_tokens = 3
+    proposer = _create_proposer("eagle", num_speculative_tokens)
+
+    # Setup batch with 3 requests
+    batch_spec = BatchSpec(
+        seq_lens=[10, 8, 12],  # Arbitrary context lengths
+        query_lens=[3, 2, 4],
+    )
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=16,
+        device=device,
+    )
+
+    # Input tensors
+    # Request 0: tokens [10, 11, 12] at positions [7, 8, 9]
+    # Request 1: tokens [20, 21] at positions [6, 7]
+    # Request 2: tokens [30, 31, 32, 33] at positions [8, 9, 10, 11]
+    target_token_ids = torch.tensor(
+        [10, 11, 12, 20, 21, 30, 31, 32, 33], dtype=torch.int32, device=device
+    )
+    target_positions = torch.tensor(
+        [7, 8, 9, 6, 7, 8, 9, 10, 11], dtype=torch.int64, device=device
+    )
+    target_hidden_states = torch.randn(
+        9, proposer.hidden_size, dtype=proposer.dtype, device=device
+    )
+    next_token_ids = torch.tensor([100, 200, 300], dtype=torch.int32, device=device)
+
+    num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(
+        target_token_ids=target_token_ids,
+        next_token_ids=next_token_ids,
+        target_positions=target_positions,
+        target_hidden_states=target_hidden_states,
+        token_indices_to_sample=None,
+        cad=common_attn_metadata,
+        num_rejected_tokens_gpu=None,
+    )
+
+    assert num_tokens == 9  # Total tokens unchanged
+
+    expected_token_indices_to_sample = torch.tensor(
+        [2, 4, 8], dtype=torch.int32, device=device
+    )
+    assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)
+
+    assert output_cad is common_attn_metadata
+
+    # Verify input_ids are rotated and next_tokens inserted
+    # Original: [10, 11, 12, 20, 21, 30, 31, 32, 33]
+    # After shift by 1: [11, 12, 12, 21, 21, 31, 32, 33, 33]
+    # After inserting at last indices [2, 4, 8]: [11, 12, 100, 21, 200, 31, 32, 33, 300]
+    expected_input_ids = torch.tensor(
+        [11, 12, 100, 21, 200, 31, 32, 33, 300], dtype=torch.int32, device=device
+    )
+    assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)
+
+    # Verify positions are copied as-is
+    assert torch.equal(proposer.positions[:num_tokens], target_positions)
+
+    # Verify hidden states are copied as-is
+    assert torch.equal(proposer.hidden_states[:num_tokens], target_hidden_states)
+
+
+def test_set_inputs_first_pass_draft_model():
+    """
+    Test for set_inputs_first_pass with a draft model (extra input slots,
+    no shift).
+
+    This tests the path where needs_extra_input_slots=True and
+    shift_input_ids=False (draft model case). In this case:
+    - Input IDs are NOT shifted
+    - Each request gets extra_slots_per_request (1) new slots
+    - The kernel handles copying tokens and inserting bonus/padding tokens
+    - A new CommonAttentionMetadata is returned with updated query_start_loc
+
+    Setup:
+    - 2 requests
+    - Request 0: tokens [10, 11, 12] at positions [0, 1, 2]
+      - Only tokens [10, 11] are "valid" (query_end_loc=1),
+        token 12 is a rejected token from previous speculation
+    - Request 1: tokens [20, 21] at positions [0, 1], both valid.
+      - Note: this is less than num_speculative_tokens (2) to ensure
+        we handle variable lengths correctly.
+    - next_token_ids: [100, 200] (bonus tokens)
+
+    With extra_slots_per_request=1 and shift=False:
+    Expected output layout:
+    Request 0 (indices 0-3):
+      - idx 0: token 10, pos 0
+      - idx 1: token 11, pos 1
+      - idx 2: token 100, pos 2 (bonus token)
+      - idx 3: padding_token_id, is_rejected=True
+    Request 1 (indices 4-6):
+      - idx 4: token 20, pos 0
+      - idx 5: token 21, pos 1
+      - idx 6: token 200, pos 2 (bonus token)
+    """
+    device = torch.device(current_platform.device_type)
+
+    num_speculative_tokens = 2
+    block_size = 16
+
+    # Create a proposer configured as a draft model (pass_hidden_states=False)
+    # We need to mock this since _create_proposer defaults to EAGLE
+    proposer = _create_proposer("draft_model", num_speculative_tokens)
+
+    proposer.parallel_drafting_token_id = 0
+    proposer.is_rejected_token_mask = torch.zeros(
+        proposer.max_num_tokens, dtype=torch.bool, device=device
+    )
+    proposer.is_masked_token_mask = torch.zeros(
+        proposer.max_num_tokens, dtype=torch.bool, device=device
+    )
+
+    # Mock the attn_metadata_builder to avoid needing the full model setup
+    mock_kv_cache_spec = mock.MagicMock()
+    mock_kv_cache_spec.block_size = block_size
+    mock_builder = mock.MagicMock()
+    mock_builder.kv_cache_spec = mock_kv_cache_spec
+    proposer.attn_metadata_builder = mock_builder
+
+    # Request 0: query_len=3 (but 1 rejected), Request 1: query_len=2
+    batch_spec = BatchSpec(
+        seq_lens=[3, 2],
+        query_lens=[3, 2],
+    )
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=block_size,
+        device=device,
+        arange_block_indices=True,  # Use predictable block indices
+    )
+
+    # Input tensors
+    target_token_ids = torch.tensor(
+        [10, 11, 12, 20, 21], dtype=torch.int32, device=device
+    )
+    target_positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device)
+    target_hidden_states = torch.randn(
+        5, proposer.hidden_size, dtype=proposer.dtype, device=device
+    )
+    next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device)
+
+    num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device)
+
+    num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(
+        target_token_ids=target_token_ids,
+        next_token_ids=next_token_ids,
+        target_positions=target_positions,
+        target_hidden_states=target_hidden_states,
+        token_indices_to_sample=None,
+        cad=common_attn_metadata,
+        num_rejected_tokens_gpu=num_rejected_tokens_gpu,
+    )
+
+    assert proposer.net_num_new_slots_per_request == 1
+    assert proposer.needs_extra_input_slots
+
+    # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size
+    assert num_tokens == 7
+
+    # Request 0: [10, 11, 100, padding_token (0)]
+    # Request 1: [20, 21, 200]
+    # Combined: [10, 11, 100, 0, 20, 21, 200]
+    expected_input_ids = torch.tensor(
+        [10, 11, 100, 0, 20, 21, 200], dtype=torch.int32, device=device
+    )
+    assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)
+
+    # Verify positions
+    # Request 0: [0, 1, 2, 0 (don't care)]
+    # Request 1: [0, 1, 2]
+    # Combined: [0, 1, 2, 0, 0, 1, 2]
+    expected_positions = torch.tensor(
+        [0, 1, 2, 0, 0, 1, 2], dtype=torch.int64, device=device
+    )
+    assert torch.equal(
+        proposer.positions[:num_tokens],
+        expected_positions,
+    )
+
+    # Verify rejection mask
+    expected_is_rejected = torch.zeros(7, dtype=torch.bool, device=device)
+    expected_is_rejected[3] = True  # padding token at index 3
+    assert torch.equal(
+        proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected
+    )
+
+    # Verify masked token mask (should all be False for non-parallel drafting)
+    expected_is_masked = torch.zeros(7, dtype=torch.bool, device=device)
+    assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked)
+
+    # Verify token_indices_to_sample (bonus tokens at indices 2 and 6)
+    expected_token_indices_to_sample = torch.tensor(
+        [2, 6], dtype=torch.int32, device=device
+    )
+    assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)
+
+    # Verify the new CAD has updated query_start_loc
+    # Original: [0, 3, 5] -> New: [0, 4, 7] (each request gains 1 slot)
+    expected_query_start_loc = torch.tensor([0, 4, 7], dtype=torch.int32, device=device)
+    assert torch.equal(output_cad.query_start_loc, expected_query_start_loc)
+
+
+def test_set_inputs_first_pass_parallel_drafting():
+    """
+    Test for set_inputs_first_pass with parallel drafting (extra input slots,
+    with shift).
+
+    This tests the path where needs_extra_input_slots=True and
+    shift_input_ids=True (parallel drafting case). In this case:
+    - Input IDs ARE shifted (like default EAGLE)
+    - Each request gets extra_slots_per_request (3) new slots
+    - Parallel drafting tokens are inserted and marked as masked
+    - Hidden states are mapped correctly
+
+    Setup:
+    - 2 requests with query_lens [4, 4] (1 bonus + 3 spec tokens each)
+    - Request 0: tokens [10, 11, 12, 13] at positions [5, 6, 7, 8]
+      - Only tokens [10, 11, 12] are "valid", token 13 is rejected
+    - Request 1: tokens [20, 21, 22, 23] at positions [10, 11, 12, 13], all valid.
+    - next_token_ids: [100, 200] (bonus tokens)
+
+    With shift_input_ids=True, extra_slots_per_request=3:
+    Expected output layout:
+    Request 0 (6 output slots = 4 - 1 + 3):
+      - idx 0-2: shifted tokens [11, 12, 100]
+      - idx 3-4: parallel_drafting_tokens, is_masked=True
+      - idx 5: padding_token, is_rejected=True
+    Request 1 (6 output slots = 4 - 1 + 3):
+      - idx 6-8: shifted tokens [21, 22, 23]
+      - idx 9: bonus token 200
+      - idx 10-11: parallel_drafting_tokens, is_masked=True
+    """
+    device = torch.device(current_platform.device_type)
+
+    num_speculative_tokens = 3
+    block_size = 16
+
+    proposer = _create_proposer("eagle", num_speculative_tokens, parallel_drafting=True)
+
+    # Override to simulate parallel drafting behavior
+    proposer.parallel_drafting_token_id = -2
+    proposer.parallel_drafting_hidden_state_tensor = torch.zeros(
+        proposer.hidden_size, dtype=proposer.dtype, device=device
+    )
+    proposer.is_rejected_token_mask = torch.zeros(
+        proposer.max_num_tokens, dtype=torch.bool, device=device
+    )
+    proposer.is_masked_token_mask = torch.zeros(
+        proposer.max_num_tokens, dtype=torch.bool, device=device
+    )
+
+    # Mock the attn_metadata_builder
+    mock_kv_cache_spec = mock.MagicMock()
+    mock_kv_cache_spec.block_size = block_size
+    mock_builder = mock.MagicMock()
+    mock_builder.kv_cache_spec = mock_kv_cache_spec
+    proposer.attn_metadata_builder = mock_builder
+
+    # Request 0: query_len=4 (1 rejected), Request 1: query_len=4 (all valid)
+    batch_spec = BatchSpec(
+        seq_lens=[9, 14],
+        query_lens=[4, 4],
+    )
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=block_size,
+        device=device,
+        arange_block_indices=True,
+    )
+
+    # Input tensors
+    target_token_ids = torch.tensor(
+        [10, 11, 12, 13, 20, 21, 22, 23], dtype=torch.int32, device=device
+    )
+    target_positions = torch.tensor(
+        [5, 6, 7, 8, 10, 11, 12, 13], dtype=torch.int64, device=device
+    )
+    target_hidden_states = torch.arange(
+        8 * proposer.hidden_size, dtype=proposer.dtype, device=device
+    ).view(8, proposer.hidden_size)
+    next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device)
+
+    num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device)
+
+    num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass(
+        target_token_ids=target_token_ids,
+        next_token_ids=next_token_ids,
+        target_positions=target_positions,
+        target_hidden_states=target_hidden_states,
+        token_indices_to_sample=None,
+        cad=common_attn_metadata,
+        num_rejected_tokens_gpu=num_rejected_tokens_gpu,
+    )
+
+    # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size
+    # = 8 + 2 * 2 = 12
+    assert num_tokens == 12
+
+    # Request 0: [11, 12, 100, -2, -2, 0(padding)]
+    # Request 1: [21, 22, 23, 200, -2, -2]
+    expected_input_ids = torch.tensor(
+        [11, 12, 100, -2, -2, 0, 21, 22, 23, 200, -2, -2],
+        dtype=torch.int32,
+        device=device,
+    )
+    assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids)
+
+    # Verify positions
+    # Request 0: [5, 6, 7, 8, 9, 0 (don't care)]
+    # Request 1: [10, 11, 12, 13, 14, 15]
+    expected_positions = torch.tensor(
+        [5, 6, 7, 8, 9, 0, 10, 11, 12, 13, 14, 15], dtype=torch.int64, device=device
+    )
+    assert torch.equal(
+        proposer.positions[:num_tokens],
+        expected_positions,
+    )
+
+    # Verify rejection mask
+    expected_is_rejected = torch.zeros(12, dtype=torch.bool, device=device)
+    expected_is_rejected[5] = True
+    assert torch.equal(
+        proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected
+    )
+
+    # Verify masked token mask (parallel drafting slots should be masked)
+    expected_is_masked = torch.zeros(12, dtype=torch.bool, device=device)
+    expected_is_masked[3] = True
+    expected_is_masked[4] = True
+    expected_is_masked[10] = True
+    expected_is_masked[11] = True
+    assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked)
+
+    # Verify token_indices_to_sample (bonus + parallel drafting tokens)
+    # Request 0: bonus at 2, parallel at 3, 4
+    # Request 1: bonus at 9, parallel at 10, 11
+    expected_token_indices_to_sample = torch.tensor(
+        [2, 3, 4, 9, 10, 11], dtype=torch.int32, device=device
+    )
+    assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample)
+
+    # Verify the new CAD has updated query_start_loc
+    # Original query_lens: [4, 4] -> Output: [6, 6]
+    expected_query_start_loc = torch.tensor(
+        [0, 6, 12], dtype=torch.int32, device=device
+    )
+    assert torch.equal(output_cad.query_start_loc, expected_query_start_loc)
+
+    # Verify masked positions have the parallel drafting hidden state (zeros)
+    parallel_drafting_hs = proposer.parallel_drafting_hidden_state_tensor
+    for i in range(num_tokens):
+        if expected_is_masked[i]:
+            assert torch.equal(proposer.hidden_states[i], parallel_drafting_hs), (
+                f"Masked position {i} should have parallel drafting hidden state"
+            )
+
+
 @pytest.mark.parametrize("method", ["eagle", "eagle3"])
 @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
 @pytest.mark.parametrize("pp_size", [1, 2])
@@ -579,7 +981,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):
         target_positions=target_positions,
         target_hidden_states=target_hidden_states,
         next_token_ids=next_token_ids,
-        last_token_indices=None,
+        token_indices_to_sample=None,
         common_attn_metadata=common_attn_metadata,
         sampling_metadata=sampling_metadata,
     )
@@ -737,7 +1139,7 @@ def test_propose_tree(spec_token_tree):
         target_positions=target_positions,
         target_hidden_states=target_hidden_states,
         next_token_ids=next_token_ids,
-        last_token_indices=None,
+        token_indices_to_sample=None,
         common_attn_metadata=common_attn_metadata,
         sampling_metadata=sampling_metadata,
     )

tests/v1/spec_decode/test_mtp.py

@@ -204,7 +204,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch):
         target_positions=target_positions,
         target_hidden_states=target_hidden_states,
         next_token_ids=next_token_ids,
-        last_token_indices=None,
+        token_indices_to_sample=None,
         common_attn_metadata=common_attn_metadata,
         sampling_metadata=sampling_metadata,
     )

tests/v1/streaming_input/test_async_llm_streaming.py

@@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
-from vllm.inputs.data import StreamingInput
+from vllm.inputs import StreamingInput
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.v1.engine.async_llm import AsyncLLM

vllm/config/__init__.py

@@ -47,6 +47,7 @@ from vllm.config.vllm import (
     get_layers_from_vllm_config,
     set_current_vllm_config,
 )
+from vllm.config.weight_transfer import WeightTransferConfig
 
 # __all__ should only contain classes and functions.
 # Types and globals should be imported from their respective modules.
@@ -111,4 +112,5 @@ __all__ = [
     "get_current_vllm_config_or_none",
     "set_current_vllm_config",
     "get_layers_from_vllm_config",
+    "WeightTransferConfig",
 ]

vllm/config/speculative.py

@@ -116,9 +116,16 @@ class SpeculativeConfig:
     """Minimum size of ngram token window when using Ngram proposer, if
     provided. Defaults to 1."""
 
+    # Alternative drafting strategies
     speculative_token_tree: str | None = None
     """Specifies the tree structure for speculative token generation.
     """
+    parallel_drafting: bool = False
+    """Enable parallel drafting, where all speculative tokens are generated
+    in parallel rather than sequentially. This can improve performance but
+    requires the speculative model be trained to support parallel drafting.
+    Only compatible with EAGLE and draft model methods."""
+
     # required configuration params passed from engine
     target_model_config: SkipValidation[ModelConfig] = None  # type: ignore
     """The configuration of the target model."""

vllm/config/vllm.py

@@ -42,6 +42,7 @@ from .scheduler import SchedulerConfig
 from .speculative import EagleModelTypes, SpeculativeConfig
 from .structured_outputs import StructuredOutputsConfig
 from .utils import SupportsHash, config, replace
+from .weight_transfer import WeightTransferConfig
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
@@ -255,6 +256,9 @@ class VllmConfig:
     performance. -02 is used by defult. See  OptimizationLevel for full
     description."""
 
+    weight_transfer_config: WeightTransferConfig | None = None
+    """The configurations for weight transfer during RL training."""
+
     def compute_hash(self) -> str:
         """
         WARNING: Whenever a new field is added to this config,
@@ -600,10 +604,13 @@ class VllmConfig:
             # Currently, async scheduling only support eagle speculative
             # decoding.
             if self.speculative_config is not None:
-                if self.speculative_config.method not in get_args(EagleModelTypes):
+                if (
+                    self.speculative_config.method not in get_args(EagleModelTypes)
+                    and self.speculative_config.method != "draft_model"
+                ):
                     raise ValueError(
                         "Currently, async scheduling is only supported "
-                        "with EAGLE/MTP kind of speculative decoding."
+                        "with EAGLE/MTP/Draft Model kind of speculative decoding."
                     )
                 if self.speculative_config.disable_padded_drafter_batch:
                     raise ValueError(
@@ -1294,16 +1301,21 @@ class VllmConfig:
         computed_compile_ranges_split_points = []
 
         # The upper bound of the compile ranges is the max_num_batched_tokens.
-        # For speculative decoding with draft model, the compile range must be extended
-        # by 1 for each sequence.
+        # For speculative decoding, the compile range must be extended
+        # - Sequential: + 1 * max_num_seqs (one draft token per iteration)
+        # - Parallel draft: + num_speculative_tokens * max_num_seqs
         compile_range_end = self.scheduler_config.max_num_batched_tokens
         if compile_range_end is not None:
-            do_extend: bool = (
-                self.speculative_config is not None
-                and self.speculative_config.uses_draft_model()
-            )
-            if do_extend:
-                compile_range_end += self.scheduler_config.max_num_seqs
+            if self.speculative_config is not None and (
+                self.speculative_config.uses_draft_model()
+                or self.speculative_config.use_eagle()
+            ):
+                multiplier = (
+                    self.speculative_config.num_speculative_tokens
+                    if self.speculative_config.parallel_drafting
+                    else 1
+                )
+                compile_range_end += multiplier * self.scheduler_config.max_num_seqs
 
             computed_compile_ranges_split_points.append(compile_range_end)
 

vllm/config/weight_transfer.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from typing import Literal
+
+from vllm.config.utils import config
+
+
+@config
+@dataclass
+class WeightTransferConfig:
+    """Configuration for weight transfer during RL training."""
+
+    backend: Literal["nccl"] = "nccl"
+    """The backend to use for weight transfer."""

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -14,6 +14,7 @@ from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vll
 from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.platforms import current_platform
 from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
@@ -192,8 +193,6 @@ def copy_kv_blocks(
         dst_device=dst_device,
     )
 
-    from vllm.platforms import current_platform
-
     if direction == "h2d":
         copy_fn = current_platform.insert_blocks_to_device
     else:
@@ -316,19 +315,53 @@ class TpKVTopology:
     attn_backend: type[AttentionBackend]
     engine_id: EngineId
     remote_block_size: dict[EngineId, int]
+    tensor_shape: torch.Size | None = None
 
     def __post_init__(self):
         # Figure out whether the first dimension of the cache is K/V
         # or num_blocks. This is used to register the memory regions correctly.
+        _MOCK_BLOCK_SIZE = 16
         kv_cache_shape = self.attn_backend.get_kv_cache_shape(
-            num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
+            num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1
         )
+        logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
         # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
         # we just mock num_blocks to 1 for the dimension check below.
         self._is_kv_layout_blocks_first = (
             len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
         )
 
+        self._cross_layers_blocks = False
+        if self.tensor_shape is not None:
+            self._cross_layers_blocks = (
+                len(self.tensor_shape) == len(kv_cache_shape) + 1
+            )
+
+        if self._cross_layers_blocks:
+            logger.debug("Using cross-layer KV cache")
+            # prepend layers dimension
+            _MOCK_NUM_LAYERS = 80
+            kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
+            try:
+                kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(
+                    include_num_layers_dimension=self._cross_layers_blocks
+                )
+            except (AttributeError, NotImplementedError):
+                kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
+
+            # In case of cross layers permute kv_cache_shape according to
+            # stride_order to retrieve physical position of block_size
+            kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+
+        # In the default non-cross layers layout the block_size position
+        # is logical while in the cross layers case it is the physical
+        # position. This matches the shape of the actual kv cache tensors
+        # passed at register_kv_caches()/register_cross_layers_kv_cache()
+        block_size_position = kv_cache_shape.index(_MOCK_BLOCK_SIZE)
+
+        assert block_size_position is not None
+        self._block_size_position = -(len(kv_cache_shape) - block_size_position)
+
     @property
     def is_kv_layout_blocks_first(self) -> bool:
         return self._is_kv_layout_blocks_first
@@ -336,7 +369,9 @@ class TpKVTopology:
     @property
     def split_k_and_v(self) -> bool:
         # Whether to register regions for K and V separately (when present).
-        return not (self.is_mla or self.is_kv_layout_blocks_first)
+        return not (
+            self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
+        )
 
     @property
     def tp_size(self) -> int:
@@ -346,6 +381,14 @@ class TpKVTopology:
     def block_size(self) -> int:
         return self.remote_block_size[self.engine_id]
 
+    @property
+    def cross_layers_blocks(self) -> bool:
+        return self._cross_layers_blocks
+
+    @property
+    def block_size_position(self) -> int:
+        return self._block_size_position
+
     def tp_ratio(
         self,
         remote_tp_size: int,

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

@@ -54,7 +54,7 @@ from vllm.forward_context import ForwardContext
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
-from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
 from vllm.v1.attention.backends.utils import get_kv_cache_layout
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.block_table import BlockTable
@@ -173,7 +173,7 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata):
 
 
 def compute_nixl_compatibility_hash(
-    vllm_config: VllmConfig, attn_backend_name: str
+    vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
 ) -> str:
     """
     Compute compatibility hash for NIXL KV transfer.
@@ -216,6 +216,7 @@ def compute_nixl_compatibility_hash(
         # Attention backend and KV cache dtype affect memory layout
         "attn_backend_name": attn_backend_name,
         "cache_dtype": str(cache_config.cache_dtype),
+        "cross_layers_blocks": cross_layers_blocks,
     }
 
     compat_hash = hash_factors(factors)
@@ -298,6 +299,26 @@ class NixlConnectorMetadata(KVConnectorMetadata):
 
 
 class NixlConnector(KVConnectorBase_V1):
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        backend = get_current_attn_backend(self._vllm_config)
+        if backend.get_name() not in (
+            "FLASH_ATTN",
+            "FLASHINFER",
+        ):
+            return False
+
+        # For now there is no benefit to run cross layers when backend
+        # does not support on HND
+        if get_kv_cache_layout() != "HND":
+            return False
+
+        extra_config = self.kv_transfer_config.kv_connector_extra_config
+        return (
+            str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
+            == "true"
+        )
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -309,7 +330,7 @@ class NixlConnector(KVConnectorBase_V1):
         assert vllm_config.kv_transfer_config is not None
         assert vllm_config.kv_transfer_config.engine_id is not None
         self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
-
+        self.kv_transfer_config = vllm_config.kv_transfer_config
         if role == KVConnectorRole.SCHEDULER:
             self.connector_scheduler: NixlConnectorScheduler | None = (
                 NixlConnectorScheduler(vllm_config, self.engine_id)
@@ -395,6 +416,16 @@ class NixlConnector(KVConnectorBase_V1):
         assert self.connector_worker is not None
         self.connector_worker.register_kv_caches(kv_caches)
 
+    def register_cross_layers_kv_cache(
+        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
+    ):
+        assert self.connector_worker is not None
+
+        cross_layer_name = "ALL_LAYERS"
+        kv_caches = {cross_layer_name: kv_cache}
+
+        self.connector_worker.register_kv_caches(kv_caches)
+
     def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
         assert self.connector_worker is not None
         self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
@@ -976,20 +1007,17 @@ class NixlConnectorWorker:
 
         # Get the attention backend from the first layer
         # NOTE (NickLucche) models with multiple backends are not supported yet
-        backend = get_current_attn_backend(vllm_config)
+        self.attn_backend = get_current_attn_backend(vllm_config)
 
-        self.backend_name = backend.get_name()
+        self.backend_name = self.attn_backend.get_name()
         self.kv_cache_layout = get_kv_cache_layout()
         self.host_buffer_kv_cache_layout = self.kv_cache_layout
         logger.debug("Detected attention backend %s", self.backend_name)
         logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
 
-        self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config, self.backend_name
-        )
-        self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
-            "enforce_handshake_compat", True
-        )
+        # lazy initialized in register_kv_caches
+        self.compat_hash: str | None = None
+        self.kv_topo: TpKVTopology | None = None
 
         self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
         self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
@@ -998,17 +1026,12 @@ class NixlConnectorWorker:
         self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
         self.xfer_stats = NixlKVConnectorStats()
 
-        self.kv_topo = TpKVTopology(
-            tp_rank=self.tp_rank,
-            engine_id=self.engine_id,
-            remote_tp_size=self._tp_size,  # shared state
-            remote_block_size=self._block_size,  # shared state
-            is_mla=self.use_mla,
-            total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
-            attn_backend=backend,
-        )
         self._physical_blocks_per_logical_kv_block = 1
 
+        self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
+            "enforce_handshake_compat", True
+        )
+
     def _nixl_handshake(
         self,
         host: str,
@@ -1022,6 +1045,7 @@ class NixlConnectorWorker:
         # Regardless, only handshake with the remote TP rank(s) that current
         # local rank will read from. Note that With homogeneous TP,
         # this happens to be the same single rank_i.
+        assert self.kv_topo is not None
         p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)
         remote_rank_to_agent_name = {}
         path = make_zmq_path("tcp", host, port)
@@ -1059,6 +1083,7 @@ class NixlConnectorWorker:
                 )
 
                 # Check compatibility hash BEFORE decoding agent metadata
+                assert self.compat_hash is not None
                 if (
                     self.enforce_compat_hash
                     and handshake_payload.compatibility_hash != self.compat_hash
@@ -1267,6 +1292,20 @@ class NixlConnectorWorker:
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
 
+        self.kv_topo = TpKVTopology(
+            tp_rank=self.tp_rank,
+            engine_id=self.engine_id,
+            remote_tp_size=self._tp_size,  # shared state
+            remote_block_size=self._block_size,  # shared state
+            is_mla=self.use_mla,
+            total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
+            attn_backend=self.attn_backend,
+            tensor_shape=next(iter(kv_caches.values())).shape,
+        )
+        self.compat_hash = compute_nixl_compatibility_hash(
+            self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
+        )
+
         if self.use_host_buffer:
             self.initialize_host_xfer_buffer(kv_caches=kv_caches)
             assert len(self.host_xfer_buffers) == len(kv_caches), (
@@ -1301,29 +1340,24 @@ class NixlConnectorWorker:
         # (roughly 8KB vs 5KB).
         # Conversely for FlashInfer, K and V are registered in the same region
         # to better exploit the memory layout (ie num_blocks is the first dim).
-        split_k_and_v = self.kv_topo.split_k_and_v
         tensor_size_bytes = None
 
-        # TODO (NickLucche): Get kernel_block_size in a cleaner way
-        # NHD default "view" for non-MLA cache
-        if self.device_type == "cpu":
-            block_size_position = -2
-        else:
-            block_size_position = -2 if self.use_mla else -3
-
         # Enable different block lengths for different layers when MLA is used.
         self.block_len_per_layer = list[int]()
         self.slot_size_per_layer = list[int]()  # HD bytes in kv terms
         for layer_name, cache_or_caches in xfer_buffers.items():
-            cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
-
+            cache_list = (
+                cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]
+            )
             for cache in cache_list:
                 base_addr = cache.data_ptr()
                 if base_addr in seen_base_addresses:
                     continue
 
-                kernel_block_size = cache.shape[block_size_position]
-
+                logger.debug(
+                    "Registering layer %s with cache shape: %s", layer_name, cache.shape
+                )
+                kernel_block_size = cache.shape[self.kv_topo.block_size_position]
                 if self.block_size != kernel_block_size:
                     logger.info_once(
                         "User-specified logical block size (%s) does not match"
@@ -1385,6 +1419,7 @@ class NixlConnectorWorker:
 
         self.device_kv_caches = kv_caches
         self.dst_num_blocks[self.engine_id] = self.num_blocks
+
         if self.kv_topo.is_kv_layout_blocks_first:
             for i in range(len(self.slot_size_per_layer)):
                 assert self.slot_size_per_layer[i] % 2 == 0
@@ -1440,6 +1475,7 @@ class NixlConnectorWorker:
             block_size=self.block_size,
         )
         # Wrap metadata in payload with hash for defensive decoding
+        assert self.compat_hash is not None
         encoder = msgspec.msgpack.Encoder()
         self.xfer_handshake_metadata = NixlHandshakePayload(
             compatibility_hash=self.compat_hash,
@@ -1461,6 +1497,8 @@ class NixlConnectorWorker:
         register another local_xfer_handler using remote block len to ensure
         data copy correctness.
         """
+        assert self.kv_topo is not None
+
         block_size_ratio = self.block_size // block_size
         blocks_data = []
         for i, base_addr in enumerate(self.seen_base_addresses):
@@ -1573,6 +1611,7 @@ class NixlConnectorWorker:
         # remote:               | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
         # local origin:|          0|          1|          8|         12|
         # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
+        assert self.kv_topo is not None
         block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)
 
         if engine_id not in self.dst_num_blocks:
@@ -1701,6 +1740,7 @@ class NixlConnectorWorker:
         remote_engine_id = nixl_agent_meta.engine_id
 
         assert self._tp_size[remote_engine_id] == remote_tp_size
+        assert self.kv_topo is not None
 
         tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
         block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
@@ -1837,6 +1877,7 @@ class NixlConnectorWorker:
         if len(self.device_kv_caches) == 0:
             return
         assert block_size_ratio >= 1, "Only nP < nD supported currently."
+        assert self.kv_topo is not None
         if self.enable_permute_local_kv and block_size_ratio > 1:
             logger.debug(
                 "Post-processing device kv cache on receive by converting "
@@ -1856,7 +1897,7 @@ class NixlConnectorWorker:
                 block_size_ratio,
             )
 
-        split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
+        split_k_and_v = self.kv_topo.split_k_and_v
 
         for block_ids in block_ids_list:
             indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
@@ -1881,6 +1922,7 @@ class NixlConnectorWorker:
         The scheduler process (via the MultiprocExecutor) will use this output
         to track which workers are done.
         """
+        assert self.kv_topo is not None
         done_sending = self._get_new_notifs()
         done_recving = self._pop_done_transfers(self._recving_transfers)
 
@@ -1950,6 +1992,7 @@ class NixlConnectorWorker:
         are reading from the same producer (heterogeneous TP scenario), wait
         for all consumers to be done pulling.
         """
+        assert self.kv_topo is not None
         notified_req_ids: set[str] = set()
         for notifs in self.nixl_wrapper.get_new_notifs().values():
             for notif in notifs:
@@ -2109,7 +2152,7 @@ class NixlConnectorWorker:
                 self._reqs_to_send[req_id] = expiration_time
 
     def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
-        assert meta.remote is not None
+        assert meta.remote is not None and self.kv_topo is not None
         remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
             meta.remote.engine_id
         )
@@ -2182,6 +2225,7 @@ class NixlConnectorWorker:
         Post a READ point-to-point xfer request from a single local worker to
         a single remote worker.
         """
+        assert self.kv_topo is not None
         block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
         if block_size_ratio > 1:
             local_block_ids = self.get_mapped_blocks(
@@ -2279,15 +2323,15 @@ class NixlConnectorWorker:
 
                 # Get descs ids for the layer.
                 layer_local_desc_ids = self._get_block_descs_ids(
-                    dst_engine_id,
+                    self.engine_id,
                     layer_local_block_ids,
                     layer_idx,
+                    block_size_ratio=block_size_ratio,
                 )
                 layer_remote_desc_ids = self._get_block_descs_ids(
-                    self.engine_id,
+                    dst_engine_id,
                     layer_remote_block_ids,
                     layer_idx,
-                    block_size_ratio=block_size_ratio,
                 )
 
                 local_descs_list.append(layer_local_desc_ids)
@@ -2414,6 +2458,7 @@ class NixlConnectorWorker:
         For FlashInfer, this is half the length of the whole block, as K and V
         share the same region.
         """
+        assert self.kv_topo is not None
         if self.kv_topo.is_kv_layout_blocks_first:
             # For indexing only half (either just the K or V part).
             block_len = self.block_len_per_layer[layer_idx] // 2

vllm/distributed/weight_transfer/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Weight transfer engines for syncing model weights from trainers
+to inference workers.
+"""
+
+from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
+
+__all__ = [
+    "WeightTransferEngineFactory",
+]

vllm/distributed/weight_transfer/base.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Base class for weight transfer engines."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import KW_ONLY, dataclass, field
+from typing import Any, Generic, TypeVar
+
+import torch
+
+from vllm.config.parallel import ParallelConfig
+from vllm.config.weight_transfer import WeightTransferConfig
+
+TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo")
+TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo")
+
+
+# Base protocols for backend-specific dataclasses
+@dataclass
+class WeightTransferInitInfo(ABC):  # noqa: B024
+    """Base class for backend-specific initialization info."""
+
+    pass
+
+
+@dataclass
+class WeightTransferUpdateInfo(ABC):  # noqa: B024
+    """Base class for backend-specific weight update info."""
+
+    _: KW_ONLY
+    is_checkpoint_format: bool = True
+    """Set to True if weights are in checkpoint/original model format and need
+    layerwise processing. Set to False if weights have already been processed
+    into kernel format (repacking, renaming, etc.)."""
+
+
+# API-level request classes (accept dicts for backend-agnostic serialization)
+@dataclass
+class WeightTransferInitRequest:
+    """API-level weight transfer initialization request."""
+
+    init_info: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class WeightTransferUpdateRequest:
+    """API-level weight update request."""
+
+    update_info: dict[str, Any] = field(default_factory=dict)
+
+
+class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
+    """
+    Base class for weight transfer engines that handle transport of model weights
+    from a trainer to inference workers.
+
+    This abstraction separates weight transfer transport logic from the worker
+    implementation, allowing different backends (NCCL, CUDA IPC[TODO], RDMA[TODO]) to be
+    plugged in.
+
+    Subclasses should define:
+        init_info_cls: Type of backend-specific initialization info
+        update_info_cls: Type of backend-specific update info
+    """
+
+    # Subclasses should override these class attributes
+    init_info_cls: type[TInitInfo]
+    update_info_cls: type[TUpdateInfo]
+
+    def __init__(
+        self, config: WeightTransferConfig, parallel_config: ParallelConfig
+    ) -> None:
+        """
+        Initialize the weight transfer engine.
+
+        Args:
+            config: The configuration for the weight transfer engine
+            parallel_config: The configuration for the parallel setup
+        """
+        self.config = config
+        self.parallel_config = parallel_config
+
+    def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
+        """
+        Construct typed init info from dict with validation.
+
+        Args:
+            init_dict: Dictionary containing backend-specific initialization parameters
+
+        Returns:
+            Typed backend-specific init info dataclass
+
+        Raises:
+            ValueError: If init_dict is invalid for this backend
+        """
+        try:
+            return self.init_info_cls(**init_dict)
+        except TypeError as e:
+            raise ValueError(
+                f"Invalid init_info for {self.__class__.__name__}: {e}"
+            ) from e
+
+    def parse_update_info(self, update_dict: dict[str, Any]) -> TUpdateInfo:
+        """
+        Construct typed update info from dict with validation.
+
+        Args:
+            update_dict: Dictionary containing backend-specific update parameters
+
+        Returns:
+            Typed backend-specific update info dataclass
+
+        Raises:
+            ValueError: If update_dict is invalid for this backend
+        """
+        try:
+            return self.update_info_cls(**update_dict)
+        except TypeError as e:
+            raise ValueError(
+                f"Invalid update_info for {self.__class__.__name__}: {e}"
+            ) from e
+
+    @abstractmethod
+    def init_transfer_engine(self, init_info: TInitInfo) -> None:
+        """
+        Initialize the weight transfer mechanism.
+        This is called once at the beginning of training.
+
+        Args:
+            init_info: Backend-specific initialization info
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def receive_weights(
+        self,
+        update_info: TUpdateInfo,
+        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
+    ) -> None:
+        """
+        Receive weights from the trainer and load them incrementally.
+
+        Args:
+            update_info: Backend-specific update info containing parameter metadata
+                        and any backend-specific data
+            load_weights: Callable that loads weights into the model. Called
+                         incrementally for each weight to avoid OOM.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """
+        Shutdown the weight transfer engine.
+        This should be called when the worker is shutting down.
+        """
+        raise NotImplementedError

vllm/distributed/weight_transfer/factory.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Factory for weight transfer engines with lazy loading."""
+
+import importlib
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from vllm.distributed.weight_transfer.base import WeightTransferEngine
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config.parallel import ParallelConfig
+    from vllm.config.weight_transfer import WeightTransferConfig
+
+logger = init_logger(__name__)
+
+
+class WeightTransferEngineFactory:
+    """Factory for creating weight transfer engines with lazy loading.
+
+    This factory implements a registry pattern that supports:
+    - Lazy loading: Engine modules are only imported when actually needed
+    - Extensibility: Custom engines can be registered at runtime
+    - Centralized registration: All built-in engines registered in one place
+    """
+
+    _registry: dict[str, Callable[[], type[WeightTransferEngine]]] = {}
+
+    @classmethod
+    def register_engine(
+        cls,
+        name: str,
+        module_path_or_cls: str | type[WeightTransferEngine],
+        class_name: str | None = None,
+    ) -> None:
+        """Register an engine with lazy-loading or direct class reference.
+
+        Supports two calling conventions:
+        1. Lazy loading: register_engine(name, module_path, class_name)
+        2. Direct class: register_engine(name, engine_cls)
+
+        Args:
+            name: The name to register the engine under (e.g., "nccl")
+            module_path_or_cls: Either a module path string for lazy loading,
+                or the engine class directly
+            class_name: Name of the engine class (required if module_path is string)
+
+        Raises:
+            ValueError: If an engine with the same name is already registered
+        """
+        if name in cls._registry:
+            raise ValueError(f"Weight transfer engine '{name}' is already registered.")
+
+        if isinstance(module_path_or_cls, str):
+            # Lazy loading path
+            module_path = module_path_or_cls
+            if class_name is None:
+                raise ValueError(
+                    "class_name is required when registering with module path"
+                )
+
+            def loader() -> type[WeightTransferEngine]:
+                module = importlib.import_module(module_path)
+                return getattr(module, class_name)
+
+            cls._registry[name] = loader
+        else:
+            # Direct class registration
+            engine_cls = module_path_or_cls
+            cls._registry[name] = lambda: engine_cls
+
+    @classmethod
+    def create_engine(
+        cls,
+        config: "WeightTransferConfig",
+        parallel_config: "ParallelConfig",
+    ) -> WeightTransferEngine:
+        """Create a weight transfer engine instance.
+
+        Args:
+            config: Weight transfer configuration containing the backend name
+            parallel_config: Parallel configuration for the engine
+
+        Returns:
+            An initialized weight transfer engine instance
+
+        Raises:
+            ValueError: If the backend is not registered
+        """
+        backend = config.backend
+        if backend not in cls._registry:
+            available = list(cls._registry.keys())
+            raise ValueError(
+                f"Invalid weight transfer backend: {backend}. "
+                f"Available engines: {available}"
+            )
+        engine_cls = cls._registry[backend]()
+
+        logger.info(
+            "Creating weight transfer engine: %s",
+            engine_cls.__name__,
+        )
+
+        return engine_cls(config, parallel_config)
+
+
+# Register built-in weight transfer engines here.
+# Registration should be centralized to ensure lazy loading -
+# engine modules are only imported when actually used.
+
+WeightTransferEngineFactory.register_engine(
+    "nccl",
+    "vllm.distributed.weight_transfer.nccl_engine",
+    "NCCLWeightTransferEngine",
+)

vllm/distributed/weight_transfer/nccl_engine.py

@@ -0,0 +1,315 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NCCL-based weight transfer engine."""
+
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+from vllm.config.parallel import ParallelConfig
+from vllm.config.weight_transfer import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferEngine,
+    WeightTransferInitInfo,
+    WeightTransferUpdateInfo,
+)
+from vllm.distributed.weight_transfer.packed_tensor import (
+    DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+    DEFAULT_PACKED_NUM_BUFFERS,
+    packed_broadcast_consumer,
+)
+
+
+@dataclass
+class NCCLWeightTransferInitInfo(WeightTransferInitInfo):
+    """Initialization info for NCCL weight transfer backend."""
+
+    master_address: str
+    master_port: int
+    rank_offset: int
+    world_size: int
+
+
+@dataclass
+class NCCLWeightTransferUpdateInfo(WeightTransferUpdateInfo):
+    """Update info for NCCL weight transfer backend."""
+
+    names: list[str]
+    dtype_names: list[str]
+    shapes: list[list[int]]
+    packed: bool = False
+    """Whether to use packed tensor broadcasting for efficiency.
+    When True, multiple tensors are batched together before broadcasting
+    to reduce NCCL communication overhead."""
+    packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES
+    """Size in bytes for each packed tensor buffer. Default is 1GB.
+    Both producer and consumer must use the same value."""
+    packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS
+    """Number of buffers for double/triple buffering during packed transfer.
+    Both producer and consumer must use the same value."""
+
+    def __post_init__(self):
+        """Validate that all lists have the same length."""
+        num_params = len(self.names)
+        if len(self.dtype_names) != num_params:
+            raise ValueError(
+                f"`dtype_names` should be of the same size as `names`: "
+                f"got {len(self.dtype_names)} and {len(self.names)}"
+            )
+        if len(self.shapes) != num_params:
+            raise ValueError(
+                f"`shapes` should be of the same size as `names`: "
+                f"got {len(self.shapes)} and {len(self.names)}"
+            )
+
+
+class NCCLWeightTransferEngine(
+    WeightTransferEngine[NCCLWeightTransferInitInfo, NCCLWeightTransferUpdateInfo]
+):
+    """
+    Weight transfer engine using NCCL for communication between trainer and workers.
+
+    This implementation uses NCCL broadcast operations to transfer weights from
+    the trainer (rank 0) to all inference workers in a process group.
+    """
+
+    # Define backend-specific dataclass types
+    init_info_cls = NCCLWeightTransferInitInfo
+    update_info_cls = NCCLWeightTransferUpdateInfo
+
+    def __init__(
+        self, config: WeightTransferConfig, parallel_config: ParallelConfig
+    ) -> None:
+        """
+        Initialize the NCCL weight transfer engine.
+
+        Args:
+            config: The configuration for the weight transfer engine
+            parallel_config: The configuration for the parallel setup
+        """
+        super().__init__(config, parallel_config)
+        self.model_update_group: PyNcclCommunicator | None = None
+
+    def init_transfer_engine(self, init_info: NCCLWeightTransferInitInfo) -> None:
+        """
+        Initialize NCCL process group with the trainer.
+
+        Args:
+            init_info: NCCL initialization info containing master address, port,
+                      rank offset, and world size
+        """
+
+        # Calculate the global rank in the trainer-worker process group
+        # Must account for data parallel to get unique ranks across all workers
+        dp_rank = self.parallel_config.data_parallel_rank
+        world_size_per_dp = self.parallel_config.world_size  # TP * PP
+        rank_within_dp = self.parallel_config.rank
+
+        # Unique rank across all DP groups
+        worker_rank = dp_rank * world_size_per_dp + rank_within_dp
+        rank = worker_rank + init_info.rank_offset
+        # Create stateless process group
+        self.model_update_group = (
+            NCCLWeightTransferEngine._stateless_init_process_group(
+                init_info.master_address,
+                init_info.master_port,
+                rank,
+                init_info.world_size,
+                torch.cuda.current_device(),
+            )
+        )
+
+    def receive_weights(
+        self,
+        update_info: NCCLWeightTransferUpdateInfo,
+        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
+    ) -> None:
+        """
+        Receive weights from trainer via NCCL broadcast and load them incrementally.
+
+        If update_info.packed is True, uses packed tensor broadcasting for
+        efficient transfer of multiple weights in batches. Otherwise, uses simple
+        one-by-one broadcasting.
+
+        Args:
+            update_info: NCCL update info containing parameter names, dtypes, shapes,
+                        and packed flag
+            load_weights: Callable that loads weights into the model. Called
+                         incrementally for each batch of weights to avoid OOM.
+        """
+        if self.model_update_group is None:
+            raise RuntimeError(
+                "NCCL weight transfer not initialized. "
+                "Call init_transfer_engine() first."
+            )
+
+        if update_info.packed:
+            # Build iterator of (name, (shape, dtype)) from update_info
+            def state_dict_info_iterator():
+                for name, dtype_name, shape in zip(
+                    update_info.names, update_info.dtype_names, update_info.shapes
+                ):
+                    dtype = getattr(torch, dtype_name)
+                    yield (name, (shape, dtype))
+
+            packed_broadcast_consumer(
+                iterator=state_dict_info_iterator(),
+                group=self.model_update_group,
+                src=0,
+                post_unpack_func=load_weights,
+                buffer_size_bytes=update_info.packed_buffer_size_bytes,
+                num_buffers=update_info.packed_num_buffers,
+            )
+        else:
+            # Use simple one-by-one broadcasting
+            for name, dtype_name, shape in zip(
+                update_info.names, update_info.dtype_names, update_info.shapes
+            ):
+                dtype = getattr(torch, dtype_name)
+                weight = torch.empty(shape, dtype=dtype, device="cuda")
+                self.model_update_group.broadcast(
+                    weight, src=0, stream=torch.cuda.current_stream()
+                )
+                load_weights([(name, weight)])
+                del weight
+
+    def shutdown(self) -> None:
+        if self.model_update_group is not None:
+            # Clean up the communicator by removing the reference
+            self.model_update_group = None
+
+    @staticmethod
+    def trainer_send_weights(
+        iterator: Iterator[tuple[str, torch.Tensor]],
+        group: Any,
+        src: int = 0,
+        post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor]
+        | None = None,
+        packed: bool = False,
+        stream: torch.cuda.Stream | None = None,
+        packed_buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+        packed_num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
+    ) -> None:
+        """Broadcast weights from trainer to vLLM workers.
+
+        Args:
+            iterator: Iterator of model parameters. Returns (name, tensor) tuples
+            group: Process group (PyNcclCommunicator)
+            src: Source rank (default 0, trainer is typically rank 0)
+            post_iter_func: Optional function to apply to each (name, tensor) pair
+                           before broadcasting. If None, extracts just the tensor.
+            packed: Whether to use packed tensor broadcasting for efficiency.
+                   When True, multiple tensors are batched together before
+                   broadcasting to reduce NCCL communication overhead.
+            stream: CUDA stream to use for broadcasting if packed is False.
+                    If packed is True, new streams will be created for each buffer.
+            packed_buffer_size_bytes: Size in bytes for each packed tensor buffer.
+                   Must match the value used in NCCLWeightTransferUpdateInfo.
+            packed_num_buffers: Number of buffers for double/triple buffering.
+                   Must match the value used in NCCLWeightTransferUpdateInfo.
+
+        Example:
+            >>> from vllm.distributed.weight_transfer.nccl_engine import (
+            ...     NCCLWeightTransferEngine,
+            ... )
+            >>> param_iter = ((n, p) for n, p in model.named_parameters())
+            >>> NCCLWeightTransferEngine.trainer_send_weights(
+            ...     param_iter, group, packed=True
+            ... )
+        """
+        if post_iter_func is None:
+            # Default: extract just the tensor from (name, tensor) tuple
+            post_iter_func = lambda x: x[1]
+
+        if packed:
+            # Use packed tensor broadcasting for efficiency
+            from vllm.distributed.weight_transfer.packed_tensor import (
+                packed_broadcast_producer,
+            )
+
+            packed_broadcast_producer(
+                iterator=iterator,
+                group=group,
+                src=src,
+                post_iter_func=post_iter_func,
+                buffer_size_bytes=packed_buffer_size_bytes,
+                num_buffers=packed_num_buffers,
+            )
+        else:
+            # Use simple one-by-one broadcasting
+            for item in iterator:
+                tensor = post_iter_func(item)
+                group.broadcast(
+                    tensor, src=src, stream=stream or torch.cuda.current_stream()
+                )
+
+    @staticmethod
+    def trainer_init(
+        init_info: NCCLWeightTransferInitInfo | dict,
+    ) -> "PyNcclCommunicator":
+        """
+        Initialize NCCL process group for trainer-side weight transfer.
+
+        The trainer is always rank 0 in the process group. Uses the current
+        CUDA device (torch.cuda.current_device()).
+
+        Args:
+            init_info: Either an NCCLWeightTransferInitInfo object or a dict with keys:
+                - master_address: str
+                - master_port: int
+                - world_size: int
+
+        Returns:
+            PyNcclCommunicator for weight transfer.
+
+        Example:
+            >>> from vllm.distributed.weight_transfer.nccl_engine import (
+            ...     NCCLWeightTransferEngine,
+            ... )
+            >>> group = NCCLWeightTransferEngine.trainer_init(
+            ...     dict(
+            ...         master_address=master_address,
+            ...         master_port=master_port,
+            ...         world_size=world_size,
+            ...     ),
+            ... )
+        """
+        if isinstance(init_info, dict):
+            master_address = init_info["master_address"]
+            master_port = init_info["master_port"]
+            world_size = init_info["world_size"]
+        else:
+            # NCCLWeightTransferInitInfo object
+            master_address = init_info.master_address
+            master_port = init_info.master_port
+            world_size = init_info.world_size
+
+        # Trainer is always rank 0
+        return NCCLWeightTransferEngine._stateless_init_process_group(
+            master_address, master_port, 0, world_size, torch.cuda.current_device()
+        )
+
+    @staticmethod
+    def _stateless_init_process_group(
+        master_address, master_port, rank, world_size, device
+    ):
+        """
+        vLLM provides `StatelessProcessGroup` to create a process group
+        without considering the global process group in torch.distributed.
+        It is recommended to create `StatelessProcessGroup`, and then initialize
+        the data-plane communication (NCCL) between external (train processes)
+        and vLLM workers.
+        """
+        from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+        from vllm.distributed.utils import StatelessProcessGroup
+
+        pg = StatelessProcessGroup.create(
+            host=master_address, port=master_port, rank=rank, world_size=world_size
+        )
+        pynccl = PyNcclCommunicator(pg, device=device)
+        return pynccl

vllm/distributed/weight_transfer/packed_tensor.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Packed tensor utilities for efficient weight transfer."""
+
+import math
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import torch
+
+# Default values for packed tensor configuration.
+# These are imported by NCCLWeightTransferUpdateInfo and trainer_send_weights.
+DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024  # 1GB
+DEFAULT_PACKED_NUM_BUFFERS = 2
+
+
+def packed_broadcast_producer(
+    iterator: Iterator[tuple[str, torch.Tensor]],
+    group: Any,
+    src: int,
+    post_iter_func: Callable[[tuple[str, torch.Tensor]], torch.Tensor],
+    buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+    num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
+) -> None:
+    """Broadcast tensors in a packed manner from trainer to workers.
+
+    Args:
+        iterator: Iterator of model parameters. Returns a tuple of (name, tensor)
+        group: Process group (PyNcclCommunicator)
+        src: Source rank (0 in current implementation)
+        post_iter_func: Function to apply to each (name, tensor) pair before
+                       packing, should return a tensor
+        buffer_size_bytes: Size in bytes for each packed tensor buffer.
+                          Both producer and consumer must use the same value.
+        num_buffers: Number of buffers for double/triple buffering.
+                    Both producer and consumer must use the same value.
+
+    """
+    target_packed_tensor_size = buffer_size_bytes
+
+    streams = [torch.cuda.Stream() for _ in range(num_buffers)]
+    buffer_idx = 0
+
+    packing_tensor_list: list[list[torch.Tensor]] = [[] for _ in range(num_buffers)]
+    packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
+    packed_tensors: list[torch.Tensor] = [
+        torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers)
+    ]
+
+    while True:
+        # Synchronize the current stream
+        streams[buffer_idx].synchronize()
+        # Start tasks for the new buffer in a new stream
+        with torch.cuda.stream(streams[buffer_idx]):
+            try:
+                # Initialize the packing tensor list and sizes
+                packing_tensor_list[buffer_idx] = []
+                packing_tensor_sizes[buffer_idx] = 0
+                # Pack the tensors
+                while True:
+                    # Apply post processing and convert to linearized uint8 tensor
+                    tensor = (
+                        post_iter_func(next(iterator))
+                        .contiguous()
+                        .view(torch.uint8)
+                        .view(-1)
+                    )
+                    packing_tensor_list[buffer_idx].append(tensor)
+                    packing_tensor_sizes[buffer_idx] += tensor.numel()
+                    if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
+                        break
+                # Pack the tensors and call broadcast collective
+                packed_tensors[buffer_idx] = torch.cat(
+                    packing_tensor_list[buffer_idx], dim=0
+                )
+                group.broadcast(packed_tensors[buffer_idx], src=src)
+                # Move to the next buffer
+                buffer_idx = (buffer_idx + 1) % num_buffers
+            except StopIteration:
+                # Do the last broadcast if there are remaining tensors
+                if len(packing_tensor_list[buffer_idx]) > 0:
+                    packed_tensors[buffer_idx] = torch.cat(
+                        packing_tensor_list[buffer_idx], dim=0
+                    )
+                    group.broadcast(packed_tensors[buffer_idx], src=src)
+                break
+
+
+def packed_broadcast_consumer(
+    iterator: Iterator[tuple[str, tuple[list[int], torch.dtype]]],
+    group: Any,
+    src: int,
+    post_unpack_func: Callable[[list[tuple[str, torch.Tensor]]], None],
+    buffer_size_bytes: int = DEFAULT_PACKED_BUFFER_SIZE_BYTES,
+    num_buffers: int = DEFAULT_PACKED_NUM_BUFFERS,
+) -> None:
+    """Consume packed tensors and unpack them into a list of tensors.
+
+    Args:
+        iterator: Iterator of parameter metadata. Returns (name, (shape, dtype))
+        group: Process group (PyNcclCommunicator)
+        src: Source rank (0 in current implementation)
+        post_unpack_func: Function to apply to each list of (name, tensor) after
+                         unpacking
+        buffer_size_bytes: Size in bytes for each packed tensor buffer.
+                          Both producer and consumer must use the same value.
+        num_buffers: Number of buffers for double/triple buffering.
+                    Both producer and consumer must use the same value.
+
+    """
+
+    def unpack_tensor(
+        packed_tensor: torch.Tensor,
+        names: list[str],
+        shapes: list[list[int]],
+        dtypes: list[torch.dtype],
+        tensor_sizes: list[int],
+    ) -> list[tuple[str, torch.Tensor]]:
+        """Unpack a single tensor into a list of tensors.
+
+        Args:
+            packed_tensor: The packed torch.uint8 tensor to unpack
+            names: List of tensor names
+            shapes: List of tensor shapes
+            dtypes: List of tensor dtypes
+            tensor_sizes: List of tensor sizes in bytes
+
+        Returns:
+            unpacked List[(name, tensor)]
+        """
+        unpacked_tensors = packed_tensor.split(tensor_sizes)
+
+        unpacked_list = [
+            (name, tensor.contiguous().view(dtype).view(*shape))
+            for name, shape, dtype, tensor in zip(
+                names, shapes, dtypes, unpacked_tensors
+            )
+        ]
+
+        return unpacked_list
+
+    target_packed_tensor_size = buffer_size_bytes
+
+    streams = [torch.cuda.Stream() for _ in range(num_buffers)]
+    buffer_idx = 0
+
+    packing_tensor_meta_data: list[list[tuple[str, list[int], torch.dtype, int]]] = [
+        [] for _ in range(num_buffers)
+    ]
+    packing_tensor_sizes: list[int] = [0 for _ in range(num_buffers)]
+    packed_tensors: list[torch.Tensor] = [
+        torch.empty(0, dtype=torch.uint8, device="cuda") for _ in range(num_buffers)
+    ]
+
+    while True:
+        # Synchronize the current stream
+        streams[buffer_idx].synchronize()
+        with torch.cuda.stream(streams[buffer_idx]):
+            # Initialize the packing tensor meta data
+            packing_tensor_meta_data[buffer_idx] = []
+            packing_tensor_sizes[buffer_idx] = 0
+            try:
+                # Form a packed tensor
+                while True:
+                    name, (shape, dtype) = next(iterator)
+                    tensor_size = math.prod(shape) * dtype.itemsize
+                    packing_tensor_meta_data[buffer_idx].append(
+                        (name, shape, dtype, tensor_size)
+                    )
+                    packing_tensor_sizes[buffer_idx] += tensor_size
+                    if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size:
+                        break
+                # Create a packed tensor and broadcast it
+                packed_tensors[buffer_idx] = torch.empty(
+                    packing_tensor_sizes[buffer_idx], dtype=torch.uint8, device="cuda"
+                )
+                group.broadcast(packed_tensors[buffer_idx], src=src)
+                # Load the packed tensor into the model
+                names, shapes, dtypes, tensor_sizes = zip(
+                    *packing_tensor_meta_data[buffer_idx]
+                )
+                post_unpack_func(
+                    unpack_tensor(
+                        packed_tensors[buffer_idx],
+                        list(names),
+                        list(shapes),
+                        list(dtypes),
+                        list(tensor_sizes),
+                    )
+                )
+                # Move to the next buffer
+                buffer_idx = (buffer_idx + 1) % num_buffers
+            except StopIteration:
+                # Do the last broadcast if there are remaining tensors
+                if len(packing_tensor_meta_data[buffer_idx]) > 0:
+                    # Create a packed tensor and broadcast it
+                    packed_tensors[buffer_idx] = torch.empty(
+                        packing_tensor_sizes[buffer_idx],
+                        dtype=torch.uint8,
+                        device="cuda",
+                    )
+                    group.broadcast(packed_tensors[buffer_idx], src=src)
+                    # Load the packed tensor into the model
+                    names, shapes, dtypes, tensor_sizes = zip(
+                        *packing_tensor_meta_data[buffer_idx]
+                    )
+                    post_unpack_func(
+                        unpack_tensor(
+                            packed_tensors[buffer_idx],
+                            list(names),
+                            list(shapes),
+                            list(dtypes),
+                            list(tensor_sizes),
+                        )
+                    )
+                break

vllm/engine/arg_utils.py

@@ -54,6 +54,7 @@ from vllm.config import (
     SpeculativeConfig,
     StructuredOutputsConfig,
     VllmConfig,
+    WeightTransferConfig,
     get_attr_docs,
 )
 from vllm.config.cache import (
@@ -581,6 +582,11 @@ class EngineArgs:
     kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend
     tokens_only: bool = False
 
+    weight_transfer_config: WeightTransferConfig | None = None
+    """Configuration for weight transfer during RL training. 
+    Accepts a JSON string or dict with backend-specific options.
+    Example: '{"backend": "nccl"}'"""
+
     def __post_init__(self):
         # support `EngineArgs(compilation_config={...})`
         # without having to manually construct a
@@ -591,6 +597,10 @@ class EngineArgs:
             self.attention_config = AttentionConfig(**self.attention_config)
         if isinstance(self.eplb_config, dict):
             self.eplb_config = EPLBConfig(**self.eplb_config)
+        if isinstance(self.weight_transfer_config, dict):
+            self.weight_transfer_config = WeightTransferConfig(
+                **self.weight_transfer_config
+            )
         # Setup plugins
         from vllm.plugins import load_general_plugins
 
@@ -1189,6 +1199,9 @@ class EngineArgs:
         vllm_group.add_argument(
             "--optimization-level", **vllm_kwargs["optimization_level"]
         )
+        vllm_group.add_argument(
+            "--weight-transfer-config", **vllm_kwargs["weight_transfer_config"]
+        )
 
         # Other arguments
         parser.add_argument(
@@ -1765,6 +1778,7 @@ class EngineArgs:
             profiler_config=self.profiler_config,
             additional_config=self.additional_config,
             optimization_level=self.optimization_level,
+            weight_transfer_config=self.weight_transfer_config,
         )
 
         return config

vllm/engine/protocol.py

@@ -6,6 +6,10 @@ from collections.abc import AsyncGenerator, Iterable, Mapping
 from typing import Any
 
 from vllm.config import ModelConfig, VllmConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
 from vllm.inputs.data import PromptType, StreamingInput
 from vllm.lora.request import LoRARequest
 from vllm.outputs import PoolingRequestOutput, RequestOutput
@@ -191,3 +195,13 @@ class EngineClient(ABC):
     async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         """Get supported tasks"""
         raise NotImplementedError
+
+    async def init_weight_transfer_engine(
+        self, init_request: WeightTransferInitRequest
+    ) -> None:
+        """Initialize weight transfer for RL training."""
+        raise NotImplementedError
+
+    async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
+        """Batched weight update for RL training."""
+        raise NotImplementedError

vllm/entrypoints/llm.py

@@ -34,6 +34,10 @@ from vllm.config.model import (
     RunnerOption,
     TokenizerMode,
 )
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
 from vllm.engine.arg_utils import EngineArgs
 from vllm.entrypoints.chat_utils import (
     ChatCompletionMessageParam,
@@ -360,6 +364,23 @@ class LLM:
     def get_tokenizer(self) -> TokenizerLike:
         return self.llm_engine.get_tokenizer()
 
+    def get_world_size(self, include_dp: bool = True) -> int:
+        """Get the world size from the parallel config.
+
+        Args:
+            include_dp: If True (default), returns the world size including
+                data parallelism (TP * PP * DP). If False, returns the world
+                size without data parallelism (TP * PP).
+
+        Returns:
+            The world size (tensor_parallel_size * pipeline_parallel_size),
+            optionally multiplied by data_parallel_size if include_dp is True.
+        """
+        parallel_config = self.llm_engine.vllm_config.parallel_config
+        if include_dp:
+            return parallel_config.world_size_across_dp
+        return parallel_config.world_size
+
     def reset_mm_cache(self) -> None:
         self.input_processor.clear_mm_cache()
         self.llm_engine.reset_mm_cache()
@@ -1135,11 +1156,12 @@ class LLM:
             # Use default pooling params.
             pooling_params = PoolingParams()
 
-        if pooling_task not in self.supported_tasks:
-            raise ValueError(f"pooling_task must be one of {self.supported_tasks}.")
-
         for param in as_iter(pooling_params):
-            param.verify(pooling_task, model_config)
+            if param.task is None:
+                param.task = pooling_task
+            elif param.task != pooling_task:
+                msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
+                raise ValueError(msg)
 
         self._validate_and_add_requests(
             prompts=prompts,
@@ -1472,8 +1494,9 @@ class LLM:
 
         if pooling_params is None:
             pooling_params = PoolingParams(task="score")
+        elif pooling_params.task is None:
+            pooling_params.task = "score"
 
-        pooling_params.verify("score", model_config)
         pooling_params_list = list[PoolingParams]()
 
         prompts = list[PromptType]()
@@ -1836,6 +1859,7 @@ class LLM:
             lora_request=lora_request,
             tokenization_kwargs=tokenization_kwargs,
             priority=priority,
+            supported_tasks=self.supported_tasks,
         )
 
         self.llm_engine.add_request(
@@ -1900,6 +1924,38 @@ class LLM:
         # its previous requests.
         return sorted(outputs, key=lambda x: int(x.request_id))
 
+    def init_weight_transfer_engine(
+        self, request: WeightTransferInitRequest | dict
+    ) -> None:
+        """
+        Initialize weight transfer for RL training.
+
+        Args:
+            request: Weight transfer initialization request with backend-specific info
+        """
+        init_info_dict = (
+            request["init_info"] if isinstance(request, dict) else request.init_info
+        )
+
+        self.llm_engine.collective_rpc(
+            "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
+        )
+
+    def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None:
+        """
+        Update the weights of the model.
+
+        Args:
+            request: Weight update request with backend-specific update info
+        """
+        update_info_dict = (
+            request["update_info"] if isinstance(request, dict) else request.update_info
+        )
+
+        self.llm_engine.collective_rpc(
+            "update_weights", kwargs={"update_info": update_info_dict}
+        )
+
     def __repr__(self) -> str:
         """Return a transformers-style hierarchical view of the model."""
         # Cache the result to avoid repeated collective_rpc calls

vllm/entrypoints/openai/api_server.py

@@ -190,11 +190,11 @@ def build_app(args: Namespace, supported_tasks: tuple["SupportedTask", ...]) ->
         register_generate_api_routers(app)
 
     if "transcription" in supported_tasks:
-        from vllm.entrypoints.openai.translations.api_router import (
-            attach_router as register_translations_api_router,
+        from vllm.entrypoints.openai.speech_to_text.api_router import (
+            attach_router as register_speech_to_text_api_router,
         )
 
-        register_translations_api_router(app)
+        register_speech_to_text_api_router(app)
 
     if "realtime" in supported_tasks:
         from vllm.entrypoints.openai.realtime.api_router import (
@@ -318,7 +318,7 @@ async def init_app_state(
         )
 
     if "transcription" in supported_tasks:
-        from vllm.entrypoints.openai.translations.api_router import (
+        from vllm.entrypoints.openai.speech_to_text.api_router import (
             init_transcription_state,
         )
 

vllm/entrypoints/openai/chat_completion/serving.py

@@ -72,6 +72,7 @@ from vllm.logger import init_logger
 from vllm.logprobs import Logprob
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.parser import ParserManager
+from vllm.reasoning import ReasoningParser
 from vllm.sampling_params import BeamSearchParams, SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.tokenizers.mistral import (
@@ -132,7 +133,7 @@ class OpenAIServingChat(OpenAIServing):
         self.logits_processors = self.model_config.logits_processors
 
         # set up reasoning parser
-        self.reasoning_parser = ParserManager.get_reasoning_parser(
+        self.reasoning_parser_cls = ParserManager.get_reasoning_parser(
             reasoning_parser_name=reasoning_parser
         )
         # set up tool use
@@ -330,6 +331,24 @@ class OpenAIServingChat(OpenAIServing):
         for the API specification. This API mimics the OpenAI
         Chat Completion API.
         """
+        # Streaming response
+        tokenizer = self.renderer.tokenizer
+        assert tokenizer is not None
+        reasoning_parser: ReasoningParser | None = None
+        try:
+            if self.reasoning_parser_cls:
+                # Pass the same chat template kwargs as used in tokenization
+                chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
+                    request.chat_template_kwargs,
+                    self.default_chat_template_kwargs,
+                )
+                reasoning_parser = self.reasoning_parser_cls(
+                    tokenizer,
+                    chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
+                )
+        except RuntimeError as e:
+            logger.exception("Error in reasoning parser creation.")
+            return self.create_error_response(str(e))
         result = await self.render_chat_request(request)
         if isinstance(result, ErrorResponse):
             return result
@@ -427,7 +446,12 @@ class OpenAIServingChat(OpenAIServing):
                         priority=request.priority,
                         data_parallel_rank=data_parallel_rank,
                     )
-
+                    reasoning_ended = None
+                    if reasoning_parser:
+                        reasoning_ended = reasoning_parser.is_reasoning_end(
+                            engine_request.prompt_token_ids or []  # type: ignore[attr-defined]
+                        )
+                        engine_request.reasoning_ended = reasoning_ended
                     generator = self.engine_client.generate(
                         engine_request,
                         sampling_params,
@@ -447,10 +471,6 @@ class OpenAIServingChat(OpenAIServing):
         assert len(generators) == 1
         (result_generator,) = generators
 
-        # Streaming response
-        tokenizer = self.renderer.tokenizer
-        assert tokenizer is not None
-
         if request.stream:
             return self.chat_completion_stream_generator(
                 request,
@@ -460,6 +480,7 @@ class OpenAIServingChat(OpenAIServing):
                 conversation,
                 tokenizer,
                 request_metadata,
+                reasoning_parser,
             )
 
         try:
@@ -471,6 +492,7 @@ class OpenAIServingChat(OpenAIServing):
                 conversation,
                 tokenizer,
                 request_metadata,
+                reasoning_parser,
             )
         except GenerationError as e:
             return self._convert_generation_error_to_response(e)
@@ -630,6 +652,7 @@ class OpenAIServingChat(OpenAIServing):
         conversation: list[ConversationMessage],
         tokenizer: TokenizerLike,
         request_metadata: RequestResponseMetadata,
+        reasoning_parser: ReasoningParser | None = None,
     ) -> AsyncGenerator[str, None]:
         from vllm.tokenizers.mistral import MistralTokenizer
 
@@ -673,7 +696,7 @@ class OpenAIServingChat(OpenAIServing):
 
         # Only one of these will be used, thus previous_texts and
         # all_previous_token_ids will not be used twice in the same iteration.
-        if tool_choice_auto or self.reasoning_parser:
+        if tool_choice_auto or reasoning_parser:
             # These are only required in "auto" tool choice case
             all_previous_token_ids = [[]] * num_choices
             # For reasoning parser and tool call all enabled
@@ -683,28 +706,6 @@ class OpenAIServingChat(OpenAIServing):
         else:
             all_previous_token_ids = None
 
-        try:
-            if self.reasoning_parser:
-                if tokenizer is None:
-                    raise ValueError(
-                        "Tokenizer not available when `skip_tokenizer_init=True`"
-                    )
-
-                # Pass the same chat template kwargs as used in tokenization
-                chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
-                    request.chat_template_kwargs,
-                    self.default_chat_template_kwargs,
-                )
-                reasoning_parser = self.reasoning_parser(
-                    tokenizer,
-                    chat_template_kwargs=chat_template_kwargs or {},  # type: ignore[call-arg]
-                )
-        except RuntimeError as e:
-            logger.exception("Error in reasoning parser creation.")
-            data = self.create_streaming_error_response(str(e))
-            yield f"data: {data}\n\n"
-            yield "data: [DONE]\n\n"
-            return
         # Prepare the tool parser if it's needed
         try:
             if tool_choice_auto and self.tool_parser:
@@ -826,7 +827,7 @@ class OpenAIServingChat(OpenAIServing):
                     tool_parser = tool_parsers[i]
 
                     if (
-                        self.reasoning_parser
+                        reasoning_parser
                         and res.prompt_token_ids
                         and prompt_is_reasoning_end_arr[i] is None
                     ):
@@ -888,7 +889,7 @@ class OpenAIServingChat(OpenAIServing):
                     delta_message: DeltaMessage | None
 
                     # just update previous_texts and previous_token_ids
-                    if tool_choice_auto or self.reasoning_parser:
+                    if tool_choice_auto or reasoning_parser:
                         assert previous_texts is not None
                         assert all_previous_token_ids is not None
                         previous_text = previous_texts[i]
@@ -915,7 +916,7 @@ class OpenAIServingChat(OpenAIServing):
                     # handle streaming deltas for tools with named tool_choice
                     elif tool_choice_function_name:
                         if (
-                            self.reasoning_parser
+                            reasoning_parser
                             and not reasoning_end_arr[i]
                             and not reasoning_parser.is_reasoning_end(
                                 previous_token_ids
@@ -952,7 +953,7 @@ class OpenAIServingChat(OpenAIServing):
                                     current_text = ""
                         else:
                             # Just to add remaining `content`
-                            if self.reasoning_parser:
+                            if reasoning_parser:
                                 delta_text = previous_text + delta_text
                                 current_text = ""
 
@@ -998,13 +999,13 @@ class OpenAIServingChat(OpenAIServing):
                         output_token_ids = as_list(output.token_ids)
 
                         if (
-                            self.reasoning_parser is not None
+                            reasoning_parser is not None
                             and not reasoning_end_arr[i]
                             and prompt_is_reasoning_end_arr[i]
                         ):
                             reasoning_end_arr[i] = True
 
-                        if self.reasoning_parser and not reasoning_end_arr[i]:
+                        if reasoning_parser and not reasoning_end_arr[i]:
                             delta_message = (
                                 reasoning_parser.extract_reasoning_streaming(
                                     previous_text,
@@ -1047,9 +1048,8 @@ class OpenAIServingChat(OpenAIServing):
 
                     # handle streaming deltas for tools with "auto" tool choice
                     # and reasoning parser
-                    elif tool_choice_auto and self.reasoning_parser:
+                    elif tool_choice_auto and reasoning_parser:
                         assert tool_parser is not None
-                        assert reasoning_parser is not None
                         assert added_content_delta_arr is not None
                         assert reasoning_end_arr is not None
                         output_token_ids = as_list(output.token_ids)
@@ -1130,7 +1130,7 @@ class OpenAIServingChat(OpenAIServing):
                             tools_streamed[i] = True
 
                     # when only reasoning
-                    elif self.reasoning_parser:
+                    elif reasoning_parser:
                         delta_message = reasoning_parser.extract_reasoning_streaming(
                             previous_text,
                             current_text,
@@ -1144,9 +1144,7 @@ class OpenAIServingChat(OpenAIServing):
                         delta_message = DeltaMessage(content=delta_text)
 
                     # update the previous values for the next iteration
-                    if (
-                        tool_choice_auto or self.reasoning_parser
-                    ) and not self.use_harmony:
+                    if (tool_choice_auto or reasoning_parser) and not self.use_harmony:
                         assert previous_texts is not None
                         assert all_previous_token_ids is not None
                         previous_texts[i] = current_text
@@ -1400,6 +1398,7 @@ class OpenAIServingChat(OpenAIServing):
         conversation: list[ConversationMessage],
         tokenizer: TokenizerLike,
         request_metadata: RequestResponseMetadata,
+        reasoning_parser: ReasoningParser | None = None,
     ) -> ErrorResponse | ChatCompletionResponse:
         from vllm.tokenizers.mistral import MistralTokenizer
 
@@ -1494,25 +1493,7 @@ class OpenAIServingChat(OpenAIServing):
                 choices.append(choice_data)
                 continue
 
-            if self.reasoning_parser:
-                try:
-                    if tokenizer is None:
-                        raise ValueError(
-                            "Tokenizer not available when `skip_tokenizer_init=True`"
-                        )
-
-                    # Pass the same chat template kwargs as used in tokenization
-                    chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
-                        request.chat_template_kwargs,
-                        self.default_chat_template_kwargs,
-                    )
-                    reasoning_parser = self.reasoning_parser(
-                        tokenizer,
-                        chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
-                    )
-                except RuntimeError as e:
-                    logger.exception("Error in reasoning parser creation.")
-                    return self.create_error_response(str(e))
+            if reasoning_parser:
                 # If the reasoning parser is enabled,
                 # tool calls are extracted exclusively from the content.
                 reasoning, content = reasoning_parser.extract_reasoning(

vllm/entrypoints/openai/engine/serving.py

@@ -57,7 +57,7 @@ from vllm.entrypoints.openai.responses.protocol import (
 from vllm.entrypoints.openai.responses.utils import (
     construct_input_messages,
 )
-from vllm.entrypoints.openai.translations.protocol import (
+from vllm.entrypoints.openai.speech_to_text.protocol import (
     TranscriptionRequest,
     TranscriptionResponse,
     TranslationRequest,

vllm/entrypoints/openai/parser/harmony_utils.py

@@ -68,6 +68,14 @@ MCP_BUILTIN_TOOLS: set[str] = {
     "container",
 }
 
+# Mapping from built-in tool recipient names to their MCP server labels.
+# This ensures consistency between streaming and non-streaming responses.
+_BUILTIN_TOOL_TO_MCP_SERVER_LABEL: dict[str, str] = {
+    "python": "code_interpreter",
+    "browser": "web_search_preview",
+    "container": "container",
+}
+
 
 def has_custom_tools(tool_types: set[str]) -> bool:
     """
@@ -601,7 +609,13 @@ def _parse_mcp_recipient(recipient: str) -> tuple[str, str]:
 
 def _parse_mcp_call(message: Message, recipient: str) -> list[ResponseOutputItem]:
     """Parse MCP calls into MCP call items."""
-    server_label, tool_name = _parse_mcp_recipient(recipient)
+    # Handle built-in tools that need server_label mapping
+    if recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL:
+        server_label = _BUILTIN_TOOL_TO_MCP_SERVER_LABEL[recipient]
+        tool_name = recipient
+    else:
+        server_label, tool_name = _parse_mcp_recipient(recipient)
+
     output_items = []
     for content in message.content:
         response_item = McpCall(
@@ -630,7 +644,7 @@ def parse_output_message(message: Message) -> list[ResponseOutputItem]:
     recipient = message.recipient
 
     if recipient is not None:
-        # Browser tool calls
+        # Browser tool calls (browser.search, browser.open, browser.find)
         if recipient.startswith("browser."):
             output_items.append(_parse_browser_tool_call(message, recipient))
 
@@ -638,10 +652,8 @@ def parse_output_message(message: Message) -> list[ResponseOutputItem]:
         elif message.channel == "commentary" and recipient.startswith("functions."):
             output_items.extend(_parse_function_call(message, recipient))
 
-        # Built-in tools are treated as reasoning
-        elif recipient.startswith(("python", "browser", "container")):
-            # Built-in tool recipients (python/browser/container)
-            # generate reasoning output
+        # Built-in MCP tools (python, browser, container)
+        elif recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL:
             output_items.extend(_parse_reasoning(message))
 
         # All other recipients are MCP calls
@@ -688,13 +700,23 @@ def parse_remaining_state(parser: StreamableParser) -> list[ResponseOutputItem]:
                     status="in_progress",
                 )
             ]
-        # Built-in tools (python, browser, container) should be treated as reasoning
-        elif not (
-            current_recipient.startswith("python")
-            or current_recipient.startswith("browser")
-            or current_recipient.startswith("container")
-        ):
-            # All other recipients are MCP calls
+        # Built-in MCP tools (python, browser, container)
+        elif current_recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL:
+            return [
+                ResponseReasoningItem(
+                    id=f"rs_{random_uuid()}",
+                    summary=[],
+                    type="reasoning",
+                    content=[
+                        ResponseReasoningTextContent(
+                            text=parser.current_content, type="reasoning_text"
+                        )
+                    ],
+                    status=None,
+                )
+            ]
+        # All other recipients are MCP calls
+        else:
             rid = random_uuid()
             server_label, tool_name = _parse_mcp_recipient(current_recipient)
             return [

vllm/entrypoints/openai/responses/serving.py

@@ -63,7 +63,6 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.chat_utils import (
     ChatCompletionMessageParam,
     ChatTemplateContentFormatOption,
-    make_tool_call_id,
 )
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import ToolServer
@@ -915,114 +914,57 @@ class OpenAIServingResponses(OpenAIServing):
         final_output: CompletionOutput,
         tokenizer: TokenizerLike,
     ) -> list[ResponseOutputItem]:
-        if self.parser and self.parser.reasoning_parser_cls:
-            try:
-                reasoning_parser = self.parser.reasoning_parser_cls(tokenizer)
-            except RuntimeError as e:
-                logger.exception("Error in reasoning parser creation.")
-                raise e
-
-            reasoning, content = reasoning_parser.extract_reasoning(
-                final_output.text, request=request
-            )
-        else:
-            reasoning = None
-            content = final_output.text
-
         # Log complete response if output logging is enabled
         if self.enable_log_outputs and self.request_logger:
-            output_text = ""
-            if content:
-                output_text = content
-            elif reasoning:
-                output_text = f"[reasoning: {reasoning}]"
-
-            if output_text:
-                self.request_logger.log_outputs(
-                    request_id=request.request_id,
-                    outputs=output_text,
-                    output_token_ids=final_output.token_ids,
-                    finish_reason=final_output.finish_reason,
-                    is_streaming=False,
-                    delta=False,
-                )
-
-        reasoning_item = None
-        message_item = None
-        if reasoning:
-            reasoning_item = ResponseReasoningItem(
-                id=f"rs_{random_uuid()}",
-                summary=[],
-                type="reasoning",
-                content=[
-                    ResponseReasoningTextContent(text=reasoning, type="reasoning_text")
-                ],
-                status=None,  # NOTE: Only the last output item has status.
+            self.request_logger.log_outputs(
+                request_id=request.request_id,
+                outputs=final_output.text,
+                output_token_ids=final_output.token_ids,
+                finish_reason=final_output.finish_reason,
+                is_streaming=False,
+                delta=False,
             )
-        tool_calls, content = self._parse_tool_calls_from_content(
-            request=request,
-            tokenizer=tokenizer,
-            content=content,
-            enable_auto_tools=self.enable_auto_tools,
-            tool_parser_cls=self.parser.tool_parser_cls if self.parser else None,
-        )
 
-        if content or (self.use_harmony and tool_calls):
-            res_text_part = None
-            if content:
-                res_text_part = ResponseOutputText(
-                    text=content,
-                    annotations=[],  # TODO
-                    type="output_text",
-                    logprobs=(
-                        self._create_response_logprobs(
-                            token_ids=final_output.token_ids,
-                            logprobs=final_output.logprobs,
-                            tokenizer=tokenizer,
-                            top_logprobs=request.top_logprobs,
-                        )
-                        if request.is_include_output_logprobs()
-                        else None
-                    ),
-                )
-            message_item = ResponseOutputMessage(
+        # Compute logprobs if requested
+        logprobs = None
+        if request.is_include_output_logprobs() and final_output.logprobs:
+            logprobs = self._create_response_logprobs(
+                token_ids=final_output.token_ids,
+                logprobs=final_output.logprobs,
+                tokenizer=tokenizer,
+                top_logprobs=request.top_logprobs,
+            )
+
+        # Use parser to extract and create response output items
+        if self.parser:
+            parser = self.parser(tokenizer)
+            return parser.extract_response_outputs(
+                model_output=final_output.text,
+                request=request,
+                enable_auto_tools=self.enable_auto_tools,
+                tool_call_id_type=self.tool_call_id_type,
+                logprobs=logprobs,
+            )
+
+        # Fallback when no parser is configured
+        return [
+            ResponseOutputMessage(
                 id=f"msg_{random_uuid()}",
-                content=[res_text_part] if res_text_part else [],
+                content=[
+                    ResponseOutputText(
+                        text=final_output.text,
+                        annotations=[],
+                        type="output_text",
+                        logprobs=logprobs,
+                    )
+                ]
+                if final_output.text
+                else [],
                 role="assistant",
                 status="completed",
                 type="message",
             )
-        outputs = []
-
-        if reasoning_item:
-            outputs.append(reasoning_item)
-        if message_item:
-            outputs.append(message_item)
-        if tool_calls:
-            # We use a simple counter for history_tool_call_count because
-            # we don't track the history of tool calls in the Responses API yet.
-            # This means that the tool call index will start from 0 for each
-            # request.
-            tool_call_items = []
-            for history_tool_call_cnt, tool_call in enumerate(tool_calls):
-                tool_call_items.append(
-                    ResponseFunctionToolCall(
-                        id=f"fc_{random_uuid()}",
-                        call_id=tool_call.id
-                        if tool_call.id
-                        else make_tool_call_id(
-                            id_type=self.tool_call_id_type,
-                            func_name=tool_call.name,
-                            idx=history_tool_call_cnt,
-                        ),
-                        type="function_call",
-                        status="completed",
-                        name=tool_call.name,
-                        arguments=tool_call.arguments,
-                    )
-                )
-            outputs.extend(tool_call_items)
-        return outputs
+        ]
 
     def _make_response_output_items_with_harmony(
         self,

vllm/entrypoints/openai/speech_to_text/api_router.py

@@ -9,13 +9,13 @@ from fastapi import APIRouter, FastAPI, Form, Request
 from fastapi.responses import JSONResponse, StreamingResponse
 
 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
-from vllm.entrypoints.openai.translations.protocol import (
+from vllm.entrypoints.openai.speech_to_text.protocol import (
     TranscriptionRequest,
     TranscriptionResponseVariant,
     TranslationRequest,
     TranslationResponseVariant,
 )
-from vllm.entrypoints.openai.translations.serving import (
+from vllm.entrypoints.openai.speech_to_text.serving import (
     OpenAIServingTranscription,
     OpenAIServingTranslation,
 )

vllm/entrypoints/openai/speech_to_text/serving.py

@@ -11,7 +11,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     RequestResponseMetadata,
 )
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.openai.translations.protocol import (
+from vllm.entrypoints.openai.speech_to_text.protocol import (
     TranscriptionRequest,
     TranscriptionResponse,
     TranscriptionResponseStreamChoice,
@@ -23,7 +23,7 @@ from vllm.entrypoints.openai.translations.protocol import (
     TranslationResponseVerbose,
     TranslationStreamResponse,
 )
-from vllm.entrypoints.openai.translations.speech_to_text import OpenAISpeechToText
+from vllm.entrypoints.openai.speech_to_text.speech_to_text import OpenAISpeechToText
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
 

vllm/entrypoints/openai/speech_to_text/speech_to_text.py

@@ -24,7 +24,7 @@ from vllm.entrypoints.openai.engine.protocol import (
 )
 from vllm.entrypoints.openai.engine.serving import OpenAIServing, SpeechToTextRequest
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.openai.translations.protocol import (
+from vllm.entrypoints.openai.speech_to_text.protocol import (
     TranscriptionResponse,
     TranscriptionResponseStreamChoice,
     TranscriptionResponseVerbose,
@@ -402,7 +402,7 @@ class OpenAISpeechToText(OpenAIServing):
         audio_data: bytes,
         request: SpeechToTextRequest,
         raw_request: Request,
-        response_class: type[T | V],
+        response_class: type[ResponseType],
         stream_generator_method: Callable[..., AsyncGenerator[str, None]],
     ) -> T | V | AsyncGenerator[str, None] | ErrorResponse:
         """Base method for speech-to-text operations like transcription and