Merge ROCm:fused_aiter_rope_kvcache (#33443)

Signed-off-by: Hashem Hashemi <hashem.hashemi@amd.com>

.buildkite/image_build/image_build.sh

@@ -8,7 +8,7 @@ clean_docker_tag() {
 }
 
 print_usage_and_exit() {
-    echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag> [<image_tag_latest>]"
+    echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
     exit 1
 }
 
@@ -154,7 +154,7 @@ print_bake_config() {
 #################################
 print_instance_info
 
-if [[ $# -lt 5 ]]; then
+if [[ $# -lt 7 ]]; then
     print_usage_and_exit
 fi
 
@@ -163,8 +163,10 @@ REGISTRY=$1
 REPO=$2
 BUILDKITE_COMMIT=$3
 BRANCH=$4
-IMAGE_TAG=$5
-IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional
+VLLM_USE_PRECOMPILED=$5
+VLLM_MERGE_BASE_COMMIT=$6
+IMAGE_TAG=$7
+IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
 
 # build config
 TARGET="test-ci"
@@ -191,6 +193,8 @@ export CACHE_FROM
 export CACHE_FROM_BASE_BRANCH
 export CACHE_FROM_MAIN
 export CACHE_TO
+export VLLM_USE_PRECOMPILED
+export VLLM_MERGE_BASE_COMMIT
 
 # print args
 echo "--- :mag: Arguments"
@@ -198,6 +202,8 @@ echo "REGISTRY: ${REGISTRY}"
 echo "REPO: ${REPO}"
 echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
 echo "BRANCH: ${BRANCH}"
+echo "VLLM_USE_PRECOMPILED: ${VLLM_USE_PRECOMPILED}"
+echo "VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
 echo "IMAGE_TAG: ${IMAGE_TAG}"
 echo "IMAGE_TAG_LATEST: ${IMAGE_TAG_LATEST}"
 

.buildkite/image_build/image_build.yaml

@@ -3,9 +3,10 @@ steps:
   - label: ":docker: Build image"
     key: image-build
     depends_on: []
+    timeout_in_minutes: 600
     commands:
-    - if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
-    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; fi
+    - if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
+    - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG $IMAGE_TAG_LATEST; fi
     retry:
       automatic:
         - exit_status: -1  # Agent was lost
@@ -41,7 +42,7 @@ steps:
           limit: 2
         - exit_status: -10  # Agent was lost
           limit: 2
-  
+
   - label: ":docker: Build CPU arm64 image"
     key: cpu-arm64-image-build
     depends_on: []

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

@@ -39,6 +39,7 @@ docker run \
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
     python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
+    python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
     python3 examples/offline_inference/basic/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4  --block-size 64 --enforce-eager
     python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2
     python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b  --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel

.buildkite/test-amd.yaml

@@ -514,7 +514,7 @@ steps:
     - python3 offline_inference/vision_language_multi_image.py --seed 0
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # for pooling models
-    - python3 pooling/pooling/vision_language_pooling.py --seed 0
+    - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # for features demo
     - python3 offline_inference/prefix_caching.py
     - python3 offline_inference/llm_engine_example.py

.buildkite/test-pipeline.yaml

@@ -453,7 +453,7 @@ steps:
     - python3 offline_inference/vision_language_multi_image.py --seed 0
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
     # for pooling models
-    - python3 pooling/pooling/vision_language_pooling.py --seed 0
+    - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # for features demo
     - python3 offline_inference/prefix_caching.py
     - python3 offline_inference/llm_engine_example.py

.buildkite/test_areas/misc.yaml

@@ -72,7 +72,7 @@ steps:
     - python3 offline_inference/vision_language_multi_image.py --seed 0
     - python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
      # for pooling models
-    - python3 pooling/pooling/vision_language_pooling.py --seed 0
+    - python3 pooling/embed/vision_embedding_offline.py --seed 0
     # for features demo
     - python3 offline_inference/prefix_caching.py
     - python3 offline_inference/llm_engine_example.py

CMakeLists.txt

@@ -56,8 +56,8 @@ endif()
 # requirements.txt files and should be kept consistent.  The ROCm torch
 # versions are derived from docker/Dockerfile.rocm
 #
-set(TORCH_SUPPORTED_VERSION_CUDA "2.9.1")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.9.1")
+set(TORCH_SUPPORTED_VERSION_CUDA "2.10.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.10.0")
 
 #
 # Try to find python package with an executable that exactly matches
@@ -433,7 +433,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_BF16_KERNEL_SRC})
     endif()
 
-    if (MARLIN_SM75_ARCHS) 
+    if (MARLIN_SM75_ARCHS)
       file(GLOB MARLIN_TEMPLATE_SM75_KERNEL_SRC "csrc/quantization/marlin/sm75_kernel_*.cu")
       set_gencode_flags_for_srcs(
         SRCS "${MARLIN_TEMPLATE_SM75_KERNEL_SRC}"
@@ -445,7 +445,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_SM75_KERNEL_SRC})
     endif()
 
-    if (MARLIN_FP8_ARCHS) 
+    if (MARLIN_FP8_ARCHS)
       file(GLOB MARLIN_TEMPLATE_FP8_KERNEL_SRC "csrc/quantization/marlin/sm89_kernel_*.cu")
       set_gencode_flags_for_srcs(
         SRCS "${MARLIN_TEMPLATE_FP8_KERNEL_SRC}"
@@ -1042,7 +1042,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       list(APPEND VLLM_MOE_EXT_SRC ${MARLIN_MOE_SRC})
     endif()
 
-    if (MARLIN_MOE_SM75_ARCHS) 
+    if (MARLIN_MOE_SM75_ARCHS)
       file(GLOB MARLIN_MOE_SM75_SRC "csrc/moe/marlin_moe_wna16/sm75_kernel_*.cu")
       set_gencode_flags_for_srcs(
         SRCS "${MARLIN_MOE_SM75_SRC}"

cmake/external_projects/triton_kernels.cmake

@@ -1,9 +1,9 @@
 # Install OpenAI triton_kernels from https://github.com/triton-lang/triton/tree/main/python/triton_kernels
 
-set(DEFAULT_TRITON_KERNELS_TAG "v3.5.0")
+set(DEFAULT_TRITON_KERNELS_TAG "v3.6.0")
 
 # Set TRITON_KERNELS_SRC_DIR for use with local development with vLLM. We expect TRITON_KERNELS_SRC_DIR to
-# be directly set to the triton_kernels python directory. 
+# be directly set to the triton_kernels python directory.
 if (DEFINED ENV{TRITON_KERNELS_SRC_DIR})
   message(STATUS "[triton_kernels] Fetch from $ENV{TRITON_KERNELS_SRC_DIR}")
   FetchContent_Declare(
@@ -24,7 +24,7 @@ else()
   )
 endif()
 
-# Fetch content 
+# Fetch content
 FetchContent_MakeAvailable(triton_kernels)
 
 if (NOT triton_kernels_SOURCE_DIR)
@@ -47,7 +47,7 @@ install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tr
 ## Copy .py files to install directory.
 install(DIRECTORY
         ${TRITON_KERNELS_PYTHON_DIR}
-        DESTINATION 
+        DESTINATION
         vllm/third_party/triton_kernels/
         COMPONENT triton_kernels
         FILES_MATCHING PATTERN "*.py")

csrc/rocm/skinny_gemms.cu

@@ -1568,8 +1568,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
         {
   #endif
           unsigned int kOff = k + (thrd * A_CHUNK);
-          unsigned int kOffcp =
-              k_str + kOff;  // min__(K - A_CHUNK, k_str + kOff);
+          unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
           for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
             __builtin_amdgcn_global_load_lds(
                 (int*)(&A[min__(

docker/Dockerfile.cpu

@@ -134,7 +134,6 @@ WORKDIR /vllm-workspace
 # Copy test requirements
 COPY requirements/test.in requirements/cpu-test.in
 
-# TODO: Update to 2.9.0 when there is a new build for intel_extension_for_pytorch for that version
 RUN \
     sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
     remove_packages_not_supported_on_aarch64() { \

docker/Dockerfile.rocm

@@ -1,7 +1,7 @@
 # default base image
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=rocm/vllm-dev:base
+ARG BASE_IMAGE=rocm/vllm-dev:base_custom_releases_rocm_v0.16.0_20260211
 
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE

docker/Dockerfile.rocm_base

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-complete
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2-complete
 ARG TRITON_BRANCH="57c693b6"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG PYTORCH_BRANCH="89075173"
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
 ARG FA_BRANCH="0e60e394"
 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
-ARG AITER_BRANCH="6af8b687"
+ARG AITER_BRANCH="v0.1.10.post2"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="2d02c6a9"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -142,6 +142,7 @@ ARG PYTORCH_VISION_REPO
 ARG PYTORCH_AUDIO_REPO
 ARG USE_SCCACHE
 
+RUN apt-get update && apt-get install -y pkg-config liblzma-dev
 RUN git clone ${PYTORCH_REPO} pytorch
 RUN cd pytorch && git checkout ${PYTORCH_BRANCH} \
     && pip install -r requirements.txt && git submodule update --init --recursive \
@@ -239,7 +240,7 @@ RUN pip install pyyaml && cd aiter \
            export HIP_CLANG_PATH=/opt/sccache-wrappers \
            && sccache --show-stats; \
        fi \
-    && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
     && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
     && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install

docs/features/multimodal_inputs.md

@@ -510,7 +510,7 @@ Our OpenAI-compatible server accepts multi-modal data via the [Chat Completions
     If no fallback is available, an error is raised and you have to provide the chat template manually via the `--chat-template` argument.
 
     For certain models, we provide alternative chat templates inside [examples](../../examples).
-    For example, VLM2Vec uses [examples/template_vlm2vec_phi3v.jinja](../../examples/template_vlm2vec_phi3v.jinja) which is different from the default one for Phi-3-Vision.
+    For example, VLM2Vec uses [examples/pooling/embed/template/vlm2vec_phi3v.jinja](../../examples/pooling/embed/template/vlm2vec_phi3v.jinja) which is different from the default one for Phi-3-Vision.
 
 ### Image Inputs
 

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

@@ -6,10 +6,11 @@ vLLM initially supports basic model inference and serving on Intel GPU platform.
 # --8<-- [start:requirements]
 
 - Supported Hardware: Intel Data Center GPU, Intel ARC GPU
-- OneAPI requirements: oneAPI 2025.1
+- OneAPI requirements: oneAPI 2025.3
+- Dependency: [vllm-xpu-kernels](https://github.com/vllm-project/vllm-xpu-kernels): a package provide all necessary vllm custom kernel when running vLLM on Intel GPU platform, 
 - Python: 3.12
 !!! warning
-    The provided IPEX whl is Python3.12 specific so this version is a MUST.
+    The provided vllm-xpu-kernels whl is Python3.12 specific so this version is a MUST.
 
 # --8<-- [end:requirements]
 # --8<-- [start:set-up-using-python]
@@ -24,7 +25,7 @@ Currently, there are no pre-built XPU wheels.
 # --8<-- [end:pre-built-wheels]
 # --8<-- [start:build-wheel-from-source]
 
-- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers) and [Intel OneAPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) 2025.1 or later.
+- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers) and [Intel OneAPI](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) 2025.3 or later.
 - Second, install Python packages for vLLM XPU backend building:
 
 ```bash
@@ -37,7 +38,7 @@ pip install -v -r requirements/xpu.txt
 - Then, build and install vLLM XPU backend:
 
 ```bash
-VLLM_TARGET_DEVICE=xpu python setup.py install
+VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
 ```
 
 # --8<-- [end:build-wheel-from-source]

docs/serving/openai_compatible_server.md

@@ -311,7 +311,7 @@ and passing a list of `messages` in the request. Refer to the examples below for
     vllm serve TIGER-Lab/VLM2Vec-Full --runner pooling \
       --trust-remote-code \
       --max-model-len 4096 \
-      --chat-template examples/template_vlm2vec_phi3v.jinja
+      --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja
     ```
 
     !!! important
@@ -319,7 +319,7 @@ and passing a list of `messages` in the request. Refer to the examples below for
         to run this model in embedding mode instead of text generation mode.
 
         The custom chat template is completely different from the original one for this model,
-        and can be found here: [examples/template_vlm2vec_phi3v.jinja](../../examples/template_vlm2vec_phi3v.jinja)
+        and can be found here: [examples/pooling/embed/template/vlm2vec_phi3v.jinja](../../examples/pooling/embed/template/vlm2vec_phi3v.jinja)
 
     Since the request schema is not defined by OpenAI client, we post a request to the server using the lower-level `requests` library:
 
@@ -359,14 +359,14 @@ and passing a list of `messages` in the request. Refer to the examples below for
     vllm serve MrLight/dse-qwen2-2b-mrl-v1 --runner pooling \
       --trust-remote-code \
       --max-model-len 8192 \
-      --chat-template examples/template_dse_qwen2_vl.jinja
+      --chat-template examples/pooling/embed/template/dse_qwen2_vl.jinja
     ```
 
     !!! important
         Like with VLM2Vec, we have to explicitly pass `--runner pooling`.
 
         Additionally, `MrLight/dse-qwen2-2b-mrl-v1` requires an EOS token for embeddings, which is handled
-        by a custom chat template: [examples/template_dse_qwen2_vl.jinja](../../examples/template_dse_qwen2_vl.jinja)
+        by a custom chat template: [examples/pooling/embed/template/dse_qwen2_vl.jinja](../../examples/pooling/embed/template/dse_qwen2_vl.jinja)
 
     !!! important
         `MrLight/dse-qwen2-2b-mrl-v1` requires a placeholder image of the minimum image size for text query embeddings. See the full code
@@ -532,7 +532,7 @@ The following [sampling parameters](../api/README.md#inference-parameters) are s
 ??? code
 
     ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:transcription-sampling-params"
+    --8<-- "vllm/entrypoints/openai/speech_to_text/protocol.py:transcription-sampling-params"
     ```
 
 The following extra parameters are supported:
@@ -540,7 +540,7 @@ The following extra parameters are supported:
 ??? code
 
     ```python
-    --8<-- "vllm/entrypoints/openai/protocol.py:transcription-extra-params"
+    --8<-- "vllm/entrypoints/openai/speech_to_text/protocol.py:transcription-extra-params"
     ```
 
 ### Translations API
@@ -560,13 +560,13 @@ Code example: [examples/online_serving/openai_translation_client.py](../../examp
 The following [sampling parameters](../api/README.md#inference-parameters) are supported.
 
 ```python
---8<-- "vllm/entrypoints/openai/protocol.py:translation-sampling-params"
+--8<-- "vllm/entrypoints/openai/speech_to_text/protocol.py:translation-sampling-params"
 ```
 
 The following extra parameters are supported:
 
 ```python
---8<-- "vllm/entrypoints/openai/protocol.py:translation-extra-params"
+--8<-- "vllm/entrypoints/openai/speech_to_text/protocol.py:translation-extra-params"
 ```
 
 ### Realtime API
@@ -954,28 +954,34 @@ You can pass multi-modal inputs to scoring models by passing `content` including
 
         ```python
         import requests
-
+        
         response = requests.post(
             "http://localhost:8000/v1/score",
             json={
                 "model": "jinaai/jina-reranker-m0",
                 "queries": "slm markdown",
-                "documents": {
-                    "content": [
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
-                            },
-                        },
-                        {
-                            "type": "image_url",
-                            "image_url": {
-                                "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
-                            },
-                        },
-                    ],
-                },
+                "documents": [
+                    {
+                        "content": [
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
+                                },
+                            }
+                        ],
+                    },
+                    {
+                        "content": [
+                            {
+                                "type": "image_url",
+                                "image_url": {
+                                    "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
+                                },
+                            }
+                        ]
+                    },
+                ],
             },
         )
         response.raise_for_status()
@@ -1001,7 +1007,6 @@ The following Score API parameters are supported:
 
 ```python
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
---8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
 ```
 
 The following extra parameters are supported:
@@ -1009,7 +1014,6 @@ The following extra parameters are supported:
 ```python
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
---8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
 ```
 
 ### Re-rank API
@@ -1092,7 +1096,6 @@ The following Re-rank API parameters are supported:
 ```python
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
---8<-- "vllm/entrypoints/pooling/score/protocol.py:score-extra-params"
 ```
 
 The following extra parameters are supported:
@@ -1100,7 +1103,6 @@ The following extra parameters are supported:
 ```python
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:pooling-common-extra-params"
 --8<-- "vllm/entrypoints/pooling/base/protocol.py:classify-extra-params"
---8<-- "vllm/entrypoints/pooling/score/protocol.py:rerank-extra-params"
 ```
 
 ## Ray Serve LLM

examples/pooling/classify/vision_classification_online.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+"""Example Python client for multimodal classification API using vLLM API server
+NOTE:
+    start a supported multimodal classification model server with `vllm serve`, e.g.
+    vllm serve muziyongshixin/Qwen2.5-VL-7B-for-VideoCls \
+         --runner pooling \
+         --max-model-len 5000 \
+         --limit-mm-per-prompt '{"video": 1}' \
+         --hf-overrides '{"text_config": {"architectures": ["Qwen2_5_VLForSequenceClassification"]}}'
+"""
+
+import argparse
+import pprint
+
+import requests
+
+from vllm.multimodal.utils import encode_image_url, fetch_image
+
+input_text = "This product was excellent and exceeded my expectations"
+image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
+image_base64 = {"url": encode_image_url(fetch_image(image_url))}
+video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"
+
+
+def parse_args():
+    parse = argparse.ArgumentParser()
+    parse.add_argument("--host", type=str, default="localhost")
+    parse.add_argument("--port", type=int, default=8000)
+    return parse.parse_args()
+
+
+def main(args):
+    base_url = f"http://{args.host}:{args.port}"
+    models_url = base_url + "/v1/models"
+    classify_url = base_url + "/classify"
+
+    response = requests.get(models_url)
+    model_name = response.json()["data"][0]["id"]
+
+    print("Text classification output:")
+    messages = [
+        {
+            "role": "assistant",
+            "content": "Please classify this text request.",
+        },
+        {
+            "role": "user",
+            "content": input_text,
+        },
+    ]
+    response = requests.post(
+        classify_url,
+        json={"model": model_name, "messages": messages},
+    )
+    pprint.pprint(response.json())
+
+    print("Image url classification output:")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Please classify this image."},
+                {"type": "image_url", "image_url": {"url": image_url}},
+            ],
+        }
+    ]
+    response = requests.post(
+        classify_url,
+        json={"model": model_name, "messages": messages},
+    )
+    pprint.pprint(response.json())
+
+    print("Image base64 classification output:")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Please classify this image."},
+                {"type": "image_url", "image_url": image_base64},
+            ],
+        }
+    ]
+    response = requests.post(
+        classify_url,
+        json={"model": model_name, "messages": messages},
+    )
+    pprint.pprint(response.json())
+
+    print("Video url classification output:")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Please classify this video."},
+                {"type": "video_url", "video_url": {"url": video_url}},
+            ],
+        }
+    ]
+    response = requests.post(
+        classify_url,
+        json={"model": model_name, "messages": messages},
+    )
+    pprint.pprint(response.json())
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/pooling/embed/vision_embedding_offline.py

@@ -11,23 +11,79 @@ on HuggingFace model repository.
 
 import argparse
 from dataclasses import asdict
+from pathlib import Path
 
 from PIL.Image import Image
 
 from vllm import LLM, EngineArgs
 from vllm.multimodal.utils import fetch_image
+from vllm.utils.print_utils import print_embeddings
+
+ROOT_DIR = Path(__file__).parent.parent.parent
+EMBED_TEMPLATE_DIR = ROOT_DIR / "pooling/embed/template/"
 
 image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
 text = "A cat standing in the snow."
 multi_modal_data = {"image": fetch_image(image_url)}
 
 
-def print_embeddings(embeds: list[float]):
-    embeds_trimmed = (str(embeds[:4])[:-1] + ", ...]") if len(embeds) > 4 else embeds
-    print(f"Embeddings: {embeds_trimmed} (size={len(embeds)})")
+def run_clip(seed: int):
+    engine_args = EngineArgs(
+        model="openai/clip-vit-base-patch32",
+        runner="pooling",
+        limit_mm_per_prompt={"image": 1},
+    )
+
+    llm = LLM(**asdict(engine_args) | {"seed": seed})
+
+    print("Text embedding output:")
+    outputs = llm.embed(text, use_tqdm=False)
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image embedding output:")
+    prompt = ""  # For image input, make sure that the prompt text is empty
+    outputs = llm.embed(
+        {
+            "prompt": prompt,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
 
 
-def run_qwen3_vl():
+def run_e5_v(seed: int):
+    engine_args = EngineArgs(
+        model="royokong/e5-v",
+        runner="pooling",
+        max_model_len=4096,
+        limit_mm_per_prompt={"image": 1},
+    )
+
+    llm = LLM(**asdict(engine_args) | {"seed": seed})
+
+    llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n"  # noqa: E501
+
+    print("Text embedding output:")
+    prompt_text = llama3_template.format(
+        f"{text}\nSummary above sentence in one word: "
+    )
+    outputs = llm.embed(prompt_text, use_tqdm=False)
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image embedding output:")
+    prompt_image = llama3_template.format("<image>\nSummary above image in one word: ")
+    outputs = llm.embed(
+        {
+            "prompt": prompt_image,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
+
+
+def run_qwen3_vl(seed: int):
     try:
         from qwen_vl_utils import smart_resize
     except ModuleNotFoundError:
@@ -61,20 +117,20 @@ def run_qwen3_vl():
     )
     default_instruction = "Represent the user's input."
     image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
-    text_prompt = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
-    image_prompt = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}<|im_end|>\n<|im_start|>assistant\n"
-    image_text_prompt = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}{text}<|im_end|>\n<|im_start|>assistant\n"
+    prompt_text = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
+    prompt_image = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}<|im_end|>\n<|im_start|>assistant\n"
+    prompt_image_text = f"<|im_start|>system\n{default_instruction}<|im_end|>\n<|im_start|>user\n{image_placeholder}{text}<|im_end|>\n<|im_start|>assistant\n"
 
-    llm = LLM(**asdict(engine_args))
+    llm = LLM(**asdict(engine_args) | {"seed": seed})
 
     print("Text embedding output:")
-    outputs = llm.embed(text_prompt, use_tqdm=False)
+    outputs = llm.embed(prompt_text, use_tqdm=False)
     print_embeddings(outputs[0].outputs.embedding)
 
     print("Image embedding output:")
     outputs = llm.embed(
         {
-            "prompt": image_prompt,
+            "prompt": prompt_image,
             "multi_modal_data": multi_modal_data,
         },
         use_tqdm=False,
@@ -84,7 +140,162 @@ def run_qwen3_vl():
     print("Image+Text embedding output:")
     outputs = llm.embed(
         {
-            "prompt": image_text_prompt,
+            "prompt": prompt_image_text,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
+
+
+def run_siglip(seed: int):
+    engine_args = EngineArgs(
+        model="google/siglip-base-patch16-224",
+        runner="pooling",
+        limit_mm_per_prompt={"image": 1},
+    )
+
+    llm = LLM(**asdict(engine_args) | {"seed": seed})
+
+    print("Text embedding output:")
+    outputs = llm.embed(text, use_tqdm=False)
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image embedding output:")
+    prompt = ""  # For image input, make sure that the prompt text is empty
+    outputs = llm.embed(
+        {
+            "prompt": prompt,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
+
+
+def run_vlm2vec_phi3v(seed: int):
+    engine_args = EngineArgs(
+        model="TIGER-Lab/VLM2Vec-Full",
+        runner="pooling",
+        max_model_len=4096,
+        trust_remote_code=True,
+        mm_processor_kwargs={"num_crops": 4},
+        limit_mm_per_prompt={"image": 1},
+    )
+
+    llm = LLM(**asdict(engine_args) | {"seed": seed})
+    image_token = "<|image_1|>"
+
+    print("Text embedding output:")
+    prompt_text = f"Find me an everyday image that matches the given caption: {text}"
+    outputs = llm.embed(prompt_text, use_tqdm=False)
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image embedding output:")
+    prompt_image = f"{image_token} Find a day-to-day image that looks similar to the provided image."  # noqa: E501
+    outputs = llm.embed(
+        {
+            "prompt": prompt_image,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image+Text embedding output:")
+    prompt_image_text = (
+        f"{image_token} Represent the given image with the following question: {text}"  # noqa: E501
+    )
+    outputs = llm.embed(
+        {
+            "prompt": prompt_image_text,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
+
+
+def run_vlm2vec_qwen2vl(seed: int):
+    # vLLM does not support LoRA adapters on multi-modal encoder,
+    # so we merge the weights first
+    from huggingface_hub.constants import HF_HUB_CACHE
+    from peft import PeftConfig, PeftModel
+    from transformers import AutoModelForImageTextToText, AutoProcessor
+
+    from vllm.entrypoints.chat_utils import load_chat_template
+
+    model_id = "TIGER-Lab/VLM2Vec-Qwen2VL-2B"
+
+    base_model = AutoModelForImageTextToText.from_pretrained(model_id)
+    lora_model = PeftModel.from_pretrained(
+        base_model,
+        model_id,
+        config=PeftConfig.from_pretrained(model_id),
+    )
+    model = lora_model.merge_and_unload().to(dtype=base_model.dtype)
+    model._hf_peft_config_loaded = False  # Needed to save the merged model
+
+    processor = AutoProcessor.from_pretrained(
+        model_id,
+        # `min_pixels` and `max_pixels` are deprecated for
+        # transformers `preprocessor_config.json`
+        size={"shortest_edge": 3136, "longest_edge": 12845056},
+    )
+    processor.chat_template = load_chat_template(
+        # The original chat template is not correct
+        EMBED_TEMPLATE_DIR / "vlm2vec_qwen2vl.jinja",
+    )
+
+    merged_path = str(
+        Path(HF_HUB_CACHE) / ("models--" + model_id.replace("/", "--") + "-vllm")
+    )
+    print(f"Saving merged model to {merged_path}...")
+    print(
+        "NOTE: This directory is not tracked by `huggingface_hub` "
+        "so you have to delete this manually if you don't want it anymore."
+    )
+    model.save_pretrained(merged_path)
+    processor.save_pretrained(merged_path)
+    print("Done!")
+
+    engine_args = EngineArgs(
+        model=merged_path,
+        runner="pooling",
+        max_model_len=4096,
+        mm_processor_kwargs={
+            "min_pixels": 3136,
+            "max_pixels": 12845056,
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+    llm = LLM(**asdict(engine_args) | {"seed": seed})
+    image_token = "<|image_pad|>"
+
+    print("Text embedding output:")
+    prompt_text = f"Find me an everyday image that matches the given caption: {text}"
+    outputs = llm.embed(prompt_text, use_tqdm=False)
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image embedding output:")
+    prompt_image = f"{image_token} Find a day-to-day image that looks similar to the provided image."  # noqa: E501
+    outputs = llm.embed(
+        {
+            "prompt": prompt_image,
+            "multi_modal_data": multi_modal_data,
+        },
+        use_tqdm=False,
+    )
+    print_embeddings(outputs[0].outputs.embedding)
+
+    print("Image+Text embedding output:")
+    prompt_image_text = (
+        f"{image_token} Represent the given image with the following question: {text}"  # noqa: E501
+    )
+    outputs = llm.embed(
+        {
+            "prompt": prompt_image_text,
             "multi_modal_data": multi_modal_data,
         },
         use_tqdm=False,
@@ -93,7 +304,12 @@ def run_qwen3_vl():
 
 
 model_example_map = {
+    "clip": run_clip,
+    "e5_v": run_e5_v,
     "qwen3_vl": run_qwen3_vl,
+    "siglip": run_siglip,
+    "vlm2vec_phi3v": run_vlm2vec_phi3v,
+    "vlm2vec_qwen2vl": run_vlm2vec_qwen2vl,
 }
 
 
@@ -103,16 +319,23 @@ def parse_args():
     )
     parser.add_argument(
         "--model",
+        "-m",
         type=str,
+        default="vlm2vec_phi3v",
         choices=model_example_map.keys(),
-        required=True,
         help="The name of the embedding model.",
     )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Set the seed when initializing `vllm.LLM`.",
+    )
     return parser.parse_args()
 
 
 def main(args):
-    model_example_map[args.model]()
+    model_example_map[args.model](args.seed)
 
 
 if __name__ == "__main__":

examples/pooling/embed/vision_embedding_online.py

@@ -17,6 +17,8 @@ from openai.types.chat import ChatCompletionMessageParam
 from openai.types.create_embedding_response import CreateEmbeddingResponse
 from PIL import Image
 
+from vllm.utils.print_utils import print_embeddings
+
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"
 openai_api_base = "http://localhost:8000/v1"
@@ -51,11 +53,6 @@ def create_chat_embeddings(
     )
 
 
-def print_embeddings(embeds):
-    embeds_trimmed = (str(embeds[:4])[:-1] + ", ...]") if len(embeds) > 4 else embeds
-    print(f"Embeddings: {embeds_trimmed} (size={len(embeds)})")
-
-
 def run_clip(client: OpenAI, model: str):
     """
     Start the server using:
@@ -105,7 +102,7 @@ def run_dse_qwen2_vl(client: OpenAI, model: str):
         --runner pooling \
         --trust-remote-code \
         --max-model-len 8192 \
-        --chat-template examples/template_dse_qwen2_vl.jinja
+        --chat-template examples/pooling/embed/template/dse_qwen2_vl.jinja
     """
     response = create_chat_embeddings(
         client,
@@ -316,7 +313,7 @@ def run_vlm2vec(client: OpenAI, model: str):
         --runner pooling \
         --trust-remote-code \
         --max-model-len 4096 \
-        --chat-template examples/template_vlm2vec_phi3v.jinja
+        --chat-template examples/pooling/embed/template/vlm2vec_phi3v.jinja
     """
 
     response = create_chat_embeddings(

examples/pooling/pooling/vision_language_pooling.py

@@ -1,441 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM for running offline inference with
-the correct prompt format on vision language models for multimodal pooling.
-
-For most models, the prompt format should follow corresponding examples
-on HuggingFace model repository.
-"""
-
-from argparse import Namespace
-from dataclasses import asdict
-from pathlib import Path
-from typing import Literal, NamedTuple, TypeAlias, TypedDict, get_args
-
-from PIL.Image import Image
-
-from vllm import LLM, EngineArgs
-from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
-from vllm.multimodal.utils import fetch_image
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-ROOT_DIR = Path(__file__).parent.parent.parent
-EXAMPLES_DIR = ROOT_DIR / "examples"
-
-
-class TextQuery(TypedDict):
-    modality: Literal["text"]
-    text: str
-
-
-class ImageQuery(TypedDict):
-    modality: Literal["image"]
-    image: Image
-
-
-class TextImageQuery(TypedDict):
-    modality: Literal["text+image"]
-    text: str
-    image: Image
-
-
-class TextImagesQuery(TypedDict):
-    modality: Literal["text+images"]
-    text: str
-    image: ScoreMultiModalParam
-
-
-QueryModality = Literal["text", "image", "text+image", "text+images"]
-Query: TypeAlias = TextQuery | ImageQuery | TextImageQuery | TextImagesQuery
-
-
-class ModelRequestData(NamedTuple):
-    engine_args: EngineArgs
-    prompt: str | None = None
-    image: Image | None = None
-    query: str | None = None
-    documents: ScoreMultiModalParam | None = None
-
-
-def run_clip(query: Query) -> ModelRequestData:
-    if query["modality"] == "text":
-        prompt = query["text"]
-        image = None
-    elif query["modality"] == "image":
-        prompt = ""  # For image input, make sure that the prompt text is empty
-        image = query["image"]
-    else:
-        modality = query["modality"]
-        raise ValueError(f"Unsupported query modality: '{modality}'")
-
-    engine_args = EngineArgs(
-        model="openai/clip-vit-base-patch32",
-        runner="pooling",
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
-def run_e5_v(query: Query) -> ModelRequestData:
-    llama3_template = "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n \n"  # noqa: E501
-
-    if query["modality"] == "text":
-        text = query["text"]
-        prompt = llama3_template.format(f"{text}\nSummary above sentence in one word: ")
-        image = None
-    elif query["modality"] == "image":
-        prompt = llama3_template.format("<image>\nSummary above image in one word: ")
-        image = query["image"]
-    else:
-        modality = query["modality"]
-        raise ValueError(f"Unsupported query modality: '{modality}'")
-
-    engine_args = EngineArgs(
-        model="royokong/e5-v",
-        runner="pooling",
-        max_model_len=4096,
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
-def run_jinavl_reranker(query: Query) -> ModelRequestData:
-    if query["modality"] != "text+images":
-        raise ValueError(f"Unsupported query modality: '{query['modality']}'")
-
-    engine_args = EngineArgs(
-        model="jinaai/jina-reranker-m0",
-        runner="pooling",
-        max_model_len=32768,
-        trust_remote_code=True,
-        mm_processor_kwargs={
-            "min_pixels": 3136,
-            "max_pixels": 602112,
-        },
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        query=query["text"],
-        documents=query["image"],
-    )
-
-
-def run_qwen3_vl(query: Query) -> ModelRequestData:
-    image_placeholder = "<vision_start><|image_pad|><vision_end>"
-    if query["modality"] == "text":
-        prompt = query["text"]
-        image = None
-    elif query["modality"] == "image":
-        prompt = image_placeholder
-        image = query["image"]
-    elif query["modality"] == "text+image":
-        text = query["text"]
-        prompt = f"{image_placeholder}\n{text}"
-        image = query["image"]
-    else:
-        modality = query["modality"]
-        raise ValueError(f"Unsupported query modality: '{modality}'")
-
-    engine_args = EngineArgs(
-        model="Qwen/Qwen3-VL-Embedding-2B",
-        runner="pooling",
-        max_model_len=8192,
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
-def run_siglip(query: Query) -> ModelRequestData:
-    if query["modality"] == "text":
-        prompt = query["text"]
-        image = None
-    elif query["modality"] == "image":
-        prompt = ""  # For image input, make sure that the prompt text is empty
-        image = query["image"]
-    else:
-        modality = query["modality"]
-        raise ValueError(f"Unsupported query modality: '{modality}'")
-
-    engine_args = EngineArgs(
-        model="google/siglip-base-patch16-224",
-        runner="pooling",
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
-def _get_vlm2vec_prompt_image(query: Query, image_token: str):
-    if query["modality"] == "text":
-        text = query["text"]
-        prompt = f"Find me an everyday image that matches the given caption: {text}"
-        image = None
-    elif query["modality"] == "image":
-        prompt = f"{image_token} Find a day-to-day image that looks similar to the provided image."  # noqa: E501
-        image = query["image"]
-    elif query["modality"] == "text+image":
-        text = query["text"]
-        prompt = f"{image_token} Represent the given image with the following question: {text}"  # noqa: E501
-        image = query["image"]
-    else:
-        modality = query["modality"]
-        raise ValueError(f"Unsupported query modality: {modality!r}")
-
-    return prompt, image
-
-
-def run_vlm2vec_phi3v(query: Query) -> ModelRequestData:
-    prompt, image = _get_vlm2vec_prompt_image(query, "<|image_1|>")
-
-    engine_args = EngineArgs(
-        model="TIGER-Lab/VLM2Vec-Full",
-        runner="pooling",
-        max_model_len=4096,
-        trust_remote_code=True,
-        mm_processor_kwargs={"num_crops": 4},
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
-def run_vlm2vec_qwen2vl(query: Query) -> ModelRequestData:
-    # vLLM does not support LoRA adapters on multi-modal encoder,
-    # so we merge the weights first
-    from huggingface_hub.constants import HF_HUB_CACHE
-    from peft import PeftConfig, PeftModel
-    from transformers import AutoModelForImageTextToText, AutoProcessor
-
-    from vllm.entrypoints.chat_utils import load_chat_template
-
-    model_id = "TIGER-Lab/VLM2Vec-Qwen2VL-2B"
-
-    base_model = AutoModelForImageTextToText.from_pretrained(model_id)
-    lora_model = PeftModel.from_pretrained(
-        base_model,
-        model_id,
-        config=PeftConfig.from_pretrained(model_id),
-    )
-    model = lora_model.merge_and_unload().to(dtype=base_model.dtype)
-    model._hf_peft_config_loaded = False  # Needed to save the merged model
-
-    processor = AutoProcessor.from_pretrained(
-        model_id,
-        # `min_pixels` and `max_pixels` are deprecated for
-        # transformers `preprocessor_config.json`
-        size={"shortest_edge": 3136, "longest_edge": 12845056},
-    )
-    processor.chat_template = load_chat_template(
-        # The original chat template is not correct
-        EXAMPLES_DIR / "template_vlm2vec_qwen2vl.jinja",
-    )
-
-    merged_path = str(
-        Path(HF_HUB_CACHE) / ("models--" + model_id.replace("/", "--") + "-vllm")
-    )
-    print(f"Saving merged model to {merged_path}...")
-    print(
-        "NOTE: This directory is not tracked by `huggingface_hub` "
-        "so you have to delete this manually if you don't want it anymore."
-    )
-    model.save_pretrained(merged_path)
-    processor.save_pretrained(merged_path)
-    print("Done!")
-
-    prompt, image = _get_vlm2vec_prompt_image(query, "<|image_pad|>")
-
-    engine_args = EngineArgs(
-        model=merged_path,
-        runner="pooling",
-        max_model_len=4096,
-        mm_processor_kwargs={
-            "min_pixels": 3136,
-            "max_pixels": 12845056,
-        },
-        limit_mm_per_prompt={"image": 1},
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompt=prompt,
-        image=image,
-    )
-
-
-def get_query(modality: QueryModality):
-    if modality == "text":
-        return TextQuery(modality="text", text="A dog sitting in the grass")
-
-    if modality == "image":
-        return ImageQuery(
-            modality="image",
-            image=fetch_image(
-                "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg"  # noqa: E501
-            ),
-        )
-
-    if modality == "text+image":
-        return TextImageQuery(
-            modality="text+image",
-            text="A cat standing in the snow.",
-            image=fetch_image(
-                "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"  # noqa: E501
-            ),
-        )
-
-    if modality == "text+images":
-        return TextImagesQuery(
-            modality="text+images",
-            text="slm markdown",
-            image={
-                "content": [
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
-                        },
-                    },
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
-                        },
-                    },
-                ]
-            },
-        )
-
-    msg = f"Modality {modality} is not supported."
-    raise ValueError(msg)
-
-
-def run_encode(model: str, modality: QueryModality, seed: int):
-    query = get_query(modality)
-    req_data = model_example_map[model](query)
-
-    # Disable other modalities to save memory
-    default_limits = {"image": 0, "video": 0, "audio": 0}
-    req_data.engine_args.limit_mm_per_prompt = default_limits | dict(
-        req_data.engine_args.limit_mm_per_prompt or {}
-    )
-
-    engine_args = asdict(req_data.engine_args) | {"seed": seed}
-    llm = LLM(**engine_args)
-
-    mm_data = {}
-    if req_data.image is not None:
-        mm_data["image"] = req_data.image
-
-    outputs = llm.embed(
-        {
-            "prompt": req_data.prompt,
-            "multi_modal_data": mm_data,
-        }
-    )
-
-    print("-" * 50)
-    for output in outputs:
-        print(output.outputs.embedding)
-        print("-" * 50)
-
-
-def run_score(model: str, modality: QueryModality, seed: int):
-    query = get_query(modality)
-    req_data = model_example_map[model](query)
-
-    engine_args = asdict(req_data.engine_args) | {"seed": seed}
-    llm = LLM(**engine_args)
-
-    outputs = llm.score(req_data.query, req_data.documents)
-
-    print("-" * 30)
-    print([output.outputs.score for output in outputs])
-    print("-" * 30)
-
-
-model_example_map = {
-    "clip": run_clip,
-    "e5_v": run_e5_v,
-    "jinavl_reranker": run_jinavl_reranker,
-    "qwen3_vl": run_qwen3_vl,
-    "siglip": run_siglip,
-    "vlm2vec_phi3v": run_vlm2vec_phi3v,
-    "vlm2vec_qwen2vl": run_vlm2vec_qwen2vl,
-}
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(
-        description="Demo on using vLLM for offline inference with "
-        "vision language models for multimodal pooling tasks."
-    )
-    parser.add_argument(
-        "--model-name",
-        "-m",
-        type=str,
-        default="vlm2vec_phi3v",
-        choices=model_example_map.keys(),
-        help="The name of the embedding model.",
-    )
-    parser.add_argument(
-        "--task",
-        "-t",
-        type=str,
-        default="embedding",
-        choices=["embedding", "scoring"],
-        help="The task type.",
-    )
-    parser.add_argument(
-        "--modality",
-        type=str,
-        default="image",
-        choices=get_args(QueryModality),
-        help="Modality of the input.",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=0,
-        help="Set the seed when initializing `vllm.LLM`.",
-    )
-    return parser.parse_args()
-
-
-def main(args: Namespace):
-    if args.task == "embedding":
-        run_encode(args.model_name, args.modality, args.seed)
-    elif args.task == "scoring":
-        run_score(args.model_name, args.modality, args.seed)
-    else:
-        raise ValueError(f"Unsupported task: {args.task}")
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

examples/pooling/score/vision_rerank_api_online.py

@@ -30,6 +30,7 @@ document = (
     "as the dog offers its paw in a heartwarming display of companionship and trust."
 )
 image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
+video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
 documents = [
     {
         "type": "text",
@@ -43,6 +44,10 @@ documents = [
         "type": "image_url",
         "image_url": {"url": encode_image_url(fetch_image(image_url))},
     },
+    {
+        "type": "video_url",
+        "video_url": {"url": video_url},
+    },
 ]
 
 
@@ -89,6 +94,15 @@ def main(args):
     response = requests.post(rerank_url, json=prompt)
     pprint.pprint(response.json())
 
+    print("Query: string & Document: video url")
+    prompt = {
+        "model": model,
+        "query": query,
+        "documents": {"content": [documents[3]]},
+    }
+    response = requests.post(rerank_url, json=prompt)
+    pprint.pprint(response.json())
+
     print("Query: string & Document: text + image url")
     prompt = {
         "model": model,

examples/pooling/score/vision_reranker_offline.py

@@ -15,20 +15,47 @@ from pathlib import Path
 from typing import NamedTuple
 
 from vllm import LLM, EngineArgs
-from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
+from vllm.multimodal.utils import encode_image_url, fetch_image
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 TEMPLATE_HOME = Path(__file__).parent / "template"
 
 
+query = "A woman playing with her dog on a beach at sunset."
+document = (
+    "A woman shares a joyful moment with her golden retriever on a sun-drenched "
+    "beach at sunset, as the dog offers its paw in a heartwarming display of "
+    "companionship and trust."
+)
+image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
+video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
+documents = [
+    {
+        "type": "text",
+        "text": document,
+    },
+    {
+        "type": "image_url",
+        "image_url": {"url": image_url},
+    },
+    {
+        "type": "image_url",
+        "image_url": {"url": encode_image_url(fetch_image(image_url))},
+    },
+    {
+        "type": "video_url",
+        "video_url": {"url": video_url},
+    },
+]
+
+
 class RerankModelData(NamedTuple):
     engine_args: EngineArgs
     chat_template: str | None = None
+    modality: set[str] = {}
 
 
-def run_jinavl_reranker(modality: str) -> RerankModelData:
-    assert modality == "image"
-
+def run_jinavl_reranker() -> RerankModelData:
     engine_args = EngineArgs(
         model="jinaai/jina-reranker-m0",
         runner="pooling",
@@ -38,19 +65,15 @@ def run_jinavl_reranker(modality: str) -> RerankModelData:
             "min_pixels": 3136,
             "max_pixels": 602112,
         },
-        limit_mm_per_prompt={modality: 1},
-    )
-    return RerankModelData(
-        engine_args=engine_args,
     )
+    return RerankModelData(engine_args=engine_args, modality={"image"})
 
 
-def run_qwen3_vl_reranker(modality: str) -> RerankModelData:
+def run_qwen3_vl_reranker() -> RerankModelData:
     engine_args = EngineArgs(
         model="Qwen/Qwen3-VL-Reranker-2B",
         runner="pooling",
         max_model_len=16384,
-        limit_mm_per_prompt={modality: 1},
         # HuggingFace model configuration overrides required for compatibility
         hf_overrides={
             # Manually route to sequence classification architecture
@@ -71,10 +94,11 @@ def run_qwen3_vl_reranker(modality: str) -> RerankModelData:
     return RerankModelData(
         engine_args=engine_args,
         chat_template=chat_template,
+        modality={"image", "video"},
     )
 
 
-model_example_map: dict[str, Callable[[str], RerankModelData]] = {
+model_example_map: dict[str, Callable[[], RerankModelData]] = {
     "jinavl_reranker": run_jinavl_reranker,
     "qwen3_vl_reranker": run_qwen3_vl_reranker,
 }
@@ -93,78 +117,67 @@ def parse_args():
         choices=model_example_map.keys(),
         help="The name of the reranker model.",
     )
-    parser.add_argument(
-        "--modality",
-        type=str,
-        default="image",
-        choices=["image", "video"],
-        help="Modality of the multimodal input (image or video).",
-    )
     return parser.parse_args()
 
 
-def get_multi_modal_input(modality: str) -> tuple[str, ScoreMultiModalParam]:
-    # Sample query for testing the reranker
-    if modality == "image":
-        query = "A woman playing with her dog on a beach at sunset."
-        # Sample multimodal documents to be scored against the query
-        # Each document contains an image URL that will be fetched and processed
-        documents: ScoreMultiModalParam = {
-            "content": [
-                {
-                    "type": "text",
-                    "text": (
-                        "A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "  # noqa: E501
-                        "as the dog offers its paw in a heartwarming display of companionship and trust."  # noqa: E501
-                    ),
-                },
-                {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
-                    },
-                },
-            ]
-        }
-    elif modality == "video":
-        query = "A girl is drawing pictures on an ipad."
-        # Sample video documents to be scored against the query
-        documents: ScoreMultiModalParam = {
-            "content": [
-                {
-                    "type": "text",
-                    "text": "A girl is drawing a guitar on her ipad with Apple Pencil.",
-                },
-                {
-                    "type": "video_url",
-                    "video_url": {
-                        "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
-                    },
-                },
-            ]
-        }
-    else:
-        raise ValueError(f"Unsupported modality: {modality}")
-    return query, documents
-
-
 def main(args: Namespace):
     # Run the selected reranker model
-    modality = args.modality
-    model_request = model_example_map[args.model_name](modality)
+    model_request = model_example_map[args.model_name]()
     engine_args = model_request.engine_args
 
     llm = LLM(**asdict(engine_args))
 
-    query, documents = get_multi_modal_input(modality)
-    outputs = llm.score(query, documents, chat_template=model_request.chat_template)
+    print("Query: string & Document: string")
+    outputs = llm.score(query, document)
+    print("Relevance scores:", [output.outputs.score for output in outputs])
 
-    print("-" * 50)
-    print(f"Model: {engine_args.model}")
-    print(f"Modality: {modality}")
-    print(f"Query: {query}")
+    print("Query: string & Document: text")
+    outputs = llm.score(
+        query, {"content": [documents[0]]}, chat_template=model_request.chat_template
+    )
+    print("Relevance scores:", [output.outputs.score for output in outputs])
+
+    print("Query: string & Document: image url")
+    outputs = llm.score(
+        query, {"content": [documents[1]]}, chat_template=model_request.chat_template
+    )
+    print("Relevance scores:", [output.outputs.score for output in outputs])
+
+    print("Query: string & Document: image base64")
+    outputs = llm.score(
+        query, {"content": [documents[2]]}, chat_template=model_request.chat_template
+    )
+    print("Relevance scores:", [output.outputs.score for output in outputs])
+
+    if "video" in model_request.modality:
+        print("Query: string & Document: video url")
+        outputs = llm.score(
+            query,
+            {"content": [documents[3]]},
+            chat_template=model_request.chat_template,
+        )
+        print("Relevance scores:", [output.outputs.score for output in outputs])
+
+    print("Query: string & Document: text + image url")
+    outputs = llm.score(
+        query,
+        {"content": [documents[0], documents[1]]},
+        chat_template=model_request.chat_template,
+    )
+    print("Relevance scores:", [output.outputs.score for output in outputs])
+
+    print("Query: string & Document: list")
+    outputs = llm.score(
+        query,
+        [
+            document,
+            {"content": [documents[0]]},
+            {"content": [documents[1]]},
+            {"content": [documents[0], documents[1]]},
+        ],
+        chat_template=model_request.chat_template,
+    )
     print("Relevance scores:", [output.outputs.score for output in outputs])
-    print("-" * 50)
 
 
 if __name__ == "__main__":

examples/pooling/score/vision_score_api_online.py

@@ -29,6 +29,7 @@ document = (
     "as the dog offers its paw in a heartwarming display of companionship and trust."
 )
 image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
+video_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4"
 documents = [
     {
         "type": "text",
@@ -42,6 +43,10 @@ documents = [
         "type": "image_url",
         "image_url": {"url": encode_image_url(fetch_image(image_url))},
     },
+    {
+        "type": "video_url",
+        "video_url": {"url": video_url},
+    },
 ]
 
 
@@ -92,6 +97,15 @@ def main(args):
     response = requests.post(score_url, json=prompt)
     pprint.pprint(response.json())
 
+    print("Query: string & Document: video url")
+    prompt = {
+        "model": model,
+        "queries": query,
+        "documents": {"content": [documents[3]]},
+    }
+    response = requests.post(score_url, json=prompt)
+    pprint.pprint(response.json())
+
     print("Query: string & Document: text + image url")
     prompt = {
         "model": model,

pyproject.toml

@@ -6,7 +6,7 @@ requires = [
     "packaging>=24.2",
     "setuptools>=77.0.3,<81.0.0",
     "setuptools-scm>=8.0",
-    "torch == 2.9.1",
+    "torch == 2.10.0",
     "wheel",
     "jinja2",
     "grpcio-tools==1.78.0",

requirements/build.txt

@@ -4,10 +4,10 @@ ninja
 packaging>=24.2
 setuptools>=77.0.3,<81.0.0
 setuptools-scm>=8
-torch==2.9.1
+torch==2.10.0
 wheel
 jinja2>=3.1.6
 regex
 build
-protobuf
+protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.*
 grpcio-tools==1.78.0 # Required for grpc entrypoints

requirements/common.txt

@@ -9,7 +9,7 @@ blake3
 py-cpuinfo
 transformers >= 4.56.0, < 5
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
-protobuf # Required by LlamaTokenizer, gRPC.
+protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
 fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
 aiohttp >= 3.13.3
 openai >= 1.99.1  # For Responses API with reasoning content

requirements/cuda.txt

@@ -5,9 +5,9 @@ numba == 0.61.2 # Required for N-gram speculative decoding
 
 # Dependencies for NVIDIA GPUs
 ray[cgraph]>=2.48.0
-torch==2.9.1
-torchaudio==2.9.1
+torch==2.10.0
+torchaudio==2.10.0
 # These must be updated alongside torch
-torchvision==0.24.1 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
+torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
 flashinfer-python==0.6.3

requirements/rocm-build.txt

@@ -1,12 +1,11 @@
 # Common dependencies
 -r common.txt
 
---extra-index-url https://download.pytorch.org/whl/rocm6.4
-torch==2.9.1
-torchvision==0.24.1
-torchaudio==2.9.1
-
-triton==3.5.1
+--extra-index-url https://download.pytorch.org/whl/test/rocm7.0
+torch==2.10.0
+torchvision==0.25.0
+torchaudio==2.10.0
+triton==3.6.0
 cmake>=3.26.1,<4
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0

requirements/test.in

@@ -24,10 +24,10 @@ sentence-transformers>=5.2.0 # required for embedding tests
 soundfile # required for audio tests
 jiwer # required for audio tests
 tblib # for pickling test exceptions
-timm==1.0.17 # required for internvl and gemma3n-mm test
-torch==2.9.1
-torchaudio==2.9.1
-torchvision==0.24.1
+timm >=1.0.17 # required for internvl and gemma3n-mm test
+torch==2.10.0
+torchaudio==2.10.0
+torchvision==0.25.0
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
 mistral_common[image,audio] >= 1.9.0 # required for voxtral test

requirements/test.txt

@@ -155,6 +155,10 @@ coverage==7.10.6
     # via pytest-cov
 cramjam==2.9.0
     # via fastparquet
+cuda-bindings==12.9.4
+    # via torch
+cuda-pathfinder==1.3.3
+    # via cuda-bindings
 cupy-cuda12x==13.6.0
     # via ray
 cycler==0.12.1
@@ -631,7 +635,7 @@ nvidia-nvjitlink-cu12==12.9.86
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
     #   torch
-nvidia-nvshmem-cu12==3.3.20
+nvidia-nvshmem-cu12==3.4.5
     # via torch
 nvidia-nvtx-cu12==12.9.79
     # via torch
@@ -1163,7 +1167,7 @@ tomli==2.2.1
     # via schemathesis
 tomli-w==1.2.0
     # via schemathesis
-torch==2.9.1+cu129
+torch==2.10.0+cu129
     # via
     #   -r requirements/test.in
     #   accelerate
@@ -1192,7 +1196,7 @@ torch==2.9.1+cu129
     #   torchvision
     #   vector-quantize-pytorch
     #   vocos
-torchaudio==2.9.1+cu129
+torchaudio==2.10.0+cu129
     # via
     #   -r requirements/test.in
     #   encodec
@@ -1205,7 +1209,7 @@ torchmetrics==1.7.4
     #   pytorch-lightning
     #   terratorch
     #   torchgeo
-torchvision==0.24.1+cu129
+torchvision==0.25.0+cu129
     # via
     #   -r requirements/test.in
     #   lightly
@@ -1247,7 +1251,7 @@ transformers==4.57.5
     #   transformers-stream-generator
 transformers-stream-generator==0.0.5
     # via -r requirements/test.in
-triton==3.5.1
+triton==3.6.0
     # via torch
 tritonclient==2.64.0
     # via -r requirements/test.in

requirements/xpu.txt

@@ -15,4 +15,4 @@ torch==2.10.0+xpu
 torchaudio
 torchvision
 
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.1/vllm_xpu_kernels-0.1.1-cp312-cp312-linux_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.2/vllm_xpu_kernels-0.1.2-cp312-cp312-linux_x86_64.whl

tests/compile/passes/test_functionalization.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import copy
+
 import pytest
 import torch
 
-import vllm.envs as envs
 from tests.compile.backend import TestBackend
 from tests.utils import TestFP8Layer
 from vllm.compilation.passes.fusion.act_quant_fusion import (
@@ -31,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
 
 TEST_FP8 = current_platform.supports_fp8()
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -198,23 +200,82 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
         return [torch.ops.aten.slice_scatter.default]
 
 
-MODELS = [
-    TestSiluMul,
-    TestFusedAddRMSNorm,
-    TestRotaryEmbedding,
-    TestRotaryEmbeddingSliceScatter,
-]
+class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
+    OP_REGISTERED = False
+
+    def __init__(self):
+        super().__init__()
+        self.register_test_custom_op()
+
+    @classmethod
+    def register_test_custom_op(cls):
+        if not cls.OP_REGISTERED:
+
+            def function_with_mutated_args_and_return_impl(
+                x: torch.Tensor,
+            ) -> torch.Tensor:
+                ret = x + 1
+                x.add_(2)
+                return ret
+
+            def function_with_mutated_args_and_return_fake(
+                x: torch.Tensor,
+            ) -> torch.Tensor:
+                return torch.empty_like(x)
+
+            direct_register_custom_op(
+                op_name="function_with_mutated_args_and_return",
+                op_func=function_with_mutated_args_and_return_impl,
+                mutates_args=["x"],
+                fake_impl=function_with_mutated_args_and_return_fake,
+            )
+
+            cls.OP_REGISTERED = True
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Clone x to avoid mutating the original tensor
+        ret = torch.ops.vllm.function_with_mutated_args_and_return(x)
+        return x, ret
+
+    def example_inputs(self, num_tokens=32):
+        hidden_states = torch.randn(num_tokens)
+        return (hidden_states,)
+
+    def ops_in_model(self, do_fusion):
+        return [torch.ops.vllm.function_with_mutated_args_and_return.default]
+
+    def ops_not_in_model(self):
+        return []
+
+
+MODELS_AND_DO_FUSION = {
+    TestSiluMul: [True, False],
+    TestFusedAddRMSNorm: [True, False],
+    TestRotaryEmbedding: [False],
+    TestRotaryEmbeddingSliceScatter: [False],
+    TestFunctionWithMutatedArgsAndReturn: [False],
+}
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("model_class", MODELS)
-@pytest.mark.parametrize("do_fusion", [True, False])
-@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
+@pytest.mark.parametrize(
+    "model_class, do_fusion",
+    [
+        (model_class, do_fusion)
+        for model_class, fusions in MODELS_AND_DO_FUSION.items()
+        for do_fusion in fusions
+    ],
+)
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on cuda and rocm platform",
+)
 def test_fix_functionalization(
     model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype
 ):
     torch.set_default_device("cuda")
     torch.set_default_dtype(dtype)
+    torch.manual_seed(0)
 
     vllm_config = VllmConfig(
         model_config=ModelConfig(dtype=dtype),
@@ -246,8 +307,14 @@ def test_fix_functionalization(
         backend_no_func = TestBackend(*passes)
 
         model = model_class()
-        torch.compile(model, backend=backend_func)(*model.example_inputs())
-        torch.compile(model, backend=backend_no_func)(*model.example_inputs())
+        inputs_func = model.example_inputs()
+        inputs_no_func = copy.deepcopy(inputs_func)
+        model_func = model_class()
+        model_no_func = copy.deepcopy(model_func)
+        model_func = torch.compile(model_func, backend=backend_func)
+        model_no_func = torch.compile(model_no_func, backend=backend_no_func)
+        model_func(*inputs_func)
+        model_no_func(*inputs_no_func)
 
         # check if the functionalization pass is applied
         for op in model.ops_in_model(do_fusion):
@@ -265,3 +332,8 @@ def test_fix_functionalization(
                     found[op] = True
         assert all(found[op] for op in model.ops_in_model(do_fusion))
         assert all(not found.get(op) for op in model.ops_not_in_model())
+
+        # TODO (Rohan138): compare the outputs from model_func and model_no_func
+        # currently runs into errors while comparing `TestFusedAddRMSNorm`
+        # Linked issue: https://github.com/vllm-project/vllm/issues/34996
+        # torch.testing.assert_close(outputs_func, outputs_no_func)

tests/compile/passes/test_fusion_attn.py

@@ -267,7 +267,7 @@ elif current_platform.is_rocm():
     PATTERN_TEST_MODELS_FP8 = [
         ("amd/Llama-3.1-8B-Instruct-FP8-KV", TestAttentionFp8StaticQuantPatternModel)
     ]
-    BACKENDS = [
+    BACKENDS_FP8 = [
         AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
         AttentionBackendEnum.ROCM_ATTN,
         AttentionBackendEnum.TRITON_ATTN,
@@ -474,6 +474,17 @@ def test_attention_quant_pattern(
     assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
         "Attention should not have output_block_scale before fusion"
     )
+
+    kv_cache_dummy_dep_pre_is_none = (
+        attn_nodes_pre[0].kwargs.get("kv_cache_dummy_dep") is None
+    )
+    kv_cache_dummy_dep_post_is_none = (
+        attn_nodes_post[0].kwargs.get("kv_cache_dummy_dep") is None
+    )
+    assert not (kv_cache_dummy_dep_pre_is_none ^ kv_cache_dummy_dep_post_is_none), (
+        "The kv_cache_dummy_dep should be consistent before and after fusion"
+    )
+
     if quant_key.dtype == FP8_DTYPE:
         assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
             "Attention should not have output_block_scale after FP8 fusion"

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+import vllm.config
+from tests.compile.backend import TestBackend
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
+from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
+from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
+from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
+from vllm.compilation.passes.utility.scatter_split_replace import (
+    ScatterSplitReplacementPass,
+)
+from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
+from vllm.config import (
+    CacheConfig,
+    CompilationConfig,
+    CompilationMode,
+    ModelConfig,
+    PassConfig,
+    VllmConfig,
+)
+from vllm.forward_context import get_forward_context, set_forward_context
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+from vllm.platforms import current_platform
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    CommonAttentionMetadata,
+)
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.kv_cache_interface import AttentionSpec
+
+INDEX_SELECT_OP = torch.ops.aten.index.Tensor
+VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
+FP8_DTYPE = current_platform.fp8_dtype()
+
+
+class QKRoPEKVCacheTestModel(torch.nn.Module):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        attn_backend: AttentionBackendEnum,
+        num_heads: int,
+        num_kv_heads: int,
+        head_size: int,
+        is_neox: bool,
+        dtype: torch.dtype,
+        device: torch.device,
+        prefix: str = "model.layers.0.self_attn.attn",
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.block_size = vllm_config.cache_config.block_size
+        self.q_size = num_heads * head_size
+        self.kv_size = num_kv_heads * head_size
+        self.is_neox = is_neox
+        self.dtype = dtype
+        self.device = device
+        self.layer_name = prefix
+
+        self.rotary_emb = RotaryEmbedding(
+            head_size,
+            rotary_dim=head_size,
+            max_position_embeddings=4096,
+            base=10000,
+            is_neox_style=is_neox,
+            dtype=self.dtype,
+        )
+
+        # Whether to check for the RoPE custom op or component index_select
+        self.enable_rope_custom_op = self.rotary_emb.enabled()
+
+        # Register layer metadata for the fusion pass via Attention.
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=1.0 / head_size**0.5,
+            num_kv_heads=num_kv_heads,
+            cache_config=vllm_config.cache_config,
+            quant_config=vllm_config.quant_config,
+            prefix=prefix,
+            attn_backend=attn_backend.get_class(),
+        )
+        self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
+        assert not self.attn_backend.forward_includes_kv_cache_update, (
+            f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache."
+        )
+        self.attn._k_scale = self.attn._k_scale.to(device)
+        self.attn._v_scale = self.attn._v_scale.to(device)
+
+        kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
+        self.kv_cache_dtype = (
+            FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
+        )
+
+        # Initialize attn MetadataBuilder
+        self.builder = self.attn.attn_backend.get_builder_cls()(
+            kv_cache_spec=AttentionSpec(
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_size=head_size,
+                dtype=self.kv_cache_dtype,
+            ),
+            layer_names=[self.attn.layer_name],
+            vllm_config=vllm_config,
+            device=device,
+        )
+
+    def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
+        """Initialize attention metadata."""
+        # Create common attn metadata
+        batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
+        common_attn_metadata = create_common_attn_metadata(
+            batch_spec, self.block_size, self.device, arange_block_indices=True
+        )
+
+        max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
+        num_blocks = batch_size * max_blocks
+
+        # Fetch the attention backend and kv cache shape and stride order
+        attn_backend = self.attn.attn_backend
+        kv_cache_shape = attn_backend.get_kv_cache_shape(
+            num_blocks, self.block_size, self.num_kv_heads, self.head_size
+        )
+        try:
+            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
+        except (AttributeError, NotImplementedError):
+            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
+
+        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+        inv_order = [
+            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
+        ]
+
+        # Create dummy KV cache
+        raw_tensor = torch.zeros(
+            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+            dtype=self.kv_cache_dtype,
+            device=self.device,
+        )
+        raw_tensor = raw_tensor.view(kv_cache_shape)
+        kv_cache = raw_tensor.permute(*inv_order)
+
+        self.attn.kv_cache = [kv_cache]
+
+        # Build attn metadata
+        attn_metadata = self.builder.build(
+            common_prefix_len=0, common_attn_metadata=common_attn_metadata
+        )
+
+        return attn_metadata
+
+    def forward(
+        self, qkv: torch.Tensor, positions: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Create copy so inplace ops do not modify the original tensors
+        qkv = qkv.clone()
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+
+        # Instead of a full forward pass, match only the KV cache update op here
+        q = q.view(-1, self.num_heads, self.head_size)
+        k = k.view(-1, self.num_kv_heads, self.head_size)
+        v = v.view(-1, self.num_kv_heads, self.head_size)
+        kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
+            k, v, self.layer_name
+        )
+        return q, k, v, kv_cache_dummy_dep
+
+    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
+        ops = []
+        if self.enable_rope_custom_op:
+            ops.append(ROTARY_OP)
+        else:
+            ops.append(INDEX_SELECT_OP)
+        ops.append(torch.ops.vllm.unified_kv_cache_update.default)
+        return ops
+
+    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
+        return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
+
+
+@pytest.mark.parametrize(
+    "attn_backend",
+    [
+        AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        AttentionBackendEnum.TRITON_ATTN,
+        AttentionBackendEnum.ROCM_ATTN,
+    ],
+)
+@pytest.mark.parametrize("enable_rope_custom_op", [True])  # [True, False])
+@pytest.mark.parametrize("num_heads", [64])
+@pytest.mark.parametrize("num_kv_heads", [8])
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("block_size", [16])
+@pytest.mark.parametrize("is_neox", [True, False])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@pytest.mark.skipif(
+    not is_aiter_found_and_supported(),
+    reason="Only test on ROCm with AITER installed and supported",
+)
+def test_rope_kvcache_fusion(
+    attn_backend: AttentionBackendEnum,
+    enable_rope_custom_op: bool,
+    num_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    block_size: int,
+    is_neox: bool,
+    dtype: torch.dtype,
+    kv_cache_dtype: str,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(dtype)
+    torch.manual_seed(0)
+
+    custom_ops: list[str] = []
+    if enable_rope_custom_op:
+        custom_ops.append("+rotary_embedding")
+
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=dtype),
+        cache_config=CacheConfig(
+            block_size=block_size,
+            cache_dtype=kv_cache_dtype,
+        ),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=custom_ops,
+            pass_config=PassConfig(
+                fuse_rope_kvcache=True,
+                eliminate_noops=True,
+            ),
+        ),
+    )
+
+    with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
+        m.setenv("VLLM_ROCM_USE_AITER", "1")
+        rocm_aiter_ops.refresh_env_variables()
+
+        model = QKRoPEKVCacheTestModel(
+            vllm_config=vllm_config,
+            attn_backend=attn_backend,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            is_neox=is_neox,
+            dtype=dtype,
+            device=torch.get_default_device(),
+        )
+
+        fusion_pass = RopeKVCacheFusionPass(vllm_config)
+        passes = [
+            NoOpEliminationPass(vllm_config),
+            SplitCoalescingPass(vllm_config),
+            ScatterSplitReplacementPass(vllm_config),
+            fusion_pass,
+            PostCleanupPass(vllm_config),
+        ]
+        backend = TestBackend(*passes)
+
+        T = 5
+
+        qkv = torch.randn(
+            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
+        )
+        pos = torch.arange(T, dtype=torch.long)
+
+        qkv_unfused = qkv.clone()
+        pos_unfused = pos.clone()
+
+        with set_forward_context(None, vllm_config):
+            forward_context = get_forward_context()
+            attn_metadata = model.build_attn_metadata(T)
+            forward_context.slot_mapping = {
+                model.layer_name: attn_metadata.slot_mapping
+            }
+            q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
+            attn_layer = forward_context.no_compile_layers[model.layer_name]
+            kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
+        del dummy
+
+        torch._dynamo.mark_dynamic(qkv, 0)
+        torch._dynamo.mark_dynamic(pos, 0)
+        with set_forward_context(None, vllm_config):
+            model_fused = torch.compile(model, backend=backend)
+            forward_context = get_forward_context()
+            attn_metadata = model_fused.build_attn_metadata(T)
+            forward_context.slot_mapping = {
+                model.layer_name: attn_metadata.slot_mapping
+            }
+            q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
+            attn_layer = forward_context.no_compile_layers[model.layer_name]
+            kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
+        del dummy
+
+        assert fusion_pass.matched_count == 1
+
+        backend.check_before_ops(model.ops_in_model_before())
+        backend.check_after_ops(model.ops_in_model_after())
+
+        if dtype == torch.float16:
+            ATOL, RTOL = (2e-3, 2e-3)
+        else:
+            ATOL, RTOL = (1e-2, 1e-2)
+
+        torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
+        torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
+        torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
+        # Cannot compare fp8_* directly here, cast to model dtype instead
+        torch.testing.assert_close(
+            kv_cache_unfused.view(dtype),
+            kv_cache_fused.view(dtype),
+            atol=ATOL,
+            rtol=RTOL,
+        )

tests/compile/passes/test_scatter_split_replace.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn as nn
+
+import vllm
+from tests.compile.backend import TestBackend
+from vllm.compilation.passes.utility.scatter_split_replace import (
+    ScatterSplitReplacementPass,
+)
+from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
+from vllm.config import CompilationConfig, CompilationMode, VllmConfig
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+
+
+class ScatterSplitReplacementModel(nn.Module):
+    """Model with a rope+getitem+slice_scatter+split_with_sizes sequence."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype: torch.dtype,
+    ):
+        super().__init__()
+        self.q_size = num_heads * head_size
+        self.kv_size = num_kv_heads * head_size
+
+        self.rotary_emb = RotaryEmbedding(
+            head_size,
+            rotary_dim=head_size,
+            max_position_embeddings=4096,
+            base=10000,
+            is_neox_style=True,
+            dtype=dtype,
+        )
+
+    def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
+        # Create copy so inplace ops do not modify the original tensors
+        qkv = qkv.clone()
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        q = q + 1
+        k = k + 2
+        v = v + 3
+        return q, k, v
+
+    def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
+        return [
+            torch.ops.aten.slice_scatter.default,
+            torch.ops.aten.split_with_sizes.default,
+            torch.ops.aten.getitem.default,
+        ]
+
+    def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
+        return [torch.ops.aten.getitem.default]
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_scatter_split_replace(dtype):
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(dtype)
+    torch.manual_seed(0)
+
+    num_heads = 8
+    num_kv_heads = 4
+    head_size = 64
+
+    vllm_config = VllmConfig(
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["+rotary_embedding"],
+        ),
+    )
+    with vllm.config.set_current_vllm_config(vllm_config):
+        # ScatterSplitReplacementPass requires SplitCoalescingPass to be run before it
+        coalesce_pass = SplitCoalescingPass(vllm_config)
+        replace_pass = ScatterSplitReplacementPass(vllm_config)
+        passes = [coalesce_pass, replace_pass]
+        backend = TestBackend(*passes)
+
+        model = ScatterSplitReplacementModel(num_heads, num_kv_heads, head_size, dtype)
+
+        T = 5
+        qkv = torch.randn(
+            T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
+        )
+        pos = torch.arange(T, dtype=torch.long)
+
+        qkv_eager = qkv.clone()
+        pos_eager = pos.clone()
+        result_eager = model(qkv_eager, pos_eager)
+
+        torch._dynamo.mark_dynamic(qkv, 0)
+        torch._dynamo.mark_dynamic(pos, 0)
+
+        model_compiled = torch.compile(model, backend=backend)
+        result_compiled = model_compiled(qkv, pos)
+
+        for eager, compiled in zip(result_eager, result_compiled):
+            torch.testing.assert_close(eager, compiled)
+
+        assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
+        assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1

tests/compile/test_aot_compile.py

@@ -90,9 +90,7 @@ def use_vllm_config(vllm_config: VllmConfig):
         yield
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
     with monkeypatch.context() as m:
         vllm_config = make_vllm_config()
@@ -116,9 +114,7 @@ def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
             assert torch.allclose(actual, expected)
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_force_aot_load(monkeypatch: pytest.MonkeyPatch):
     with tempfile.TemporaryDirectory() as tmpdirname, monkeypatch.context() as m:
         args = (torch.randn(10, 10),)
@@ -132,9 +128,7 @@ def test_force_aot_load(monkeypatch: pytest.MonkeyPatch):
             CompiledMod(vllm_config=vllm_config)(*args)
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
     with monkeypatch.context() as m:
         args = (torch.randn(10, 10),)
@@ -162,9 +156,7 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
             assert torch.allclose(ret, expected)
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
     """
     Test that cache loading correctly handles the returns_tuple logic.
@@ -223,9 +215,7 @@ def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
             )
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_cache_load_returns_tuple_consistency_tuple_output(
     monkeypatch: pytest.MonkeyPatch,
 ):
@@ -294,9 +284,7 @@ def test_cache_load_returns_tuple_consistency_tuple_output(
             )
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_shape_env(monkeypatch: pytest.MonkeyPatch):
     """
     Test that the shape environment is correctly serialized and preserved
@@ -333,9 +321,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
                 assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_partition_wrapper_applied_on_aot_load(
     monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker
 ):
@@ -426,9 +412,7 @@ def test_partition_wrapper_applied_on_aot_load(
         )
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 @create_new_process_for_each_test("spawn")
 def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
     """
@@ -492,9 +476,7 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
         symbolic_shapes_module.make_symbol = original_make_symbol
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 class TestStandaloneCompiledArtifacts:
     def test_init(self):
         cache = StandaloneCompiledArtifacts()
@@ -668,9 +650,7 @@ class TestStandaloneCompiledArtifacts:
         assert len(restored_cache.loaded_submodule_store) == 0
 
 
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 class TestStandaloneCompiledArtifactsIntegration:
     def test_add_pickle_unpickle(self):
         cache = StandaloneCompiledArtifacts()

tests/compile/test_dynamic_shapes_compilation.py

@@ -39,9 +39,7 @@ def get_test_models():
 @pytest.mark.parametrize("use_aot_compile", ["0", "1"])
 @pytest.mark.parametrize("use_bytecode_hook", [True, False])
 @pytest.mark.parametrize("evaluate_guards", [False, True])
-@pytest.mark.skipif(
-    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
-)
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
 def test_dynamic_shapes_compilation(
     monkeypatch,
     model_name,

tests/config/test_config_generation.py

@@ -78,3 +78,27 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
     )
 
     ray.shutdown()
+
+
+def test_unrecognized_env():
+    import os
+
+    # Test that if fail_on_environ_validation is True, then an error
+    # is raised when an unrecognized vLLM environment variable is set
+    os.environ["VLLM_UNRECOGNIZED_ENV_VAR"] = "some_value"
+    engine_args = EngineArgs(
+        fail_on_environ_validation=True,
+    )
+    with pytest.raises(ValueError, match="Unknown vLLM environment variable detected"):
+        engine_args.create_engine_config()
+
+    # Test that if fail_on_environ_validation is False, then no error is raised
+    engine_args = EngineArgs()
+    engine_args.create_engine_config()
+
+    # Test that when the unrecognized env var is removed, no error is raised
+    os.environ.pop("VLLM_UNRECOGNIZED_ENV_VAR", None)
+    engine_args = EngineArgs(
+        fail_on_environ_validation=True,
+    )
+    engine_args.create_engine_config()

tests/entrypoints/pooling/classify/test_offline.py

@@ -7,12 +7,15 @@ import pytest
 import torch
 
 from tests.models.utils import softmax
-from vllm import LLM, PoolingParams
+from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.tasks import PoolingTask
 
 MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
 
-prompts = ["The chef prepared a delicious meal."]
+prompt = "The chef prepared a delicious meal."
+prompt_token_ids = [785, 29706, 10030, 264, 17923, 15145, 13]
+num_labels = 2
 
 
 @pytest.fixture(scope="module")
@@ -35,11 +38,48 @@ def llm():
     cleanup_dist_env_and_memory()
 
 
+@pytest.mark.skip_global_cleanup
+def test_str_prompts(llm: LLM):
+    outputs = llm.classify(prompt, use_tqdm=False)
+    assert len(outputs) == 1
+    assert isinstance(outputs[0], ClassificationRequestOutput)
+    assert outputs[0].prompt_token_ids == prompt_token_ids
+    assert len(outputs[0].outputs.probs) == num_labels
+
+
+@pytest.mark.skip_global_cleanup
+def test_token_ids_prompts(llm: LLM):
+    outputs = llm.classify([prompt_token_ids], use_tqdm=False)
+    assert len(outputs) == 1
+    assert isinstance(outputs[0], ClassificationRequestOutput)
+    assert outputs[0].prompt_token_ids == prompt_token_ids
+    assert len(outputs[0].outputs.probs) == num_labels
+
+
+@pytest.mark.skip_global_cleanup
+def test_list_prompts(llm: LLM):
+    outputs = llm.classify([prompt, prompt_token_ids], use_tqdm=False)
+    assert len(outputs) == 2
+    for i in range(len(outputs)):
+        assert isinstance(outputs[i], ClassificationRequestOutput)
+        assert outputs[i].prompt_token_ids == prompt_token_ids
+        assert len(outputs[i].outputs.probs) == num_labels
+
+
+@pytest.mark.skip_global_cleanup
+def test_token_classify(llm: LLM):
+    outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
+    assert len(outputs) == 1
+    assert isinstance(outputs[0], PoolingRequestOutput)
+    assert outputs[0].prompt_token_ids == prompt_token_ids
+    assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
+
+
 @pytest.mark.skip_global_cleanup
 def test_pooling_params(llm: LLM):
     def get_outputs(use_activation):
         outputs = llm.classify(
-            prompts,
+            prompt,
             pooling_params=PoolingParams(use_activation=use_activation),
             use_tqdm=False,
         )
@@ -61,11 +101,14 @@ def test_pooling_params(llm: LLM):
 
 
 @pytest.mark.skip_global_cleanup
-def test_token_classify(llm: LLM):
-    llm.encode(prompts, pooling_task="token_classify", use_tqdm=False)
-
-
 def test_score_api(llm: LLM):
     err_msg = "Score API is only enabled for num_labels == 1."
     with pytest.raises(ValueError, match=err_msg):
         llm.score("ping", "pong", use_tqdm=False)
+
+
+@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
+def test_unsupported_tasks(llm: LLM, task: PoolingTask):
+    err_msg = f"Unsupported task: '{task}' Supported tasks.+"
+    with pytest.raises(ValueError, match=err_msg):
+        llm.encode(prompt, pooling_task=task, use_tqdm=False)

tests/entrypoints/pooling/embed/test_online_vision.py

@@ -10,12 +10,12 @@ from transformers import AutoProcessor
 from tests.utils import VLLM_PATH, RemoteOpenAIServer
 from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
 from vllm.multimodal.media import MediaWithBytes
-from vllm.multimodal.utils import fetch_image
+from vllm.multimodal.utils import encode_image_url, fetch_image
 
 MODEL_NAME = "TIGER-Lab/VLM2Vec-Full"
 MAXIMUM_IMAGES = 2
 
-vlm2vec_jinja_path = VLLM_PATH / "examples/template_vlm2vec_phi3v.jinja"
+vlm2vec_jinja_path = VLLM_PATH / "examples/pooling/embed/template/vlm2vec_phi3v.jinja"
 assert vlm2vec_jinja_path.exists()
 
 # Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
@@ -26,6 +26,10 @@ TEST_IMAGE_ASSETS = [
     "RGBA_comp.png",  # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
 ]
 
+input_text = "The best thing about vLLM is that it supports many different models"
+image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
+image_base64 = {"url": encode_image_url(fetch_image(image_url))}
+
 
 @pytest.fixture(scope="module")
 def server():
@@ -48,6 +52,81 @@ def server():
         yield remote_server
 
 
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_text_request(server: RemoteOpenAIServer, model_name: str):
+    messages = [
+        {
+            "role": "user",
+            "content": input_text,
+        },
+    ]
+
+    # note: vlm2vec_phi3v.jinja
+    # Embedding models should only embed one message at a time.
+
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={"model": model_name, "messages": messages},
+    )
+    response.raise_for_status()
+
+    output = EmbeddingResponse.model_validate(response.json())
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert len(output.data[0].embedding) == 3072
+    assert output.usage.prompt_tokens == 14
+
+
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_image_url_request(server: RemoteOpenAIServer, model_name: str):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Represent the user's input."},
+                {"type": "image_url", "image_url": {"url": image_url}},
+            ],
+        }
+    ]
+
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={"model": model_name, "messages": messages},
+    )
+    response.raise_for_status()
+
+    output = EmbeddingResponse.model_validate(response.json())
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert len(output.data[0].embedding) == 3072
+    assert output.usage.prompt_tokens == 767
+
+
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Represent the user's input."},
+                {"type": "image_url", "image_url": image_base64},
+            ],
+        }
+    ]
+
+    response = requests.post(
+        server.url_for("v1/embeddings"),
+        json={"model": model_name, "messages": messages},
+    )
+    response.raise_for_status()
+
+    output = EmbeddingResponse.model_validate(response.json())
+    assert len(output.data) == 1
+    assert output.model == MODEL_NAME
+    assert len(output.data[0].embedding) == 3072
+    assert output.usage.prompt_tokens == 767
+
+
 def get_hf_prompt_tokens(model_name, content, image_url):
     processor = AutoProcessor.from_pretrained(
         model_name, trust_remote_code=True, num_crops=4

tests/kernels/moe/test_gpt_oss_triton_kernels.py

@@ -22,7 +22,7 @@ from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
 from triton_kernels.tensor_details import layout
 from triton_kernels.testing import assert_close
 
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
     triton_kernel_moe_forward,
 )
@@ -298,12 +298,18 @@ def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):
         pc2,
     ) = init_compute_data(M, K, N, E, a_dtype, w_dtype, num_warps=8)
 
-    quant_config = FusedMoEQuantConfig.make(
-        w1_bias=w1_bias_tri,
-        w2_bias=w2_bias_tri,
-        w1_scale=pc1,
-        w2_scale=pc2,
-    )
+    if a_dtype == "bf16" and w_dtype == "mx4":
+        quant_config = mxfp4_w4a16_moe_quant_config(
+            w1_scale=pc1,
+            w2_scale=pc2,
+            w1_bias=w1_bias_tri,
+            w2_bias=w2_bias_tri,
+        )
+    else:
+        raise NotImplementedError(
+            f"Quantization configuration for activation={a_dtype} and weight={w_dtype} "
+            f"has not been implemented."
+        )
 
     out_triton_monolithic = triton_kernel_moe_forward(
         hidden_states=x_tri,

tests/kernels/moe/test_shared_fused_moe_routed_transform.py

@@ -14,6 +14,7 @@ import torch.nn as nn
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
+from vllm.utils.torch_utils import is_torch_equal_or_newer
 
 
 class SimpleLinear(nn.Module):
@@ -60,6 +61,10 @@ def setup_cuda():
 @pytest.mark.parametrize("num_tokens", [1, 32])
 @pytest.mark.parametrize("hidden_size,latent_size", [(256, 128), (128, 64)])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.skipif(
+    is_torch_equal_or_newer("2.10.0"),
+    reason="Test fails with PyTorch 2.10.0 see: https://github.com/vllm-project/vllm/issues/33995",
+)
 def test_routed_input_transform_inside_vs_outside(
     num_tokens: int,
     hidden_size: int,

tests/models/quantization/test_gpt_oss.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end accuracy test for GPT-OSS model quantization.
+
+Config:
+    Task:   gsm8k_platinum
+    Filter: flexible-extract
+    n-shot: 5
+    Metric: exact_match
+
+Run: pytest tests/models/quantization/test_gpt_oss.py
+"""
+
+import importlib
+import importlib.metadata
+from dataclasses import dataclass
+
+import huggingface_hub
+import lm_eval
+import pytest
+from packaging import version
+
+MODEL_ACCURACIES = {
+    # Full quantization: attention linears and MoE linears
+    "amd/gpt-oss-20b-WFP8-AFP8-KVFP8": 0.89,
+    # MoE linears only quantization
+    "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8": 0.89,
+    # MoE linears only quantization
+    # "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-MXFP4-KV-FP8": 0.90,
+}
+
+QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse(
+    importlib.metadata.version("amd-quark")
+) >= version.parse("0.9.0")
+
+
+def has_huggingface_access(repo):
+    try:
+        huggingface_hub.list_repo_refs(repo)
+        return True
+    except huggingface_hub.errors.RepositoryNotFoundError:
+        return False
+
+
+HF_HUB_AMD_ORG_ACCESS = all(
+    [has_huggingface_access(model_name) for model_name in MODEL_ACCURACIES]
+)
+
+
+@dataclass
+class ModelCase:
+    model_id: str
+    tp: int
+
+
+@dataclass
+class EvaluationConfig:
+    model_name: str
+
+    def get_model_args(self, tp_size: int):
+        return {
+            "pretrained": self.model_name,
+            "chat_template_args": {"reasoning_effort": "low"},
+            "enable_thinking": True,
+            "think_end_token": "200008",
+            "tensor_parallel_size": tp_size,
+            "dtype": "auto",
+            "gpu_memory_utilization": 0.95,
+            "trust_remote_code": False,
+            "enable_prefix_caching": False,
+            "enforce_eager": False,
+        }
+
+
+@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
+@pytest.mark.skipif(
+    not HF_HUB_AMD_ORG_ACCESS,
+    reason="Read access to huggingface.co/amd is required for this test.",
+)
+@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
+@pytest.mark.parametrize("model_name, expected_accuracy", MODEL_ACCURACIES.items())
+def test_gpt_oss_attention_quantization(
+    model_name: str, tp_size: int, expected_accuracy: float
+):
+    model_args = EvaluationConfig(model_name).get_model_args(tp_size)
+
+    extra_run_kwargs = {
+        "gen_kwargs": {"max_gen_toks": 8000},
+        "apply_chat_template": True,
+        "fewshot_as_multiturn": True,
+        "num_fewshot": 5,
+    }
+
+    lm_eval_out = lm_eval.simple_evaluate(
+        model="vllm",
+        model_args=model_args,
+        tasks="gsm8k_platinum",
+        batch_size="auto",
+        **extra_run_kwargs,
+    )
+    measured_accuracy = float(
+        lm_eval_out["results"]["gsm8k_platinum"]["exact_match,flexible-extract"]
+    )
+
+    rtol = 0.02
+    assert (
+        measured_accuracy - rtol < expected_accuracy
+        and measured_accuracy + rtol > expected_accuracy
+    ), f"Expected: {expected_accuracy} |  Measured: {measured_accuracy}"

tests/models/quantization/test_gpt_oss_attn_quantization.py

@@ -1,80 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test attention quantization of gpt-oss model.
-The qkv_proj and o_proj in self_attention can be either quantized or excluded.
-
-Run `pytest tests/models/quantization/test_gpt_oss_attn_quantization.py`.
-
-"""
-
-import importlib
-import importlib.metadata
-from dataclasses import dataclass
-
-import huggingface_hub
-import lm_eval
-import pytest
-from packaging import version
-
-MODEL_NAMES = ["amd/gpt-oss-20b-customized-attention-quantization"]
-
-QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse(
-    importlib.metadata.version("amd-quark")
-) >= version.parse("0.8.99")
-
-
-def has_huggingface_access(repo):
-    try:
-        huggingface_hub.list_repo_refs(repo)
-        return True
-    except huggingface_hub.errors.RepositoryNotFoundError:
-        return False
-
-
-HF_HUB_AMD_ORG_ACCESS = all(
-    [has_huggingface_access(model_name) for model_name in MODEL_NAMES]
-)
-
-
-@dataclass
-class ModelCase:
-    model_id: str
-    tp: int
-
-
-@dataclass
-class EvaluationConfig:
-    model_name: str
-
-    def get_model_args(self) -> str:
-        return (
-            f"pretrained={self.model_name},"
-            "tensor_parallel_size=4,dtype=auto,gpu_memory_utilization=0.9,trust_remote_code=False"
-        )
-
-
-EXPECTED_ACCURACIES = {"arc_challenge": 0.20}
-
-
-@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
-@pytest.mark.skipif(
-    not HF_HUB_AMD_ORG_ACCESS,
-    reason="Read access to huggingface.co/amd is required for this test.",
-)
-@pytest.mark.parametrize("model_name", MODEL_NAMES)
-@pytest.mark.parametrize("task_name, expected_accuracy", EXPECTED_ACCURACIES.items())
-def test_gpt_oss_attention_quantization(
-    model_name: str, task_name: str, expected_accuracy: float
-):
-    measured_accuracy = lm_eval.simple_evaluate(
-        model="vllm",
-        model_args=EvaluationConfig(model_name).get_model_args(),
-        tasks=task_name,
-        batch_size="auto",
-    )["results"][task_name]["acc,none"]
-
-    rtol = 0.05
-    assert (
-        measured_accuracy - rtol < expected_accuracy
-        and measured_accuracy + rtol > expected_accuracy
-    ), f"Expected: {expected_accuracy} |  Measured: {measured_accuracy}"

tests/quantization/test_cpu_wna16.py

@@ -17,7 +17,7 @@ DTYPE = ["bfloat16"]
 
 @pytest.mark.parametrize("model", MODELS)
 @pytest.mark.parametrize("dtype", DTYPE)
-def test_ipex_quant(vllm_runner, model, dtype):
+def test_cpu_quant(vllm_runner, model, dtype):
     with vllm_runner(model, dtype=dtype) as llm:
         output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
     assert output

tests/quantization/test_ipex_quant.py

@@ -1,32 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Test model set-up and inference for quantized HF models supported
-on the CPU/GPU backend using IPEX (including AWQ/GPTQ).
-
-Validating the configuration and printing results for manual checking.
-
-Run `pytest tests/quantization/test_ipex_quant.py`.
-"""
-
-import pytest
-
-from vllm.platforms import current_platform
-
-MODELS = [
-    "AMead10/Llama-3.2-1B-Instruct-AWQ",
-    "shuyuej/Llama-3.2-1B-Instruct-GPTQ",  # with g_idx
-]
-DTYPE = ["bfloat16"]
-
-
-@pytest.mark.skipif(
-    not current_platform.is_cpu() and not current_platform.is_xpu(),
-    reason="only supports Intel CPU/XPU backend.",
-)
-@pytest.mark.parametrize("model", MODELS)
-@pytest.mark.parametrize("dtype", DTYPE)
-def test_ipex_quant(vllm_runner, model, dtype):
-    with vllm_runner(model, dtype=dtype, enforce_eager=True) as llm:
-        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
-    assert output
-    print(output)

tests/renderers/test_hf.py

@@ -428,13 +428,13 @@ def test_resolve_content_format_fallbacks(model, expected_format):
         ("template_chatglm.jinja", "string"),
         ("template_chatglm2.jinja", "string"),
         ("template_chatml.jinja", "string"),
-        ("template_dse_qwen2_vl.jinja", "openai"),
         ("template_falcon_180b.jinja", "string"),
         ("template_falcon.jinja", "string"),
         ("template_inkbot.jinja", "string"),
         ("template_teleflm.jinja", "string"),
-        ("template_vlm2vec_phi3v.jinja", "openai"),
-        ("template_vlm2vec_qwen2vl.jinja", "openai"),
+        ("pooling/embed/template/dse_qwen2_vl.jinja", "openai"),
+        ("pooling/embed/template/vlm2vec_phi3v.jinja", "openai"),
+        ("pooling/embed/template/vlm2vec_qwen2vl.jinja", "openai"),
         ("tool_chat_template_granite_20b_fc.jinja", "string"),
         ("tool_chat_template_hermes.jinja", "string"),
         ("tool_chat_template_internlm2_tool.jinja", "string"),

tests/v1/attention/test_attention_backends.py

@@ -179,7 +179,7 @@ def create_and_prepopulate_kv_cache(
         block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
         start_block_idx += num_blocks_for_seq
 
-        # Create a realistic slot mapping that corresponds to the block table
+    # Create a realistic slot mapping that corresponds to the block table
     for i in range(batch_size):
         token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
         block_indices = token_offsets // block_size

tests/v1/core/test_async_scheduler.py

@@ -236,7 +236,7 @@ def test_prefix_caching_for_multi_turn():
         req._all_token_ids = req.prompt_token_ids.copy()
         req.all_token_ids = ConstantList(req._all_token_ids)
         req.block_hashes = []
-        req.block_hashes = req.get_hash_new_full_blocks()
+        req.update_block_hashes()
 
     # Schedule the next-turn requests.
     for req in next_turn_requests:

vllm/_aiter_ops.py

@@ -87,6 +87,10 @@ def _rocm_aiter_fused_moe_impl(
     a2_scale: torch.Tensor | None = None,
     num_local_tokens: torch.Tensor | None = None,
     output_dtype: torch.dtype | None = None,
+    hidden_pad: int = 0,
+    intermediate_pad: int = 0,
+    bias1: torch.Tensor | None = None,
+    bias2: torch.Tensor | None = None,
 ) -> torch.Tensor:
     from aiter import ActivationType, QuantType
     from aiter.fused_moe import fused_moe
@@ -110,6 +114,10 @@ def _rocm_aiter_fused_moe_impl(
         a2_scale,
         num_local_tokens=num_local_tokens,
         dtype=output_dtype,
+        hidden_pad=hidden_pad,
+        intermediate_pad=intermediate_pad,
+        bias1=bias1,
+        bias2=bias2,
     )
 
 
@@ -307,6 +315,28 @@ def _rocm_aiter_grouped_topk_fake(
     pass
 
 
+def _rocm_aiter_fused_topk_impl(
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    top_k: int,
+    gate_up: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    from aiter.fused_moe import fused_topk
+
+    # fused_topk returns (topk_weights, topk_indices)
+    return fused_topk(x, router_logits, top_k, gate_up)
+
+
+def _rocm_aiter_fused_topk_fake(
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    top_k: int,
+    gate_up: bool,
+) -> None:
+    # tuple[torch.Tensor, torch.Tensor]:
+    pass
+
+
 # Cache whether aiter supports FP8 MLA parameters
 _AITER_MLA_SUPPORTS_FP8: bool | None = None
 
@@ -941,6 +971,70 @@ class rocm_aiter_ops:
         cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
         cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
 
+    @staticmethod
+    def get_aiter_activation_type(activation_str: str):
+        """
+        Given an activation type as a string, returns the corresponding aiter ActivationType enum.
+        Supported activation types: "no", "none", "silu", "gelu", "swiglu".
+        Returns None if the mapping fails.
+
+        Args:
+            activation_str (str): Activation type as string.
+
+        Returns:
+            Aiter ActivationType enum value, or None if not found.
+        """
+        # Import only locally, since aiter may not always be available.
+        try:
+            from aiter import ActivationType
+        except ImportError:
+            return None
+
+        if not isinstance(activation_str, str):
+            return None
+
+        name = activation_str.strip().lower()
+        mapping = {
+            "none": ActivationType.No,
+            "no": ActivationType.No,
+            "silu": ActivationType.Silu,
+            "gelu": ActivationType.Gelu,
+            "swiglu": ActivationType.Swiglu,
+        }
+        return mapping.get(name)
+
+    @staticmethod
+    def get_aiter_quant_type(quant_type_str: str):
+        """
+        Given a quantization type as a string, returns the corresponding aiter QuantType enum.
+        Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128".
+        Returns None if the mapping fails.
+
+        Args:
+            quant_type_str (str): Quantization type as string.
+
+        Returns:
+            Aiter QuantType enum value, or None if not found.
+        """
+        try:
+            from aiter import QuantType
+        except ImportError:
+            return None
+
+        if not isinstance(quant_type_str, str):
+            return None
+
+        name = quant_type_str.strip().lower()
+        mapping = {
+            "no": QuantType.No,
+            "per_tensor": QuantType.per_Tensor,
+            "per_token": QuantType.per_Token,
+            "per_1x32": QuantType.per_1x32,
+            "per_1x128": QuantType.per_1x128,
+            "per_128x128": QuantType.per_128x128,
+        }
+        return mapping.get(name)
+
     @classmethod
     @if_aiter_supported
     def is_enabled(cls) -> bool:
@@ -1070,6 +1164,14 @@ class rocm_aiter_ops:
                 dispatch_key=current_platform.dispatch_key,
             )
 
+            direct_register_custom_op(
+                op_name="rocm_aiter_fused_topk",
+                op_func=_rocm_aiter_fused_topk_impl,
+                mutates_args=[],
+                fake_impl=_rocm_aiter_fused_topk_fake,
+                dispatch_key=current_platform.dispatch_key,
+            )
+
             direct_register_custom_op(
                 op_name="rocm_aiter_mla_decode_fwd",
                 op_func=_rocm_aiter_mla_decode_fwd_impl,
@@ -1291,6 +1393,10 @@ class rocm_aiter_ops:
         a2_scale: torch.Tensor | None = None,
         num_local_tokens: torch.Tensor | None = None,
         output_dtype: torch.dtype | None = None,
+        hidden_pad: int = 0,
+        intermediate_pad: int = 0,
+        bias1: torch.Tensor | None = None,
+        bias2: torch.Tensor | None = None,
     ) -> torch.Tensor:
         return torch.ops.vllm.rocm_aiter_fused_moe(
             hidden_states,
@@ -1308,6 +1414,10 @@ class rocm_aiter_ops:
             a2_scale,
             num_local_tokens,
             output_dtype,
+            hidden_pad,
+            intermediate_pad,
+            bias1,
+            bias2,
         )
 
     @staticmethod
@@ -1412,6 +1522,15 @@ class rocm_aiter_ops:
             routed_scaling_factor,
         )
 
+    @staticmethod
+    def fused_topk(
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        top_k: int,
+        gate_up: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
+
     @staticmethod
     def mla_decode_fwd(
         q: torch.Tensor,
@@ -1518,6 +1637,45 @@ class rocm_aiter_ops:
         query = query.view(query_shape)
         key = key.view(key_shape)
 
+    @staticmethod
+    def triton_rope_and_cache(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        positions: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        key_cache: torch.Tensor,
+        value_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+        k_scale: torch.Tensor,
+        v_scale: torch.Tensor,
+        flash_layout: bool,
+        apply_scale: bool,
+    ):
+        from aiter.ops.triton.fused_kv_cache import fused_qk_rope_reshape_and_cache
+
+        cos, sin = cos_sin_cache.chunk(2, dim=-1)
+        fused_qk_rope_reshape_and_cache(
+            query,
+            key,
+            value,
+            key_cache,
+            value_cache,
+            layer_slot_mapping,
+            positions,
+            cos,
+            sin,
+            k_scale,
+            v_scale,
+            is_neox,
+            flash_layout=flash_layout,
+            apply_scale=apply_scale,
+            q_out=query,
+            k_out=key,
+            output_zeros=False,
+        )
+
     @staticmethod
     def batched_gemm_a16wfp4(
         X: torch.Tensor,
@@ -1629,6 +1787,47 @@ class rocm_aiter_ops:
 
         return shuffle_weight(tensor, layout=layout)
 
+    @staticmethod
+    def shuffle_weight_a16w4(
+        tensor: "torch.Tensor",
+        nLane: int,
+        gate_up: bool,
+    ) -> "torch.Tensor":
+        """
+        Shuffles the weight tensor into (A16W4) layout for AITER kernels.
+
+        Args:
+            tensor: The input weight tensor to be shuffled.
+            layout: The block layout to use, defaults to (16, 4).
+
+        Returns:
+            torch.Tensor: The shuffled tensor.
+        """
+        from aiter.ops.shuffle import shuffle_weight_a16w4
+
+        return shuffle_weight_a16w4(tensor, nLane, gate_up)
+
+    @staticmethod
+    def shuffle_scale_a16w4(
+        tensor: "torch.Tensor",
+        num_experts: int,
+        gate_up: bool,
+    ) -> "torch.Tensor":
+        """
+        Shuffles the scale tensor into (A16W4) layout for AITER kernels.
+
+        Args:
+            tensor: The input scale tensor to be shuffled.
+            num_experts: Number of experts, needed for reshaping logic.
+            gate_up: Whether the scale is for w13 (True) or w2 (False).
+
+        Returns:
+            torch.Tensor: The shuffled scale tensor.
+        """
+        from aiter.ops.shuffle import shuffle_scale_a16w4
+
+        return shuffle_scale_a16w4(tensor, num_experts, gate_up)
+
     @staticmethod
     def shuffle_weights(
         *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)

vllm/_xpu_ops.py

@@ -53,7 +53,7 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
         return torch.empty((M, N), dtype=input.dtype, device=input.device)
 
 
-class ipex_ops:
+class xpu_ops:
     @staticmethod
     def flash_attn_varlen_func(
         q: torch.Tensor,
@@ -73,7 +73,7 @@ class ipex_ops:
         cu_seqlens_k: torch.Tensor | None = None,
         # passed in qwen vl
         dropout_p: float = 0.0,
-        # The following parameters are not used in ipex kernel currently,
+        # The following parameters are not used in xpu kernel currently,
         # we keep API compatible to CUDA's.
         scheduler_metadata=None,
         fa_version: int = 2,
@@ -153,6 +153,6 @@ class ipex_ops:
         sm_margin=0,  # Can be tuned if some SMs are used for communication
     ) -> None:
         logger.warning_once(
-            "get_scheduler_metadata is not implemented for ipex_ops, returning None."
+            "get_scheduler_metadata is not implemented for xpu_ops, returning None."
         )
         return None

vllm/compilation/compiler_interface.py

@@ -233,7 +233,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
 
         from torch._inductor import standalone_compile
 
-        supports_aot = is_torch_equal_or_newer("2.10.0.dev")
+        supports_aot = is_torch_equal_or_newer("2.10.0")
 
         if not supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT:
             logger.error(

vllm/compilation/decorators.py

@@ -333,7 +333,7 @@ def _support_torch_compile(
     ) -> None:
         def mark_dynamic(arg: torch.Tensor, dims: list[int]) -> None:
             if ds_type == DynamicShapesType.UNBACKED:
-                if is_torch_equal_or_newer("2.10.0.dev"):
+                if is_torch_equal_or_newer("2.10.0"):
                     for dim in dims:
                         torch._dynamo.decorators.mark_unbacked(
                             arg, dim, hint_override=arg.size()[dim]
@@ -373,7 +373,7 @@ def _support_torch_compile(
                     if isinstance(arg, torch.Tensor):
                         # In case dims is specified with negative indexing
                         dims = [arg.ndim + dim if dim < 0 else dim for dim in dims]
-                        if is_torch_equal_or_newer("2.10.0.dev"):
+                        if is_torch_equal_or_newer("2.10.0"):
                             for dim in dims:
                                 torch._dynamo.decorators.mark_unbacked(
                                     arg, dim, hint_override=arg.size()[dim]
@@ -525,9 +525,9 @@ def _support_torch_compile(
             fx_config_patches["backed_size_oblivious"] = True
 
         # Prepare inductor config patches
-        # assume_32bit_indexing is only available in torch 2.10.0.dev+
+        # assume_32bit_indexing is only available in torch 2.10.0+
         inductor_config_patches = {}
-        if is_torch_equal_or_newer("2.10.0.dev"):
+        if is_torch_equal_or_newer("2.10.0"):
             inductor_config_patches["assume_32bit_indexing"] = (
                 self.compilation_config.dynamic_shapes_config.assume_32_bit_indexing
             )

vllm/compilation/passes/fusion/attn_quant_fusion.py

@@ -142,6 +142,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
             v: torch.Tensor,
             output_attn: torch.Tensor,
             scale: torch.Tensor,
+            kv_cache_dummy_dep: torch.Tensor,
         ) -> torch.Tensor:
             at1 = auto_functionalized(
                 ATTN_OP,
@@ -152,6 +153,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
                 layer_name=self.layer_name,
                 output_scale=None,
                 output_block_scale=None,
+                kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
             attn_out_view = RESHAPE_OP(
                 at1[1], [q.shape[0], self.num_heads * self.head_size]
@@ -165,6 +167,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
             v: torch.Tensor,
             output_attn: torch.Tensor,
             scale: torch.Tensor,
+            kv_cache_dummy_dep: torch.Tensor,
         ) -> torch.Tensor:
             # attn output in quant_dtype
             output_attn = torch.ops.aten.full.default(
@@ -182,6 +185,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
                 layer_name=self.layer_name,
                 output_scale=scale,
                 output_block_scale=None,
+                kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
             return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
 
@@ -191,6 +195,7 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
             self.empty(5, self.num_heads, self.head_size),  # v
             self.empty(5, self.num_heads, self.head_size),  # attn_output
             empty_fp32(1, 1),  # scale
+            self.empty(0),  # kv_cache_dummy_dep
         ]
 
         pm.register_replacement(
@@ -228,6 +233,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
             output_quant: torch.Tensor,
             output_scale: torch.Tensor,
             input_scale: torch.Tensor,
+            kv_cache_dummy_dep: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             at1 = auto_functionalized(
                 ATTN_OP,
@@ -238,6 +244,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
                 layer_name=self.layer_name,
                 output_scale=None,
                 output_block_scale=None,
+                kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
             attn_out_view = RESHAPE_OP(
                 at1[1], [q.shape[0], self.num_heads * self.head_size]
@@ -261,6 +268,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
             output_quant: torch.Tensor,
             output_scale: torch.Tensor,
             input_scale: torch.Tensor,
+            kv_cache_dummy_dep: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor]:
             # attention output in quant_dtype
             output_attn = torch.ops.aten.full.default(
@@ -280,6 +288,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
                 layer_name=self.layer_name,
                 output_scale=input_scale,
                 output_block_scale=output_scale_view,
+                kv_cache_dummy_dep=kv_cache_dummy_dep,
             )
             output = RESHAPE_OP(at2[1], [-1, self.num_heads * self.head_size // 2])
             return output, at2[2]
@@ -294,6 +303,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
                 128, round_up(self.num_heads * self.head_size // 16, 4)
             ),  # output_scale
             empty_fp32(1, 1),  # input_scale
+            self.empty(0),  # kv_cache_dummy_dep
         ]
 
         pm.register_replacement(

vllm/compilation/passes/fusion/rope_kvcache_fusion.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch._inductor.pattern_matcher as pm
+from torch import fx
+from torch._higher_order_ops import auto_functionalized
+from torch._inductor.fx_passes.post_grad import view_to_reshape
+from torch._inductor.pattern_matcher import PatternMatcherPass
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config.utils import Range
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.attention import (
+    Attention,
+    get_attention_context,
+)
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from ..inductor_pass import enable_fake_mode
+from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
+from .matcher_utils import (
+    MatcherRotaryEmbedding,
+)
+from .rms_quant_fusion import (
+    empty_bf16,
+    empty_i64,
+)
+
+logger = init_logger(__name__)
+
+
+def fused_rope_and_unified_kv_cache_update_impl(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    layer_name: str = "",
+) -> torch.Tensor:
+    """
+    This impl fetches the KV cache and slot mapping from the forward context,
+    then calls the layer impl's `AttentionImpl.do_rope_and_kv_cache_update` method.
+    It also returns a dummy tensor, similar to `Attention.unified_kv_cache_update`,
+    that is passed to unified_attention to signal a side effect and
+    the data dependency between them to ensure torch.compile preserves ordering.
+    """
+    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
+    if layer_slot_mapping is not None:
+        attn_layer.impl.do_rope_and_kv_cache_update(
+            attn_layer,
+            query,
+            key,
+            value,
+            positions,
+            cos_sin_cache,
+            is_neox,
+            kv_cache,
+            layer_slot_mapping,
+        )
+
+    return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
+
+
+def fused_rope_and_unified_kv_cache_update_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    layer_name: str = "",
+) -> torch.Tensor:
+    return torch.empty(0, device=query.device, dtype=query.dtype)
+
+
+direct_register_custom_op(
+    op_name="fused_rope_and_unified_kv_cache_update",
+    op_func=fused_rope_and_unified_kv_cache_update_impl,
+    mutates_args=["query", "key"],
+    fake_impl=fused_rope_and_unified_kv_cache_update_fake,
+)
+
+
+class RopeReshapeKVCachePattern:
+    """
+    This pattern matches the following unfused inplace ops:
+      q, k = rotary_embedding(positions, q, k, head_size, cos_sin_cache, is_neox)
+      kv_cache_dummy = unified_kv_cache_update(k, v, layer_name)
+
+    and replaces it with the fused inplace op:
+      kv_cache_dummy = fused_rope_and_unified_kv_cache_update(
+        q, k, v, positions, cos_sin_cache, is_neox, layer_name
+      )
+    """
+
+    FUSED_OP = torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
+
+    def __init__(
+        self,
+        layer: Attention,
+        is_neox: bool,
+    ) -> None:
+        self.layer_name = layer.layer_name
+        self.num_heads = layer.num_heads
+        self.num_kv_heads = layer.num_kv_heads
+        self.head_size = layer.head_size
+        self.head_size_v = layer.head_size_v
+        self.is_neox = is_neox
+
+        self.q_size = self.num_heads * self.head_size
+        self.k_size = self.num_kv_heads * self.head_size
+        self.v_size = self.num_kv_heads * self.head_size_v
+
+        self.rope_matcher = MatcherRotaryEmbedding(
+            is_neox=self.is_neox,
+            head_size=self.head_size,
+            num_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
+        )
+
+    def get_inputs(self) -> list[torch.Tensor]:
+        # Sample inputs to help pattern tracing
+        T = 5
+        L = 4096
+        qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
+        positions = empty_i64(T)
+        cos_sin_cache = empty_bf16(L, self.head_size)
+        return [
+            qkv,
+            positions,
+            cos_sin_cache,
+        ]
+
+    def register(self, pm_pass: PatternMatcherPass) -> None:
+        def pattern(
+            qkv: torch.Tensor,
+            positions: torch.Tensor,
+            cos_sin_cache: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+            q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
+            q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
+            q = q.view(-1, self.num_heads, self.head_size)
+            k = k.view(-1, self.num_kv_heads, self.head_size)
+            v = v.view(-1, self.num_kv_heads, self.head_size_v)
+            dummy = torch.ops.vllm.unified_kv_cache_update(k, v, self.layer_name)
+            return dummy, q, k, v
+
+        def replacement(
+            qkv: torch.Tensor,
+            positions: torch.Tensor,
+            cos_sin_cache: torch.Tensor,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+            q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
+            q = q.view(-1, self.num_heads, self.head_size)
+            k = k.view(-1, self.num_kv_heads, self.head_size)
+            v = v.view(-1, self.num_kv_heads, self.head_size_v)
+            results = auto_functionalized(
+                self.FUSED_OP,
+                query=q,
+                key=k,
+                value=v,
+                positions=positions,
+                cos_sin_cache=cos_sin_cache,
+                is_neox=self.is_neox,
+                layer_name=self.layer_name,
+            )
+            return results[0], results[1], results[2], v
+
+        # NOTE: use view_to_reshape to unify view/reshape to simplify
+        # pattern and increase matching opportunities
+        def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
+            gm = pm.fwd_only(*args, **kwargs)
+            view_to_reshape(gm)
+            return gm
+
+        pm.register_replacement(
+            pattern, replacement, self.get_inputs(), fwd_and_view_to_reshape, pm_pass
+        )
+
+
+class RopeKVCacheFusionPass(VllmPatternMatcherPass):
+    """
+    This pass fuses the rotary embedding and KV cache update operations
+    into a single fused kernel if available.
+
+    It uses the pattern matcher and matches each layer manually, as strings
+    cannot be wildcarded. This also lets us check support on attention layers
+    upon registration instead of during pattern matching.
+
+    This fusion eliminates the need for separate kernel launches and
+    intermediate memory operations between the RoPE and cache update steps.
+    """
+
+    @enable_fake_mode
+    def __init__(self, config: VllmConfig) -> None:
+        super().__init__(config)
+
+        self.patterns: PatternMatcherPass = PatternMatcherPass(
+            pass_name="rope_kv_cache_fusion_pass"
+        )
+
+        cc = config.compilation_config
+        self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
+
+        attn_layers = get_layers_from_vllm_config(config, Attention)
+        for _, layer in attn_layers.items():
+            if layer.impl.fused_rope_kvcache_supported():
+                for is_neox in [True, False]:
+                    RopeReshapeKVCachePattern(
+                        layer=layer,
+                        is_neox=is_neox,
+                    ).register(self.patterns)
+
+        self.dump_patterns(config, self.patterns)
+
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        self.matched_count = self.patterns.apply(graph)
+        logger.debug("Replaced %s patterns", self.matched_count)
+
+    def is_applicable_for_range(self, compile_range: Range) -> bool:
+        # This pass works best for the small-batch decode setting.
+        # For large-batch e.g. prefill, it is better to use two separate kernels
+        # since they are compute bound and the fused kernels require further tuning.
+        return compile_range.end <= self.max_token_num
+
+    def uuid(self) -> str:
+        return VllmInductorPass.hash_source(self, RopeReshapeKVCachePattern)

vllm/compilation/passes/pass_manager.py

@@ -28,7 +28,9 @@ if current_platform.is_cuda_alike():
     from .fusion.attn_quant_fusion import AttnFusionPass
     from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
     from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
+    from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
     from .fusion.sequence_parallelism import SequenceParallelismPass
+    from .utility.scatter_split_replace import ScatterSplitReplacementPass
     from .utility.split_coalescing import SplitCoalescingPass
 
 if current_platform.is_cuda():
@@ -136,6 +138,11 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
             if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
                 self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
 
+            if self.pass_config.fuse_rope_kvcache:
+                self.passes += [SplitCoalescingPass(config)]
+                self.passes += [ScatterSplitReplacementPass(config)]
+                self.passes += [RopeKVCacheFusionPass(config)]
+
             if self.pass_config.fuse_attn_quant:
                 self.passes += [AttnFusionPass(config)]
 

vllm/compilation/passes/utility/fix_functionalization.py

@@ -162,6 +162,24 @@ class FixFunctionalizationPass(VllmInductorPass):
                     "position_ids",
                 )
                 self.defunctionalize(graph, node, mutated_args=mutated_args, args=args)
+            elif (
+                hasattr(torch.ops.vllm, "fused_rope_and_unified_kv_cache_update")
+                and at_target
+                == torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
+            ):
+                mutated_args = {
+                    1: "query",
+                    2: "key",
+                }
+                self.defunctionalize(graph, node, mutated_args=mutated_args)
+            # only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn
+            elif (
+                hasattr(torch.ops.vllm, "function_with_mutated_args_and_return")
+                and at_target
+                == torch.ops.vllm.function_with_mutated_args_and_return.default
+            ):
+                mutated_args = {1: "x"}
+                self.defunctionalize(graph, node, mutated_args=mutated_args)
             else:
                 continue  # skip the count
 
@@ -208,13 +226,20 @@ class FixFunctionalizationPass(VllmInductorPass):
         self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]
     ) -> None:
         """
-        Replace all getitem users of the auto-functionalized node with the
+        Replace mutated getitem users of the auto-functionalized node with the
         mutated arguments.
         :param node: The auto-functionalized node
         :param mutated_args: The mutated arguments, indexed by getitem index.
         If the value of an arg is a string, `node.kwargs[arg]` is used.
         """
         for idx, user in self.getitem_users(node).items():
+            # Some functionalized nodes may return both a result at getitem[0]
+            # as well as mutated args at getitem[1:...]
+            if idx == 0:
+                assert idx not in mutated_args, (
+                    f"result at getitem[0] should not be in mutated_args for {node}"
+                )
+                continue
             arg = mutated_args[idx]
             arg = node.kwargs[arg] if isinstance(arg, str) else arg
             user.replace_all_uses_with(arg)
@@ -257,10 +282,20 @@ class FixFunctionalizationPass(VllmInductorPass):
         with graph.inserting_before(node):
             function = node.args[0]
             if args is None:
-                graph.call_function(function, kwargs=node.kwargs)
+                fn_node = graph.call_function(function, kwargs=node.kwargs)
             else:
                 # Args passed as strings refer to items in node.kwargs
                 args = tuple(
                     node.kwargs[arg] if isinstance(arg, str) else arg for arg in args
                 )
-                graph.call_function(function, args=args)
+                fn_node = graph.call_function(function, args=args)
+
+        # If the function returns a value as well as mutating args inplace,
+        # the functionalized node will have a getitem[0] user that holds this value
+        # Replace getitem[0] user of the auto-functionalized node
+        # with the new defunctionalized node directly if it exists
+        users = self.getitem_users(node)
+        if 0 in users:
+            user = users[0]
+            user.replace_all_uses_with(fn_node)
+            self._remove(user)

vllm/compilation/passes/utility/scatter_split_replace.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single
+assignment if there are no users for the inplace tensor written to by
+the slice_scatter call.
+
+The inplace rotary_embedding custom op takes in mutable query and key inputs
+that are split+getitem outputs of a single qkv tensor.
+When functionalized, we fetch the rotated query and key from the functionalized op
+using `getitem` calls. However, we also write to the qkv tensor inplace using a
+`slice_scatter`, then split the inplace tensor to get the output tensors again.
+Instead, if the inplace tensor has no subsequent users, we can just replace the
+`slice_scatter` and `split_with_sizes` nodes with the `getitem` calls.
+
+This is already done in fix_functionalization::FixFunctionalizationPass, but
+writing a custom pass for it before defunctionalization allows matching against the
+qkv split+rotary_embedding subpattern as part of e.g. the RoPE+KVCache fusion pass.
+"""
+
+import operator
+
+import torch
+from torch import fx
+from torch._higher_order_ops.auto_functionalize import auto_functionalized
+
+from vllm.logger import init_logger
+
+from ..fx_utils import is_func
+from ..vllm_inductor_pass import VllmInductorPass
+
+logger = init_logger(__name__)
+
+
+class ScatterSplitReplacementPass(VllmInductorPass):
+    """Replace getitem+slice_scatter+split nodes with a single getitem when
+    the inplace subtensor written to by the slice_scatter has no other users.
+
+    Here's an example graph with q_size = 512, kv_size = 64:
+    split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
+    at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k))
+    q = operator.getitem(at, 1)
+    k = operator.getitem(at, 2)
+    torch.ops.aten.slice_scatter.default(qkv, q, [0, 512], -1)
+    torch.ops.aten.slice_scatter.default(qkv, k, [512, 512 + 64], -1)
+    split_with_sizes_2 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
+    q = operator.getitem(split_with_sizes_2, 0)
+    k = operator.getitem(split_with_sizes_2, 1)
+    v = operator.getitem(split_with_sizes_2, 2)
+
+    After this pass, this sequence of nodes is replaced with:
+    split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
+    at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k))
+    q = operator.getitem(at, 1)
+    k = operator.getitem(at, 2)
+    v = operator.getitem(split_with_sizes_1, 2)
+    """
+
+    @VllmInductorPass.time_and_log
+    def __call__(self, graph: fx.Graph) -> None:
+        count = 0
+
+        for node in graph.nodes:
+            if not is_func(node, auto_functionalized):
+                continue
+
+            kwargs = node.kwargs
+            at_target = node.args[0]
+
+            if at_target == torch.ops._C.rotary_embedding.default:
+                query = kwargs["query"]
+                key = kwargs["key"]
+                getitem_nodes = {}
+                for user in node.users:
+                    if is_func(user, operator.getitem):
+                        getitem_nodes[user.args[1]] = user
+
+                if (
+                    is_func(query, operator.getitem)
+                    and is_func(key, operator.getitem)
+                    and query.args[0] == key.args[0]
+                    and is_func(query.args[0], torch.ops.aten.split_with_sizes.default)
+                    and all(
+                        is_func(user, torch.ops.aten.slice_scatter.default)
+                        for getitem_node in getitem_nodes.values()
+                        for user in getitem_node.users
+                    )
+                ):
+                    # Pattern where query and key are slices of a qkv tensor.
+                    # While functionalized, results at [1] and [2] are scattered
+                    # back into qkv, then split again to get query and key.
+                    # If the inplace tensor has no other users, we can replace
+                    # the slice_scatter+split nodes with the original results.
+                    for user in getitem_nodes[1].users:
+                        slice_scatter_1_node = user
+                    if not is_func(
+                        slice_scatter_1_node, torch.ops.aten.slice_scatter.default
+                    ):
+                        continue
+
+                    for user in getitem_nodes[2].users:
+                        slice_scatter_2_node = user
+                    if not is_func(
+                        slice_scatter_2_node, torch.ops.aten.slice_scatter.default
+                    ):
+                        continue
+
+                    for user in slice_scatter_2_node.users:
+                        split_node = user
+                    if not is_func(split_node, torch.ops.aten.split_with_sizes.default):
+                        continue
+
+                    split_getitem_users = {}
+                    for user in split_node.users:
+                        if is_func(user, operator.getitem):
+                            split_getitem_users[user.args[1]] = user
+
+                    # Replace query node
+                    split_getitem_users[0].replace_all_uses_with(getitem_nodes[1])
+                    graph.erase_node(split_getitem_users[0])
+                    # Replace key node
+                    split_getitem_users[1].replace_all_uses_with(getitem_nodes[2])
+                    graph.erase_node(split_getitem_users[1])
+                    # Redirect value node to original qkv tensor
+                    split_getitem_users[2].replace_input_with(split_node, query.args[0])
+
+                    # Erase unused nodes
+                    graph.erase_node(split_node)
+                    graph.erase_node(slice_scatter_2_node)
+                    graph.erase_node(slice_scatter_1_node)
+
+                    count += 1
+
+        logger.debug("Eliminated %d slice_scatter+split nodes", count)

vllm/config/compilation.py

@@ -127,6 +127,13 @@ class PassConfig:
     # ROCm/AITER specific fusions
     fuse_act_padding: bool = Field(default=None)
     """Fuse the custom RMSNorm + padding ops."""
+    fuse_rope_kvcache: bool = Field(default=None)
+    """Fuse the QK rope + KV cache ops."""
+
+    rope_kvcache_fusion_max_token_num: int = 256
+    """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
+    Larger batch sizes e.g. during prefill will use the unfused kernels.
+    """
 
     fi_allreduce_fusion_max_size_mb: float | None = None
     """The threshold of the communicated tensor sizes under which
@@ -199,6 +206,7 @@ class PassConfig:
         "fuse_gemm_comms",
         "fuse_allreduce_rms",
         "fuse_act_padding",
+        "fuse_rope_kvcache",
         mode="wrap",
     )
     @classmethod
@@ -244,6 +252,12 @@ class PassConfig:
                 "The fusion will be disabled."
             )
             self.fuse_act_padding = False
+        if self.fuse_rope_kvcache and not current_platform.is_rocm():
+            logger.warning_once(
+                "KV cache fusion currently only enabled on ROCm. "
+                "The fusion will be disabled."
+            )
+            self.fuse_rope_kvcache = False
 
 
 class DynamicShapesType(str, enum.Enum):
@@ -825,6 +839,19 @@ class CompilationConfig:
             # TODO(zhuhaoran): support rope native forward match and remove this.
             # Linked issue: https://github.com/vllm-project/vllm/issues/28042
             self.custom_ops.append("+rotary_embedding")
+        if self.pass_config.fuse_rope_kvcache:
+            from vllm._aiter_ops import rocm_aiter_ops
+
+            if rocm_aiter_ops.is_triton_rotary_embed_enabled():
+                logger.warning(
+                    "Cannot use VLLM_ROCM_USE_AITER_TRITON_ROPE with "
+                    "fuse_rope_kvcache. Disabling fuse_rope_kvcache."
+                )
+                self.pass_config.fuse_rope_kvcache = False
+            else:
+                # TODO(Rohan138): support rope native forward match and remove this.
+                # Linked issue: https://github.com/vllm-project/vllm/issues/28042
+                self.custom_ops.append("+rotary_embedding")
 
         if (
             is_torch_equal_or_newer("2.9.0.dev")

vllm/config/vllm.py

@@ -1365,6 +1365,20 @@ class VllmConfig:
                         "allreduce-rms fusion will be enabled for all num_tokens."
                     )
 
+        if compilation_config.pass_config.fuse_rope_kvcache:
+            max_token_num = (
+                compilation_config.pass_config.rope_kvcache_fusion_max_token_num
+            )
+            if max_token_num is not None:
+                if compile_range_end is not None and max_token_num < compile_range_end:
+                    computed_compile_ranges_split_points.append(max_token_num)
+                else:
+                    logger.debug(
+                        "Max num batched tokens below rope+kvcache fusion threshold, "
+                        "rope+kvcache fusion enabled for num_tokens <= %d.",
+                        compile_range_end,
+                    )
+
         if compilation_config.compile_ranges_split_points is not None:
             for x in compilation_config.compile_ranges_split_points:
                 assert isinstance(x, int)

vllm/engine/arg_utils.py

@@ -592,6 +592,8 @@ class EngineArgs:
         "weight_transfer_config",
     )
 
+    fail_on_environ_validation: bool = False
+
     def __post_init__(self):
         # support `EngineArgs(compilation_config={...})`
         # without having to manually construct a
@@ -1235,6 +1237,14 @@ class EngineArgs:
             help="Log aggregate rather than per-engine statistics "
             "when using data parallelism.",
         )
+
+        parser.add_argument(
+            "--fail-on-environ-validation",
+            help="If set, the engine will raise an error if "
+            "environment validation fails.",
+            default=False,
+            action=argparse.BooleanOptionalAction,
+        )
         return parser
 
     @classmethod
@@ -1391,6 +1401,8 @@ class EngineArgs:
 
         device_config = DeviceConfig(device=cast(Device, current_platform.device_type))
 
+        envs.validate_environ(self.fail_on_environ_validation)
+
         # Check if the model is a speculator and override model/tokenizer/config
         # BEFORE creating ModelConfig, so the config is created with the target model
         # Skip speculator detection for cloud storage models (eg: S3, GCS) since

vllm/entrypoints/openai/api_server.py

@@ -8,6 +8,7 @@ import os
 import signal
 import socket
 import tempfile
+import warnings
 from argparse import Namespace
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
@@ -62,6 +63,8 @@ prometheus_multiproc_dir: tempfile.TemporaryDirectory
 # Cannot use __name__ (https://github.com/vllm-project/vllm/pull/4765)
 logger = init_logger("vllm.entrypoints.openai.api_server")
 
+_FALLBACK_SUPPORTED_TASKS: tuple[SupportedTask, ...] = ("generate",)
+
 
 @asynccontextmanager
 async def build_async_engine_client(
@@ -152,7 +155,19 @@ async def build_async_engine_client_from_engine_args(
             async_llm.shutdown()
 
 
-def build_app(args: Namespace, supported_tasks: tuple["SupportedTask", ...]) -> FastAPI:
+def build_app(
+    args: Namespace, supported_tasks: tuple["SupportedTask", ...] | None = None
+) -> FastAPI:
+    if supported_tasks is None:
+        warnings.warn(
+            "The 'supported_tasks' parameter was not provided to "
+            "build_app and will be required in a future version. "
+            "Defaulting to ('generate',).",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        supported_tasks = _FALLBACK_SUPPORTED_TASKS
+
     if args.disable_fastapi_docs:
         app = FastAPI(
             openapi_url=None, docs_url=None, redoc_url=None, lifespan=lifespan
@@ -263,9 +278,18 @@ async def init_app_state(
     engine_client: EngineClient,
     state: State,
     args: Namespace,
-    supported_tasks: tuple["SupportedTask", ...],
+    supported_tasks: tuple["SupportedTask", ...] | None = None,
 ) -> None:
     vllm_config = engine_client.vllm_config
+    if supported_tasks is None:
+        warnings.warn(
+            "The 'supported_tasks' parameter was not provided to "
+            "init_app_state and will be required in a future version. "
+            "Please pass 'supported_tasks' explicitly.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        supported_tasks = _FALLBACK_SUPPORTED_TASKS
 
     if args.served_model_name is not None:
         served_model_names = args.served_model_name

vllm/entrypoints/pooling/base/protocol.py

@@ -40,6 +40,21 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
             "if the served model does not use priority scheduling."
         ),
     )
+    mm_processor_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        description="Additional kwargs to pass to the HF processor.",
+    )
+    cache_salt: str | None = Field(
+        default=None,
+        description=(
+            "If specified, the prefix cache will be salted with the provided "
+            "string to prevent an attacker to guess prompts in multi-user "
+            "environments. The salt should be random, protected from "
+            "access by 3rd parties, and long enough to be "
+            "unpredictable (e.g., 43 characters base64-encoded, corresponding "
+            "to 256 bit)."
+        ),
+    )
     # --8<-- [end:pooling-common-extra-params]
 
 

vllm/entrypoints/pooling/classify/protocol.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import time
-from typing import Any, TypeAlias
+from typing import TypeAlias
 
 from pydantic import Field
 
@@ -48,12 +48,6 @@ class ClassificationCompletionRequest(
 class ClassificationChatRequest(
     PoolingBasicRequestMixin, ChatRequestMixin, ClassifyRequestMixin
 ):
-    # --8<-- [start:chat-classification-extra-params]
-    mm_processor_kwargs: dict[str, Any] | None = Field(
-        default=None,
-        description=("Additional kwargs to pass to the HF processor."),
-    )
-
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
         encoder_config = model_config.encoder_config or {}
 

vllm/entrypoints/pooling/embed/protocol.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import time
-from typing import Any, TypeAlias
+from typing import TypeAlias
 
 from pydantic import Field
 
@@ -78,11 +78,6 @@ class EmbeddingCompletionRequest(
 class EmbeddingChatRequest(
     PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin
 ):
-    mm_processor_kwargs: dict[str, Any] | None = Field(
-        default=None,
-        description=("Additional kwargs to pass to the HF processor."),
-    )
-
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
         encoder_config = model_config.encoder_config or {}
 

vllm/entrypoints/pooling/pooling/protocol.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import time
-from typing import Any, Generic, TypeAlias, TypeVar
+from typing import Generic, TypeAlias, TypeVar
 
 from pydantic import Field
 
@@ -65,11 +65,6 @@ class PoolingChatRequest(
 ):
     task: PoolingTask | None = None
 
-    mm_processor_kwargs: dict[str, Any] | None = Field(
-        default=None,
-        description=("Additional kwargs to pass to the HF processor."),
-    )
-
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
         encoder_config = model_config.encoder_config or {}
 

vllm/entrypoints/pooling/score/protocol.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import time
-from typing import Any, TypeAlias
+from typing import TypeAlias
 
 from pydantic import BaseModel, Field
 
@@ -23,13 +23,6 @@ from vllm.utils import random_uuid
 
 
 class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
-    # --8<-- [start:score-extra-params]
-    mm_processor_kwargs: dict[str, Any] | None = Field(
-        default=None,
-        description=("Additional kwargs to pass to the HF processor."),
-    )
-    # --8<-- [end:score-extra-params]
-
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
         encoder_config = model_config.encoder_config or {}
 
@@ -106,13 +99,6 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
     documents: ScoreInputs
     top_n: int = Field(default_factory=lambda: 0)
 
-    # --8<-- [start:rerank-extra-params]
-    mm_processor_kwargs: dict[str, Any] | None = Field(
-        default=None,
-        description=("Additional kwargs to pass to the HF processor."),
-    )
-    # --8<-- [end:rerank-extra-params]
-
     def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
         encoder_config = model_config.encoder_config or {}
 

vllm/envs.py

@@ -98,7 +98,6 @@ if TYPE_CHECKING:
     VLLM_DISABLED_KERNELS: list[str] = []
     VLLM_DISABLE_PYNCCL: bool = False
     VLLM_ROCM_USE_AITER: bool = False
-    VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
     VLLM_ROCM_USE_AITER_LINEAR: bool = True
     VLLM_ROCM_USE_AITER_MOE: bool = True
     VLLM_ROCM_USE_AITER_RMSNORM: bool = True
@@ -259,6 +258,14 @@ def maybe_convert_bool(value: str | None) -> bool | None:
     return bool(int(value))
 
 
+def use_aiter() -> bool:
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
+    return is_aiter_found_and_supported() and os.getenv(
+        "VLLM_ROCM_USE_AITER", "True"
+    ).lower() in ("true", "1")
+
+
 def disable_compile_cache() -> bool:
     return bool(int(os.getenv("VLLM_DISABLE_COMPILE_CACHE", "0")))
 
@@ -271,7 +278,7 @@ def use_aot_compile() -> bool:
 
     default_value = (
         "1"
-        if is_torch_equal_or_newer("2.10.0.dev") and not disable_compile_cache()
+        if is_torch_equal_or_newer("2.11.0.dev") and not disable_compile_cache()
         else "0"
     )
 
@@ -882,14 +889,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     ),
     # Disable aiter ops unless specifically enabled.
     # Acts as a parent switch to enable the rest of the other operations.
-    "VLLM_ROCM_USE_AITER": lambda: (
-        os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
-    ),
-    # Whether to use aiter paged attention.
-    # By default is disabled.
-    "VLLM_ROCM_USE_AITER_PAGED_ATTN": lambda: (
-        os.getenv("VLLM_ROCM_USE_AITER_PAGED_ATTN", "False").lower() in ("true", "1")
-    ),
+    "VLLM_ROCM_USE_AITER": use_aiter,
     # use aiter linear op if aiter ops are enabled
     # The following list of related ops
     # - scaled_mm (per-tensor / rowwise)
@@ -911,9 +911,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
         os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in ("true", "1")
     ),
     # Whether to use aiter mha ops.
-    # By default is enabled.
+    # By default is disabled.
     "VLLM_ROCM_USE_AITER_MHA": lambda: (
-        os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1")
+        os.getenv("VLLM_ROCM_USE_AITER_MHA", "False").lower() in ("true", "1")
     ),
     # Whether to use aiter fp4 gemm asm.
     # By default is disabled.
@@ -1606,6 +1606,15 @@ def is_set(name: str):
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
+def validate_environ(hard_fail: bool) -> None:
+    for env in os.environ:
+        if env.startswith("VLLM_") and env not in environment_variables:
+            if hard_fail:
+                raise ValueError(f"Unknown vLLM environment variable detected: {env}")
+            else:
+                logger.warning("Unknown vLLM environment variable detected: %s", env)
+
+
 def compile_factors() -> dict[str, object]:
     """Return env vars used for torch.compile cache keys.
 

vllm/lora/model_manager.py

@@ -35,7 +35,7 @@ from vllm.model_executor.models.interfaces import is_pooling_model
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.utils import PPMissingLayer
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.budget import MultiModalBudget
+from vllm.multimodal.encoder_budget import MultiModalBudget
 from vllm.utils.cache import LRUCache
 from vllm.utils.platform_utils import is_pin_memory_available
 

vllm/model_executor/layers/attention/attention.py

@@ -570,11 +570,11 @@ direct_register_custom_op(
 
 def get_attention_context(
     layer_name: str,
-) -> tuple[Any, "Attention | MLAAttention", torch.Tensor]:
+) -> tuple[Any, "Attention | MLAAttention", torch.Tensor, torch.Tensor]:
     """Extract attention context for a given layer.
 
     This helper function extracts the attention metadata, attention layer
-    instance, and KV cache tensor for a specific layer.
+    instance, KV cache tensor, and slot mapping for a specific layer.
 
     Args:
         layer_name: The name/identifier of the attention layer.
@@ -585,6 +585,7 @@ def get_attention_context(
             no metadata available
         - attn_layer: The attention layer instance (Attention or MLAAttention)
         - kv_cache: The KV cache tensor for current virtual engine
+        - slot_mapping: The slot mapping for this specific layer
 
         Note: attn_metadata may be None, but attn_layer and kv_cache are always
         extracted from the forward context.
@@ -593,9 +594,14 @@ def get_attention_context(
     attn_metadata = forward_context.attn_metadata
     if isinstance(attn_metadata, dict):
         attn_metadata = attn_metadata[layer_name]
-    attn_layer = forward_context.no_compile_layers[layer_name]
+    attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
     kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
-    return attn_metadata, attn_layer, kv_cache
+    slot_mapping = forward_context.slot_mapping
+    assert isinstance(slot_mapping, dict), (
+        f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
+    )
+    layer_slot_mapping = slot_mapping.get(layer_name)
+    return attn_metadata, attn_layer, kv_cache, layer_slot_mapping
 
 
 @maybe_transfer_kv_layer
@@ -605,7 +611,7 @@ def unified_attention(
     value: torch.Tensor,
     layer_name: str,
 ) -> torch.Tensor:
-    attn_metadata, self, kv_cache = get_attention_context(layer_name)
+    attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
     output = self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
 
     return output
@@ -636,15 +642,7 @@ def unified_kv_cache_update(
     Returns a dummy that is passed to unified_attention to signal a side effect and
     the data dependency between them to ensure torch.compile preserves ordering.
     """
-    forward_context = get_forward_context()
-    attn_layer = forward_context.no_compile_layers[layer_name]
-    kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
-
-    slot_mapping = forward_context.slot_mapping
-    assert isinstance(slot_mapping, dict), (
-        f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
-    )
-    layer_slot_mapping = slot_mapping.get(layer_name)
+    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
     if layer_slot_mapping is not None:
         assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
             f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
@@ -691,7 +689,7 @@ def unified_attention_with_output(
     # that ensures torch.compile preserves ordering between KV cache update and
     # attention forward.
     del kv_cache_dummy_dep
-    attn_metadata, self, kv_cache = get_attention_context(layer_name)
+    attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
 
     self.impl.forward(
         self,

vllm/model_executor/layers/attention/kv_transfer_utils.py

@@ -40,8 +40,8 @@ def maybe_transfer_kv_layer(func: Callable) -> Callable:
 
         layer_name: str = args[layer_name_index]
 
-        # Extract attention context (layer-specific metadata, layer, and kv_cache)
-        attn_metadata, attn_layer, kv_cache = get_attention_context(layer_name)
+        # Extract attention context (metadata, layer, kv_cache, layer_slot_mapping)
+        attn_metadata, _, kv_cache, _ = get_attention_context(layer_name)
         connector = get_kv_transfer_group()
         if attn_metadata is None or not connector.has_connector_metadata():
             return func(*args, **kwargs)

vllm/model_executor/layers/attention/mla_attention.py

@@ -827,7 +827,7 @@ def unified_mla_attention(
     k_pe: torch.Tensor,
     layer_name: str,
 ) -> torch.Tensor:
-    attn_metadata, layer, kv_cache = get_attention_context(layer_name)
+    attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
     output = layer.forward_impl(q, kv_c_normed, k_pe, kv_cache, attn_metadata)
 
     return output
@@ -861,7 +861,7 @@ def unified_mla_attention_with_output(
     output_scale: torch.Tensor | None = None,
     output_block_scale: torch.Tensor | None = None,
 ) -> None:
-    attn_metadata, layer, kv_cache = get_attention_context(layer_name)
+    attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
     layer.forward_impl(
         q,
         kv_c_normed,

vllm/model_executor/layers/batch_invariant.py

@@ -974,7 +974,7 @@ def enable_batch_invariant_mode():
     )
 
     reduced_precision_val = (
-        (False, False) if is_torch_equal_or_newer("2.10.0.dev") else False
+        (False, False) if is_torch_equal_or_newer("2.10.0") else False
     )
     torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
         reduced_precision_val

vllm/model_executor/layers/fused_moe/__init__.py

@@ -102,6 +102,7 @@ if HAS_TRITON:
     )
     from vllm.model_executor.layers.fused_moe.xpu_fused_moe import (
         XPUExperts,
+        XPUExpertsFp8,
     )
 
     __all__ += [
@@ -121,6 +122,7 @@ if HAS_TRITON:
         "BatchedDeepGemmExperts",
         "TritonOrDeepGemmExperts",
         "XPUExperts",
+        "XPUExpertsFp8",
     ]
 else:
     # Some model classes directly use the custom ops. Add placeholders

vllm/model_executor/layers/fused_moe/config.py

@@ -386,6 +386,10 @@ class FusedMoEQuantConfig:
     def use_nvfp4_w4a4(self) -> bool:
         return self.quant_dtype == "nvfp4"
 
+    @property
+    def use_mxfp4_w4a8(self) -> bool:
+        return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4"
+
     def config_name(self, dtype: torch.dtype) -> str | None:
         """
         Return a string used to construct the filename that contains the
@@ -532,6 +536,8 @@ def fp8_w8a8_moe_quant_config(
     w2_scale: torch.Tensor,
     a1_scale: torch.Tensor | None = None,
     a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
     per_act_token_quant: bool = False,
     per_out_ch_quant: bool = False,
     block_shape: list[int] | None = None,
@@ -549,6 +555,8 @@ def fp8_w8a8_moe_quant_config(
         g1_alphas=g1_alphas,
         w2_scale=w2_scale,
         g2_alphas=g2_alphas,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
         a1_scale=a1_scale,
         a1_gscale=a1_gscale,
         a2_scale=a2_scale,
@@ -564,6 +572,8 @@ def int8_w8a8_moe_quant_config(
     w2_scale: torch.Tensor,
     a1_scale: torch.Tensor | None,
     a2_scale: torch.Tensor | None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
     per_act_token_quant: bool = False,
 ) -> FusedMoEQuantConfig:
     """
@@ -575,6 +585,8 @@ def int8_w8a8_moe_quant_config(
         w2_scale=w2_scale,
         a1_scale=a1_scale,
         a2_scale=a2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
         per_act_token_quant=per_act_token_quant,
         per_out_ch_quant=False,
         block_shape=None,
@@ -654,6 +666,26 @@ def mxfp4_mxfp8_moe_quant_config(
     )
 
 
+def mxfp4_w4a8_moe_quant_config(
+    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for fp8 activations and mxfp4 weights.
+    """
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None),
+        _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None),
+        _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
+        _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
+    )
+
+
 def ocp_mx_moe_quant_config(
     quant_dtype: str,
     w1_scale: Union[torch.Tensor, "PrecisionConfig"],
@@ -691,6 +723,8 @@ def nvfp4_moe_quant_config(
     a2_gscale: torch.Tensor,
     w1_scale: torch.Tensor,
     w2_scale: torch.Tensor,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for mxfp4 activations and nvp4 weights.
@@ -699,6 +733,8 @@ def nvfp4_moe_quant_config(
         "nvfp4",
         w1_scale=w1_scale,
         w2_scale=w2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
         a1_gscale=a1_gscale,
         a2_gscale=a2_gscale,
         g1_alphas=g1_alphas,
@@ -787,6 +823,32 @@ def int8_w8a16_moe_quant_config(
     )
 
 
+def mxfp4_w4a4_moe_quant_config(
+    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for mxfp4 activations and mxfp4 weights.
+    """
+    return FusedMoEQuantConfig.make(
+        "mxfp4",
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=block_shape,
+    )
+
+
 def int4_w4afp8_moe_quant_config(
     w1_scale: torch.Tensor,
     w2_scale: torch.Tensor,

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -38,7 +38,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
 )
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
 from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
-from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
     kFp8Dynamic128Sym,
@@ -1583,6 +1582,11 @@ def _get_config_quant_dtype(
         return "mxfp6_e3m2"
     elif ocp_mx_scheme in {"w_mxfp4_a_mxfp6_e2m3", "w_mxfp6_e2m3_a_mxfp6_e2m3"}:
         return "mxfp6_e2m3"
+    elif ocp_mx_scheme in {"w_mxfp4", "w_mxfp6_e3m2", "w_mxfp6_e2m3"}:
+        return torch.bfloat16
+    elif ocp_mx_scheme in {"w_mxfp4_a_fp8", "w_mxfp6_e3m2_a_fp8", "w_mxfp6_e2m3_a_fp8"}:
+        return torch.float8_e4m3fn
+
     return None
 
 
@@ -1617,17 +1621,10 @@ def fused_experts_impl(
     if use_int4_w4a16:
         assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"
     elif ocp_mx_scheme is not None:
-        if ocp_mx_scheme in {
-            "w_mxfp4_a_mxfp4",
-            "w_mxfp4_a_mxfp6_e3m2",
-            "w_mxfp4_a_mxfp6_e2m3",
-        }:
+        if ocp_mx_scheme.startswith("w_mxfp4"):
             # 16bit activation and fp4x2 packed weight
             assert hidden_states.size(1) == w1.size(2) * 2, "hidden size mismatch"
-        elif ocp_mx_scheme in {
-            "w_mxfp6_e3m2_a_mxfp6_e3m2",
-            "w_mxfp6_e2m3_a_mxfp6_e2m3",
-        }:
+        elif ocp_mx_scheme.startswith("w_mxfp6"):
             assert hidden_states.size(1) == (w1.size(2) * 4) // 3, (
                 "hidden size mismatch"
             )
@@ -1717,17 +1714,13 @@ def fused_experts_impl(
         # TODO: On platforms for which `current_platform.supports_mx()` is True
         # and for which we have a native OCP mx fused MOE kernel,
         # this dequantization step should not be done.
-        if ocp_mx_scheme in {
-            OCP_MX_Scheme.w_mxfp4_a_mxfp4,
-            OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
-            OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
-        }:
+        if ocp_mx_scheme.startswith("w_mxfp4"):
             # Weight has to be dequantized for mxfp4 emulation.
             w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
             w1_scale = None
             w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
             w2_scale = None
-        elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2:
+        elif ocp_mx_scheme.startswith("w_mxfp6_e3m2"):
             w1 = dequant_mxfp6(
                 w1, w1_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
             )
@@ -1736,7 +1729,7 @@ def fused_experts_impl(
                 w2, w2_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
             )
             w2_scale = None
-        elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3:
+        elif ocp_mx_scheme.startswith("w_mxfp6_e2m3"):
             w1 = dequant_mxfp6(
                 w1, w1_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
             )
@@ -1779,6 +1772,7 @@ def fused_experts_impl(
             quant_dtype=quant_dtype,
             per_act_token_quant=per_channel_quant,
             block_shape=block_shape,
+            ocp_mx_scheme=ocp_mx_scheme,
         )
 
         # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k
@@ -1846,6 +1840,7 @@ def fused_experts_impl(
             quant_dtype=quant_dtype,
             per_act_token_quant=per_channel_quant,
             block_shape=block_shape,
+            ocp_mx_scheme=ocp_mx_scheme,
         )
 
         if expert_map is not None:

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

@@ -19,17 +19,42 @@ from vllm.model_executor.layers.fused_moe.utils import _resize_cache
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
 )
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.import_utils import has_triton_kernels
 
 logger = init_logger(__name__)
 
+use_legacy_triton_kernels = False
+
 if has_triton_kernels():
     try:
         import triton_kernels.swiglu
-        from triton_kernels.matmul_ogs import FnSpecs, FusedActivation, matmul_ogs
-        from triton_kernels.routing import RoutingData, routing, routing_from_bitmatrix
-        from triton_kernels.tensor import Bitmatrix
+        from triton_kernels.matmul_ogs import (
+            FnSpecs,
+            FusedActivation,
+            GatherIndx,
+            RoutingData,
+            ScatterIndx,
+            matmul_ogs,
+        )
+        from triton_kernels.tensor import (
+            BIT,
+            Bitmatrix,
+        )
+        from triton_kernels.topk import topk
+
+        try:
+            from triton_kernels.tensor import (
+                SparseMatrix,
+                make_ragged_tensor_metadata,
+            )
+        except ImportError:
+            if current_platform.is_rocm():
+                logger.warning_once("Using legacy triton_kernels on ROCm")
+                use_legacy_triton_kernels = True
+            else:
+                raise
     except (AttributeError, ImportError) as e:
         logger.error(
             "Failed to import Triton kernels. Please make sure your triton "
@@ -78,6 +103,68 @@ def pack_bitmatrix(
         tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows)
 
 
+def legacy_routing_from_bitmatrix(
+    bitmatrix: "Bitmatrix",
+    expt_scal: torch.Tensor,
+    expt_indx: torch.Tensor,
+    n_expts_tot: int,
+    n_expts_act: int,
+) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
+    """
+    Replacement for the removed triton_kernels.routing.routing_from_bitmatrix.
+    Creates routing data from a bitmatrix representation.
+    """
+    if use_legacy_triton_kernels:
+        from triton_kernels.routing import routing_from_bitmatrix
+
+        return routing_from_bitmatrix(
+            bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act
+        )
+    sparse_logits = SparseMatrix(indx=expt_indx, vals=expt_scal, mask=bitmatrix)
+    dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
+    combine_indx = sparse_logits.mask_metadata.col_sorted_indx
+    ragged_batch_metadata = make_ragged_tensor_metadata(
+        sparse_logits.mask_metadata.col_sum,
+        dispatch_indx.shape[0],
+    )
+    gate_scal = sparse_logits.vals.flatten()[combine_indx]
+    routing_data = RoutingData(
+        gate_scal,
+        ragged_batch_metadata.block_sizes,
+        n_expts_tot,
+        n_expts_act,
+        ragged_batch_metadata,
+    )
+    gather_idx = GatherIndx(combine_indx, dispatch_indx)
+    scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
+    return routing_data, gather_idx, scatter_idx
+
+
+def legacy_routing(
+    logits: torch.Tensor,
+    n_expts_act: int,
+    sm_first: bool = False,
+) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
+    """
+    Replacement for the removed triton_kernels.routing.routing function.
+    Computes routing data from gating logits.
+    """
+    if use_legacy_triton_kernels:
+        from triton_kernels.routing import routing
+
+        return routing(logits, n_expts_act, sm_first=sm_first)
+    if sm_first:
+        logits = torch.softmax(logits, dim=-1)
+    sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first)
+    return legacy_routing_from_bitmatrix(
+        sparse_logits.mask,
+        sparse_logits.vals,
+        sparse_logits.indx,
+        logits.shape[-1],
+        n_expts_act,
+    )
+
+
 def triton_kernel_moe_forward(
     hidden_states: torch.Tensor,
     w1,  # Tensor or triton_kernels.Tensor
@@ -91,7 +178,7 @@ def triton_kernel_moe_forward(
     global_num_experts: int = -1,
     expert_map: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    routing_data, gather_idx, scatter_idx = routing(
+    routing_data, gather_idx, scatter_idx = legacy_routing(
         gating_output, topk, sm_first=not renormalize
     )
 
@@ -167,10 +254,22 @@ def triton_kernel_fused_experts(
     )
     output_tensor = _resize_cache(output_tensor, (batch_dim, M, K))
 
-    act = FusedActivation(
-        FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
-        (swiglu_alpha, swiglu_limit),
-        2,
+    act = (
+        FusedActivation(
+            FnSpecs(
+                "swiglu",
+                triton_kernels.swiglu.swiglu_fn,
+                ("alpha", "limit"),
+                reduction_n=2,
+            ),
+            (swiglu_alpha, swiglu_limit),
+        )
+        if not use_legacy_triton_kernels
+        else FusedActivation(
+            FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
+            (swiglu_alpha, swiglu_limit),
+            2,
+        )
     )
     gammas = routing_data.gate_scal if routing_data else None
 
@@ -200,6 +299,182 @@ def triton_kernel_fused_experts(
     return output_tensor
 
 
+def triton_kernel_moe_oss_forward(
+    hidden_states: torch.Tensor,
+    w1,  # Tensor or triton_kernels.Tensor
+    w2,  # Tensor or triton_kernels.Tensor
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    activation: str = "silu",
+    quant_config: FusedMoEQuantConfig | None = None,
+    apply_router_weight_on_input: bool = False,
+    global_num_experts: int = -1,
+    expert_map: torch.Tensor | None = None,
+    unpadded_N_w1=None,
+    unpadded_K_w1=None,
+    unpadded_N_w2=None,
+    unpadded_K_w2=None,
+) -> torch.Tensor:
+    assert quant_config is not None
+
+    if quant_config.use_mxfp4_w4a16:
+        from triton_kernels.routing import routing
+
+        routing_data, gather_idx, scatter_idx = routing(
+            gating_output, topk, sm_first=not renormalize
+        )
+    elif quant_config.use_mxfp4_w4a4:
+        from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
+
+        routing_data, gather_idx, scatter_idx = aiter_routing(
+            gating_output, topk, sm_first=not renormalize
+        )
+
+    return triton_kernel_fused_oss_experts(
+        None,
+        hidden_states,
+        w1,
+        w2,
+        routing_data,
+        gather_idx,
+        scatter_idx,
+        activation=activation,
+        quant_config=quant_config,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        global_num_experts=global_num_experts,
+        expert_map=expert_map,
+        unpadded_N_w1=unpadded_N_w1,
+        unpadded_K_w1=unpadded_K_w1,
+        unpadded_N_w2=unpadded_N_w2,
+        unpadded_K_w2=unpadded_K_w2,
+    )
+
+
+# This is a triton implementation of the fused_experts function
+def triton_kernel_fused_oss_experts(
+    output_tensor: torch.Tensor,
+    hidden_states: torch.Tensor,
+    w1,  # Tensor or triton_kernels.Tensor
+    w2,  # Tensor or triton_kernels.Tensor
+    routing_data,  # RoutingData
+    gather_indx,  # GatherIndx
+    scatter_indx,  # ScatterIndx
+    activation: str = "silu",
+    quant_config: FusedMoEQuantConfig | None = None,
+    swiglu_alpha: float = 1.702,
+    swiglu_limit: float = 7.0,
+    apply_router_weight_on_input: bool = False,
+    global_num_experts: int = -1,
+    expert_map: torch.Tensor | None = None,
+    a1q_scale: torch.Tensor | None = None,
+    unpadded_N_w1=None,
+    unpadded_K_w1=None,
+    unpadded_N_w2=None,
+    unpadded_K_w2=None,
+) -> torch.Tensor:
+    if quant_config is None:
+        quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
+
+    # type check, uint8 means mxfp4
+    assert hidden_states.dtype == torch.bfloat16
+    assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
+    assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
+
+    # Shape check, only check non-mxfp4
+    assert hidden_states.shape[-1] == w1.shape[-2]
+    assert w2.shape[-1] == w1.shape[1]
+
+    E, _, N = w1.shape
+
+    if global_num_experts == -1:
+        global_num_experts = E
+
+    gammas = routing_data.gate_scal if routing_data else None
+
+    if quant_config.use_mxfp4_w4a16:
+        act = FusedActivation(
+            FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
+            (swiglu_alpha, swiglu_limit),
+            2,
+        )
+        intermediate_cache1 = matmul_ogs(
+            hidden_states,
+            w1,
+            quant_config.w1_bias,
+            routing_data,
+            gather_indx=gather_indx,
+            precision_config=quant_config.w1_precision,
+            gammas=gammas if apply_router_weight_on_input else None,
+            fused_activation=act,
+        )
+        intermediate_cache3 = matmul_ogs(
+            intermediate_cache1,
+            w2,
+            quant_config.w2_bias,
+            routing_data,
+            scatter_indx=scatter_indx,
+            precision_config=quant_config.w2_precision,
+            gammas=None if apply_router_weight_on_input else gammas,
+            y=output_tensor,
+        )
+
+    elif quant_config.use_mxfp4_w4a4:
+        from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
+        from aiter.ops.triton.quant_moe import downcast_to_static_fp8
+
+        assert quant_config.w1_precision is not None, (
+            "w1_precision in quant config can't be None"
+        )
+        assert quant_config.w2_precision is not None, (
+            "w2_precision in quant config can't be None"
+        )
+
+        hidden_states = downcast_to_static_fp8(
+            hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
+        )
+
+        intermediate_cache1 = moe_gemm_a8w4(
+            hidden_states,
+            w1.storage.data,
+            None,
+            quant_config.w1_precision.weight_scale.storage.data,
+            quant_config.w1_precision.flex_ctx.lhs_data.scale,
+            quant_config.w2_precision.flex_ctx.lhs_data.scale,
+            quant_config.w1_bias,
+            routing_data,
+            gather_indx=gather_indx,
+            gammas=gammas if apply_router_weight_on_input else None,
+            swizzle_mx_scale="CDNA4_SCALE",
+            out_dtype=torch.float8_e4m3fn,
+            apply_swiglu=True,
+            alpha=swiglu_alpha,
+            limit=swiglu_limit,
+            unpadded_N=unpadded_N_w1,
+            unpadded_K=unpadded_K_w1,
+        )
+
+        intermediate_cache3 = moe_gemm_a8w4(
+            intermediate_cache1,
+            w2.storage.data,
+            None,
+            quant_config.w2_precision.weight_scale.storage.data,
+            quant_config.w2_precision.flex_ctx.lhs_data.scale,
+            None,
+            quant_config.w2_bias,
+            routing_data,
+            scatter_indx=scatter_indx,
+            gammas=None if apply_router_weight_on_input else gammas,
+            swizzle_mx_scale="CDNA4_SCALE",
+            unpadded_N=unpadded_N_w2,
+            unpadded_K=unpadded_K_w2,
+        )
+    else:
+        raise AssertionError(f"Non supported {quant_config=} in fused MoE op")
+
+    return intermediate_cache3
+
+
 def make_routing_data(
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
@@ -231,13 +506,22 @@ def make_routing_data(
 
     bitmatrix_shape = [n_rows, bm_cols * 32]
     bitmatrix_shape_max = [n_rows, None]
-    bitmatrix = Bitmatrix(
-        bitmatrix, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max, scratchpad=None
+    bitmatrix = (
+        Bitmatrix(
+            bitmatrix, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
+        )
+        if not use_legacy_triton_kernels
+        else Bitmatrix(
+            bitmatrix,
+            shape=bitmatrix_shape,
+            shape_max=bitmatrix_shape_max,
+            scratchpad=None,
+        )
     )
 
     # matmul_ogs expects invalid topk_weights to be -1s
     topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
-    routing_data, gather_indx, scatter_indx = routing_from_bitmatrix(
+    routing_data, gather_indx, scatter_indx = legacy_routing_from_bitmatrix(
         bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
     )
 
@@ -379,6 +663,9 @@ class OAITritonExperts(BaseOAITritonExperts):
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
+        if self.quant_config is None:
+            self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG
+
         if expert_map is not None:
             topk_ids = expert_map[topk_ids]
 

vllm/model_executor/layers/fused_moe/layer.py

@@ -221,12 +221,14 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
     )
 
 
+# TODO(rob): move this down to the kernel.
 def maybe_roundup_hidden_size(
     hidden_size: int,
     act_dtype: torch.dtype,
-    quant_config: QuantizationConfig | None,
     moe_parallel_config: FusedMoEParallelConfig,
     is_lora_enabled: bool,
+    model_type: str | None,
+    is_mxfp4_quant: bool,
 ) -> int:
     """
     Given layer hidden size and MoE configurations, round up hidden_size
@@ -235,11 +237,12 @@ def maybe_roundup_hidden_size(
     Args:
         hidden_size: Layer hidden-size
         act_dtype: Data type of the layer activations.
-        quant_config: Fused MoE quantization configuration.
         moe_parallel_config: Fused MoE parallelization strategy configuration.
         is_lora_enabled: True if the engine is enabled with LoRA. This
             is used in the case of mxfp4 quantization in selecting the
             MxFP4Backend.
+        model_type: for checking if gpt-oss
+        is_mxfp4_quant: whether the layer is quantized with mxfp4
 
     Return:
         Rounded up hidden_size if rounding up is required based on the configs.
@@ -254,7 +257,7 @@ def maybe_roundup_hidden_size(
     )
 
     # we are padding globally so EP buffer allocation works
-    if quant_config and quant_config.get_name() == "mxfp4":
+    if model_type == "gpt_oss" and is_mxfp4_quant:
         from vllm.model_executor.layers.quantization.mxfp4 import (
             Mxfp4Backend,
             get_mxfp4_backend,
@@ -398,15 +401,6 @@ class FusedMoE(CustomOp):
         # Expert mapping used in self.load_weights
         self.expert_mapping = expert_mapping
 
-        # Round up hidden size if needed.
-        hidden_size = maybe_roundup_hidden_size(
-            hidden_size,
-            moe_in_dtype,
-            quant_config,
-            self.moe_parallel_config,
-            is_lora_enabled=self.vllm_config.lora_config is not None,
-        )
-
         # For smuggling this layer into the fused moe custom op
         compilation_config = vllm_config.compilation_config
         if prefix in compilation_config.static_forward_context:
@@ -508,7 +502,6 @@ class FusedMoE(CustomOp):
             ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."
 
         assert intermediate_size % self.tp_size == 0
-        self.hidden_size = hidden_size
         self.intermediate_size_per_partition = intermediate_size // self.tp_size
         self.reduce_results = reduce_results
         self.renormalize = renormalize
@@ -548,6 +541,26 @@ class FusedMoE(CustomOp):
         )
         self.routing_method_type: RoutingMethodType = self.router.routing_method_type
 
+        # Round up hidden size before creating moe_config.
+        # This way moe_config is created with the correct hidden_size from the start.
+        unpadded_hidden_size = hidden_size
+        self.model_type = (
+            self.vllm_config.model_config.hf_config.model_type
+            if self.vllm_config.model_config is not None
+            else None
+        )
+        hidden_size = maybe_roundup_hidden_size(
+            hidden_size=hidden_size,
+            act_dtype=moe_in_dtype,
+            moe_parallel_config=self.moe_parallel_config,
+            is_lora_enabled=vllm_config.lora_config is not None,
+            model_type=self.model_type,
+            is_mxfp4_quant=(
+                quant_config is not None and quant_config.is_mxfp4_quant(prefix, self)
+            ),
+        )
+        self.hidden_size = hidden_size
+
         self.moe_config: FusedMoEConfig = FusedMoEConfig(
             num_experts=self.global_num_experts,
             experts_per_token=top_k,
@@ -615,6 +628,7 @@ class FusedMoE(CustomOp):
         moe_quant_params = {
             "num_experts": self.local_num_experts,
             "hidden_size": hidden_size,
+            "unpadded_hidden_size": unpadded_hidden_size,
             "intermediate_size_per_partition": self.intermediate_size_per_partition,
             "params_dtype": params_dtype,
             "weight_loader": self.weight_loader,
@@ -1147,16 +1161,43 @@ class FusedMoE(CustomOp):
         expert_id: int,
         return_success: bool = False,
     ) -> bool | None:
-        if self.quant_config and self.quant_config.get_name() == "mxfp4":
-            # (FIXME) for gpt-oss all experts are combined
-            if "bias" in weight_name:
-                dim1 = loaded_weight.shape[1]
-                param.data[:, :dim1].copy_(loaded_weight)
-            else:
-                dim1 = loaded_weight.shape[1]
-                dim2 = loaded_weight.shape[2]
-                param.data[:, :dim1, :dim2].copy_(loaded_weight)
-            return True if return_success else None
+        if self.quant_config is not None:
+            if self.quant_config.get_name() == "mxfp4":
+                # (FIXME) for gpt-oss all experts are combined
+                if "bias" in weight_name:
+                    dim1 = loaded_weight.shape[1]
+                    param.data[:, :dim1].copy_(loaded_weight)
+                else:
+                    dim1 = loaded_weight.shape[1]
+                    dim2 = loaded_weight.shape[2]
+                    param.data[:, :dim1, :dim2].copy_(loaded_weight)
+                return True if return_success else None
+            elif (
+                self.quant_config.get_name() == "quark" and self.model_type == "gpt_oss"
+            ):
+                # When self._is_mxfp4 is true, model_dtype must be gpt_oss
+                expert_data = param.data[expert_id]
+                if "input_scale" in weight_name:
+                    assert loaded_weight.numel() == 1
+                    expert_data.data.copy_(loaded_weight)
+                    return True if return_success else None
+
+                shard_dim = (
+                    0 if shard_id in ("w1", "w3") or "bias" in weight_name else 1
+                )
+                if shard_id == "w2":
+                    shard_size = loaded_weight.shape[shard_dim] // self.tp_size
+                    loaded_weight = loaded_weight.narrow(
+                        shard_dim, shard_size * self.tp_rank, shard_size
+                    )
+                if "bias" in weight_name:
+                    dim1 = loaded_weight.shape[0]
+                    expert_data.data[:dim1].copy_(loaded_weight)
+                else:
+                    dim1 = loaded_weight.shape[0]
+                    dim2 = loaded_weight.shape[1]
+                    expert_data.data[:dim1, :dim2].copy_(loaded_weight)
+                return True if return_success else None
 
         quant_method_name = self.quant_method.__class__.__name__
         global_expert_id = expert_id

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

@@ -52,6 +52,7 @@ class Fp8MoeBackend(Enum):
     AITER = "AITER"
     VLLM_CUTLASS = "VLLM_CUTLASS"
     BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
+    XPU = "XPU"
 
 
 def backend_to_kernel_cls(
@@ -123,6 +124,13 @@ def backend_to_kernel_cls(
 
         return CutlassBatchedExpertsFp8
 
+    elif backend == Fp8MoeBackend.XPU:
+        from vllm.model_executor.layers.fused_moe.xpu_fused_moe import (
+            XPUExpertsFp8,
+        )
+
+        return XPUExpertsFp8
+
     else:
         raise ValueError(f"Unknown FP8 MoE backend: {backend.value}")
 
@@ -154,6 +162,7 @@ def select_fp8_moe_backend(
         Fp8MoeBackend.TRITON,
         Fp8MoeBackend.BATCHED_TRITON,
         Fp8MoeBackend.MARLIN,
+        Fp8MoeBackend.XPU,
     ]
 
     # NOTE(rob): We need to peak into the P/F selection to determine
@@ -393,6 +402,7 @@ def convert_to_fp8_moe_kernel_format(
             Fp8MoeBackend.BATCHED_TRITON,
             Fp8MoeBackend.VLLM_CUTLASS,
             Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
+            Fp8MoeBackend.XPU,
         ]:
             raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}")
 

vllm/model_executor/layers/fused_moe/utils.py

@@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     mxfp8_e4m3_quantize,
 )
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    per_tensor_dequantize,
+)
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -241,7 +244,27 @@ def moe_kernel_quantize_input(
     per_act_token_quant: bool,
     block_shape: list[int] | None = None,
     is_fp4_scale_swizzled: bool = True,
+    ocp_mx_scheme: str | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
+    # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation
+    if ocp_mx_scheme is not None:
+        if ocp_mx_scheme in {"w_mxfp4", "w_mxfp4_a_mxfp4"}:
+            pass  # No QDQ needed for these schemes
+        elif ocp_mx_scheme.endswith("a_fp8"):
+            # Perform QDQ (quantize and dequantize) on activation for emulation
+            # purpose, because there is no native kernel for weight in ocp_mx_scheme
+            # and activation in FP8. The implementation is based on existing
+            # non-emulation ops.
+            qA, qA_scale = ops.scaled_fp8_quant(
+                A, A_scale, use_per_token_if_dynamic=False
+            )
+            A = per_tensor_dequantize(qA, qA_scale).to(A.dtype)
+            # After QDQ, we don't need further quantization
+            return A, None
+        # else: For other schemes (e.g., *_a_mxfp6_e3m2, *_a_mxfp6_e2m3),
+        # weights are already dequantized, and we proceed with normal
+        # activation quantization below.
+
     if quant_dtype == torch.float8_e4m3fn:
         return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
     elif quant_dtype == torch.int8:

vllm/model_executor/layers/fused_moe/xpu_fused_moe.py

@@ -4,13 +4,16 @@ import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
     FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceNoOP,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
+    kFp8DynamicTensorSym,
     kFp8StaticTensorSym,
 )
 from vllm.platforms import current_platform
@@ -20,6 +23,21 @@ if current_platform.is_xpu():
 
 
 class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int | None = None,
+        num_dispatchers: int | None = None,
+    ):
+        super().__init__(
+            moe_config,
+            quant_config,
+            max_num_tokens,
+            num_dispatchers,
+        )
+        self.is_fp8 = False
+
     @property
     def expects_unquantized_inputs(self) -> bool:
         return True
@@ -49,10 +67,10 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        # TODO: dispatch based on device.
         SUPPORTED_W_A = [
             (None, None),
             (kFp8StaticTensorSym, None),
+            (kFp8StaticTensorSym, kFp8DynamicTensorSym),
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A
 
@@ -103,10 +121,10 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):
         xpu_fused_moe(
             hidden_states=hidden_states,
             w13=w1,
-            w13_scales=a1q_scale,
+            w13_scales=self.w1_scale,
             w13_bias=self.w1_bias,
             w2=w2,
-            w2_scales=a2_scale,
+            w2_scales=self.w2_scale,
             w2_bias=self.w2_bias,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
@@ -116,5 +134,22 @@ class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):
             ep_rank=self.moe_config.ep_rank,
             ep_size=self.moe_config.ep_size,
             output=output,
+            is_fp8=self.is_fp8,
         )
-        return
+
+
+class XPUExpertsFp8(XPUExperts):
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int | None = None,
+        num_dispatchers: int | None = None,
+    ):
+        super().__init__(
+            moe_config,
+            quant_config,
+            max_num_tokens,
+            num_dispatchers,
+        )
+        self.is_fp8 = True

vllm/model_executor/layers/quantization/base_config.py

@@ -168,3 +168,19 @@ class QuantizationConfig(ABC):
         Interface to update values after config initialization.
         """
         pass
+
+    def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
+        """
+        Determine if mxfp4 quantization will be used for this config.
+
+        This allows hidden_size rounding to happen before moe_config creation
+        without needing to instantiate quant_method first.
+
+        Args:
+            prefix: The layer prefix/name in the model
+            layer: The layer module
+
+        Returns:
+            True if this config uses MXFP4 quantization, False otherwise
+        """
+        return False

vllm/model_executor/layers/quantization/fp8.py

@@ -180,18 +180,9 @@ class Fp8Config(QuantizationConfig):
             weight_block_size=weight_block_size,
         )
 
-    def get_xpu_quant_method(
-        self, layer: torch.nn.Module, prefix: str
-    ) -> "QuantizeMethodBase | None":
-        raise NotImplementedError(
-            "FP8 quantization is not supported during xpu kernel migration."
-        )
-
     def get_quant_method(
         self, layer: torch.nn.Module, prefix: str
     ) -> "QuantizeMethodBase | None":
-        if current_platform.is_xpu():
-            return self.get_xpu_quant_method(layer, prefix)
         if isinstance(layer, LinearBase):
             if is_layer_skipped(
                 prefix=prefix,
@@ -300,7 +291,7 @@ class Fp8LinearMethod(LinearMethodBase):
             or envs.VLLM_TEST_FORCE_FP8_MARLIN
         )
         # Disable marlin for rocm
-        if current_platform.is_rocm():
+        if current_platform.is_rocm() or current_platform.is_xpu():
             self.use_marlin = False
         if vllm_is_batch_invariant():
             self.use_marlin = False

vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py

@@ -39,6 +39,9 @@ from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKer
 from vllm.model_executor.layers.quantization.kernels.scaled_mm.triton import (
     TritonInt8ScaledMMLinearKernel,
 )
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.xpu import (
+    XPUFP8ScaledMMLinearKernel,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
 from vllm.platforms import PlatformEnum, current_platform
 
@@ -72,6 +75,9 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
         PerTensorTorchFP8ScaledMMLinearKernel,
         ChannelWiseTorchFP8ScaledMMLinearKernel,
     ],
+    PlatformEnum.XPU: [
+        XPUFP8ScaledMMLinearKernel,
+    ],
 }
 
 _KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel)

vllm/model_executor/layers/quantization/kernels/scaled_mm/xpu.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+
+import torch
+
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import (  # noqa: E501
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+from vllm.platforms import current_platform
+
+
+class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUFP8ScaledMM only support on XPU"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
+            return False, "XPUFP8ScaledMM only support FP8 weight dtype"
+        return True, None
+
+    def __init__(
+        self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
+    ) -> None:
+        assert self.can_implement(c)[0]
+        assert self.is_supported()[0]
+        self.config = c
+        self.layer_param_names = layer_param_names
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        weight = layer.weight
+        weight_scale = layer.weight_scale
+        return torch.ops._xpu_C.fp8_gemm_w8a16(x, weight, weight_scale, bias)
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        pass

vllm/model_executor/layers/quantization/mxfp4.py

@@ -6,6 +6,7 @@ import torch
 from torch.nn.parameter import Parameter
 
 from vllm import envs
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import get_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
@@ -49,9 +50,10 @@ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
     get_padding_alignment,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
-from vllm.model_executor.utils import set_weight_attrs
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
+from vllm.platforms.rocm import on_gfx950
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.utils.import_utils import has_triton_kernels
 from vllm.utils.math_utils import round_up
@@ -75,6 +77,8 @@ class Mxfp4Backend(Enum):
     # Triton Backend
     TRITON = 6
 
+    CK = 7
+
 
 def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
     """
@@ -160,11 +164,15 @@ def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
             logger.info_once("Using Triton backend")
             return Mxfp4Backend.TRITON
     elif current_platform.is_xpu():
-        logger.info_once("Using ipex marlin backend on XPU")
+        logger.info_once("Using xpu backend on XPU")
         return Mxfp4Backend.MARLIN
-    elif current_platform.is_rocm() and has_triton_kernels():
-        logger.info_once("Using Triton backend")
-        return Mxfp4Backend.TRITON
+    elif current_platform.is_rocm():
+        if rocm_aiter_ops.is_enabled() and on_gfx950():
+            logger.info_once("Using CK MXFP4 MoE backend (Aiter ROCm)")
+            return Mxfp4Backend.CK
+        elif has_triton_kernels():
+            logger.info_once("Using Triton backend")
+            return Mxfp4Backend.TRITON
 
     return Mxfp4Backend.NONE
 
@@ -229,10 +237,15 @@ class Mxfp4Config(QuantizationConfig):
             )
         return None
 
+    def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
+        """MXFP4 config always uses MXFP4 quantization."""
+        return True
+
 
 class Mxfp4MoEMethod(FusedMoEMethodBase):
     def __init__(self, moe: FusedMoEConfig):
         super().__init__(moe)
+        self.weight_dtype = "mxfp4"
         self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
 
         self.marlin_input_dtype = None
@@ -326,6 +339,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
 
         self.intermediate_size = intermediate_size_per_partition_after_pad
         self.hidden_size = hidden_size
+        self.hidden_pad = extra_weight_attrs.get("hidden_pad", 0)
+        self.intermediate_pad = (
+            intermediate_size_per_partition_after_pad - intermediate_size_per_partition
+        )
         # Fused gate_up_proj (column parallel)
         w13_weight = torch.nn.Parameter(
             torch.zeros(
@@ -738,46 +755,99 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 layer.w2_weight_scale = torch.nn.Parameter(
                     w2_scales_interleaved, requires_grad=False
                 )
-        elif self.mxfp4_backend == Mxfp4Backend.TRITON:
-            from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
-
+        elif (
+            self.mxfp4_backend == Mxfp4Backend.TRITON
+            or self.mxfp4_backend == Mxfp4Backend.CK
+        ):
             w13_bias = layer.w13_bias.to(torch.float32)
             w2_bias = layer.w2_bias.to(torch.float32)
 
             layer.w13_bias = Parameter(w13_bias, requires_grad=False)
             layer.w2_bias = Parameter(w2_bias, requires_grad=False)
 
-            # Ideally we'd use FusedMoEModularKernel.prepare_finalize object
-            # (stored in self.fused_experts) to determine if the MoE has a
-            # batched activation format. As self.fused_experts is not
-            # initialized at this point, we resort to checking the MoE config
-            # directly.
-            is_batched_moe = self.moe.use_pplx_kernels or self.moe.use_deepep_ll_kernels
-            if is_batched_moe:
-                num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
+            if self.mxfp4_backend == Mxfp4Backend.CK:
+                w13_aiter_weight = layer.w13_weight.contiguous()
+                w13_aiter_scale = layer.w13_weight_scale.contiguous()
+                w2_aiter_weight = layer.w2_weight.contiguous()
+                w2_aiter_scale = layer.w2_weight_scale.contiguous()
+
+                e, n, k = w13_aiter_weight.shape
+                w13_aiter_weight = (
+                    w13_aiter_weight.view(e, n // 2, 2, k)
+                    .permute(0, 2, 1, 3)
+                    .contiguous()
+                    .view(e, n, k)
+                )
+                w13_aiter_scale = (
+                    w13_aiter_scale.view(e, n // 2, 2, -1)
+                    .permute(0, 2, 1, 3)
+                    .contiguous()
+                    .view(e, n, -1)
+                )
+
+                w13_aiter_weight = w13_aiter_weight.view(torch.float4_e2m1fn_x2)
+                w13_aiter_scale = w13_aiter_scale.view(-1, w13_aiter_scale.shape[-1])
+                w2_aiter_weight = w2_aiter_weight.view(torch.float4_e2m1fn_x2)
+                w2_aiter_scale = w2_aiter_scale.view(-1, w2_aiter_scale.shape[-1])
+
+                w13_weight = rocm_aiter_ops.shuffle_weight_a16w4(
+                    w13_aiter_weight, 16, True
+                )
+                w13_weight_scale = rocm_aiter_ops.shuffle_scale_a16w4(
+                    w13_aiter_scale, self.num_experts, True
+                )
+                w2_weight = rocm_aiter_ops.shuffle_weight_a16w4(
+                    w2_aiter_weight, 16, False
+                )
+                w2_weight_scale = rocm_aiter_ops.shuffle_scale_a16w4(
+                    w2_aiter_scale, self.num_experts, False
+                )
+                w13_bias = (
+                    layer.w13_bias.view(-1, n // 2, 2)
+                    .permute(0, 2, 1)
+                    .contiguous()
+                    .view(-1, n)
+                )
+                replace_parameter(layer, "w13_bias", w13_bias)
+                replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
+                replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
+                replace_parameter(layer, "w13_weight", w13_weight)
+                replace_parameter(layer, "w2_weight", w2_weight)
             else:
-                num_warps = 8
+                from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
 
-            w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
-                layer.w13_weight, layer.w13_weight_scale, num_warps
-            )
-            w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
-                layer.w2_weight, layer.w2_weight_scale, num_warps
-            )
+                # Ideally we'd use FusedMoEModularKernel.prepare_finalize object
+                # (stored in self.fused_experts) to determine if the MoE has a
+                # batched activation format. As self.fused_experts is not
+                # initialized at this point, we resort to checking the MoE config
+                # directly.
+                is_batched_moe = (
+                    self.moe.use_pplx_kernels or self.moe.use_deepep_ll_kernels
+                )
+                if is_batched_moe:
+                    num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
+                else:
+                    num_warps = 8
+                w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
+                    layer.w13_weight, layer.w13_weight_scale, num_warps
+                )
+                w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
+                    layer.w2_weight, layer.w2_weight_scale, num_warps
+                )
 
-            self.w13_precision_config = PrecisionConfig(
-                weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
-            )
-            self.w2_precision_config = PrecisionConfig(
-                weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
-            )
+                self.w13_precision_config = PrecisionConfig(
+                    weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
+                )
+                self.w2_precision_config = PrecisionConfig(
+                    weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
+                )
+                self.w13_weight = w13_weight
+                self.w2_weight = w2_weight
+                del layer.w13_weight
+                del layer.w2_weight
+                layer.w13_weight = w13_weight
+                layer.w2_weight = w2_weight
 
-            self.w13_weight = w13_weight
-            self.w2_weight = w2_weight
-            del layer.w13_weight
-            del layer.w2_weight
-            layer.w13_weight = w13_weight
-            layer.w2_weight = w2_weight
         else:
             raise ValueError(
                 f"Unsupported mxfp4_backend: {self.mxfp4_backend}: "
@@ -813,7 +883,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 w1_scale=layer.w13_weight_scale,
                 w2_scale=layer.w2_weight_scale,
             )
-        elif self.mxfp4_backend in [Mxfp4Backend.SM100_FI_MXFP4_BF16]:
+        elif (
+            self.mxfp4_backend in [Mxfp4Backend.SM100_FI_MXFP4_BF16]
+            or self.mxfp4_backend == Mxfp4Backend.CK
+        ):
             return mxfp4_w4a16_moe_quant_config(
                 w1_bias=layer.w13_bias,
                 w2_bias=layer.w2_bias,
@@ -887,6 +960,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
             or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
             or self.mxfp4_backend == Mxfp4Backend.TRITON
+            or self.mxfp4_backend == Mxfp4Backend.CK
         )
 
     def apply(
@@ -1075,6 +1149,27 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 tune_max_num_tokens=max(self.max_capture_size, 1),
             )[0]
             return trtllm_gen_output
+        elif self.mxfp4_backend == Mxfp4Backend.CK:
+            topk_weights, topk_ids = rocm_aiter_ops.fused_topk(
+                x, router_logits, layer.top_k, True
+            )
+            output = rocm_aiter_ops.fused_moe(
+                x,
+                layer.w13_weight,
+                layer.w2_weight,
+                topk_weights,
+                topk_ids,
+                activation_method=rocm_aiter_ops.get_aiter_activation_type("swiglu"),
+                quant_method=rocm_aiter_ops.get_aiter_quant_type("per_1x32"),
+                w1_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+                doweight_stage1=False,
+                hidden_pad=self.hidden_pad // 128 * 128,
+                intermediate_pad=self.intermediate_pad // 64 * 64 * 2,
+                bias1=layer.w13_bias,
+                bias2=layer.w2_bias,
+            )
+            return output
         elif self.mxfp4_backend == Mxfp4Backend.TRITON:
             from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (  # noqa: E501
                 triton_kernel_moe_forward,

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

@@ -35,6 +35,7 @@ from vllm.model_executor.layers.quantization.quark.utils import (
 )
 from vllm.model_executor.models.utils import WeightsMapper
 from vllm.platforms import current_platform
+from vllm.transformers_utils.config import get_config
 
 if TYPE_CHECKING:
     from vllm.model_executor.models.utils import WeightsMapper
@@ -59,6 +60,35 @@ class QuarkConfig(QuantizationConfig):
         self.kv_cache_group = kv_cache_group
         self.kv_cache_config = kv_cache_config
         self.pack_method = pack_method
+        self.dynamic_mxfp4_quant = False
+        self._is_global_mxfp4()
+
+    def _is_global_mxfp4(self):
+        # Check if it is MXFP4 to determine if pre-padding should be applied.
+        # This must be created during the initialization of moe.
+        global_quant_config = cast(
+            dict[str, Any], self.quant_config.get("global_quant_config")
+        )
+        weight_quant = global_quant_config.get("weight")
+        input_quant = global_quant_config.get("input_tensors")
+        self.is_global_mxfp4 = self._is_mx_fp4(
+            weight_quant=weight_quant, input_quant=input_quant
+        )
+
+    def maybe_update_config(self, model_name: str, revision: str | None = None):
+        self.hf_config = get_config(
+            model=model_name,
+            trust_remote_code=False,  # or get from model_config if available
+            revision=revision,
+            config_format="auto",
+        )
+
+        quant_config = getattr(self.hf_config, "quantization_config", None)
+        if quant_config is not None:
+            quant_dtype = quant_config["global_quant_config"]["weight"]["dtype"]
+            model_type = self.hf_config.model_type
+            if quant_dtype == "fp4" and model_type == "deepseek_v3":
+                self.dynamic_mxfp4_quant = True
 
     def get_linear_method(self) -> "QuarkLinearMethod":
         return QuarkLinearMethod(self)
@@ -108,7 +138,20 @@ class QuarkConfig(QuantizationConfig):
         if should_ignore_layer(
             prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping
         ):
-            return UnquantizedLinearMethod()
+            if (
+                "self_attn" not in prefix  # only quantize attention projections
+                or not getattr(self, "dynamic_mxfp4_quant", False)
+                or not isinstance(layer, LinearBase)  # Ignore other methods
+            ):
+                return UnquantizedLinearMethod()
+
+            scheme = self.get_scheme(
+                layer=layer,
+                layer_name=prefix,
+                dynamic_mxfp4_quant=True,
+            )
+            layer.scheme = scheme
+            return QuarkLinearMethod(self)
         if isinstance(layer, LinearBase):
             scheme = self.get_scheme(layer=layer, layer_name=prefix)
             layer.scheme = scheme
@@ -320,38 +363,83 @@ class QuarkConfig(QuantizationConfig):
         # Only symmetric weight quantization supported.
         return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
 
-    def _is_ocp_mx(
-        self,
-        weight_quant: dict[str, Any] | None,
-        input_quant: dict[str, Any] | None,
+    def _is_mx_fp4(
+        self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
     ) -> bool:
+        # Confirm weights quantized.
         # Confirm weights and input quantized.
         if weight_quant is None or input_quant is None:
+            return False
+
+        # Input and weight dtype needs to be fp4.
+        if weight_quant.get("dtype") != "fp4":
+            logger.debug("Quark model is not in MX-FP4 format: weight dtype not fp4")
+            return False
+
+        # Input and weight qscheme needs to be per group.
+        if weight_quant.get("qscheme") != "per_group":
+            logger.debug("Quark model is not in MX-FP4 format: not per_group")
+            return False
+
+        # Input and weight group size needs to be 32.
+        if weight_quant.get("group_size") != 32:
+            logger.debug("Quark model is not in MX-FP4 format: not group_size=32")
+            return False
+
+        # Activations and weight scales need to be in e8m0 format.
+        if weight_quant.get("scale_format") != "e8m0":
+            logger.debug("Quark model is not in MX-FP4 format: not scale_format e8m0")
+            return False
+
+        # Input dtype needs to be one of {'fp4', 'fp6_e2m3', 'fp8_e4m3'}.
+        if input_quant.get("dtype") not in ("fp4", "fp6_e2m3", "fp8_e4m3"):
             logger.debug(
-                "Quark model is not in OCP MX format: "
-                "weight_quant or input_quant not set"
+                "Quark model is not in MX-FP4 format: expected input dtype "
+                "to be one of {'fp4', 'fp6_e2m3', 'fp8_e4m3'}"
+            )
+            return False
+
+        return True
+
+    def _is_w_ocp_mx_a_x(
+        self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
+    ) -> bool:
+        """
+        This check returns True only if it is an OCP-MX weight quantization.
+        The activation can be any data type (e.g., FP16/BF16, FP8, or OCP-MX format).
+        The rationale for checking only the weight type is that
+        the model loading concept and process primarily concerns the weights themselves.
+        """
+        # Confirm weights quantized.
+        if weight_quant is None:
+            logger.debug(
+                "Quark model's weight quantization is incompatible with OCP_MX format: "
+                "weight_quant is not set."
             )
             return False
 
         # Input and weight qscheme needs to be per group.
-        if (
-            weight_quant.get("qscheme") != "per_group"
-            or input_quant.get("qscheme") != "per_group"
-        ):
-            logger.debug("Quark model is not in OCP MX format: not per_group")
+        if weight_quant.get("qscheme") != "per_group":
+            logger.debug(
+                "Quark model's weight quantization is incompatible with OCP MX format: "
+                "weight is not per_group."
+            )
             return False
 
         # Input and weight group size needs to be 32.
-        if weight_quant.get("group_size") != 32 or input_quant.get("group_size") != 32:
-            logger.debug("Quark model is not in OCP MX format: not group_size=32")
+        if weight_quant.get("group_size") != 32:
+            logger.debug(
+                "Quark model's weight quantization is incompatible with OCP MX format: "
+                "group_size of weight is not 32."
+            )
             return False
 
         # Activations and weight scales need to be in e8m0 format.
-        if (
-            weight_quant.get("scale_format") != "e8m0"
-            or input_quant.get("scale_format") != "e8m0"
-        ):
-            logger.debug("Quark model is not in OCP MX format: not scale_format e8m0")
+        if weight_quant.get("scale_format") != "e8m0":
+            logger.debug(
+                "Quark model's weight quantization is incompatible with OCP MX format: "
+                "scale_format of weight is not e8m0."
+            )
             return False
 
         # Input and weight dtypes need to be any of fp4,
@@ -360,14 +448,31 @@ class QuarkConfig(QuantizationConfig):
             "fp4",
             "fp6_e3m2",
             "fp6_e2m3",
-        } or input_quant.get("dtype") not in {"fp4", "fp6_e3m2", "fp6_e2m3"}:
+        }:
             logger.debug(
-                "Quark model is not in OCP MX format: dtype not fp4, fp6_e3m2, fp6_e2m3"
+                "Quark model's weight quantization is incompatible with OCP MX format: "
+                "dtype is not in {fp4, fp6_e3m2, fp6_e2m3}."
             )
             return False
 
         return True
 
+    def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
+        """
+        For Quark, determine if it's OCP MXFP4 by checking config directly.
+        This allows hidden_size rounding to happen before moe_config creation.
+        """
+        layer_quant_config = self._find_matched_config(prefix, layer)
+        weight_config = layer_quant_config.get("weight")
+        input_config = layer_quant_config.get("input_tensors")
+
+        return (
+            self._is_w_ocp_mx_a_x(weight_config, input_config)
+            and weight_config is not None
+            and weight_config.get("dtype") == "fp4"
+            and getattr(torch, "float4_e2m1fn_x2", None) is not None
+        )
+
     def _find_matched_config(
         self, layer_name: str, module: torch.nn.Module
     ) -> dict[str, Any]:
@@ -419,7 +524,9 @@ class QuarkConfig(QuantizationConfig):
             )
             return global_quant_config
 
-    def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkScheme":
+    def _get_scheme_from_config(
+        self, config: dict[str, Any], dynamic_mxfp4_quant: bool = False
+    ) -> "QuarkScheme":
         if config.get("output_tensors") or config.get("bias"):
             raise NotImplementedError(
                 "Currently, Quark models with output_tensors "
@@ -441,8 +548,10 @@ class QuarkConfig(QuantizationConfig):
                 is_static_input_scheme=True,
                 input_symmetric=input_config.get("symmetric"),
             )
-        elif self._is_ocp_mx(weight_config, input_config):
-            return QuarkOCP_MX(weight_config, input_config)
+        elif self._is_w_ocp_mx_a_x(weight_config, input_config):
+            return QuarkOCP_MX(
+                weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
+            )
 
         raise NotImplementedError(
             "No quark compatible scheme was found. "
@@ -450,11 +559,15 @@ class QuarkConfig(QuantizationConfig):
             f"Input config: {input_config}"
         )
 
-    def get_scheme(self, layer: torch.nn.Module, layer_name: str) -> "QuarkScheme":
+    def get_scheme(
+        self, layer: torch.nn.Module, layer_name: str, dynamic_mxfp4_quant: bool = False
+    ) -> "QuarkScheme":
         layer_quant_config = self._find_matched_config(layer_name, layer)
 
         # Find the quant_scheme
-        scheme = self._get_scheme_from_config(layer_quant_config)
+        scheme = self._get_scheme_from_config(
+            layer_quant_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
+        )
         # Raise error if device does not support the scheme
         # (e.g. fp8 needs ada lovelace)
         self._check_scheme_supported(scheme.get_min_capability())

vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py

@@ -24,7 +24,12 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
     OCP_MX_BLOCK_SIZE,
     OCP_MX_Scheme,
 )
-from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
+from vllm.model_executor.parameter import (
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
+    PackedvLLMParameter,
+)
+from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 
 from .quark_scheme import QuarkScheme
@@ -169,13 +174,16 @@ except (ImportError, AttributeError, RuntimeError):
 
 class QuarkOCP_MX(QuarkScheme):
     def __init__(
-        self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any]
+        self,
+        weight_quant_spec: dict[str, Any],
+        input_quant_spec: dict[str, Any],
+        dynamic_mxfp4_quant: bool = False,
     ):
         self.out_dtype = torch.get_default_dtype()
         self.qscheme = "per_group"
         self.weight_quant_spec = weight_quant_spec
         self.input_quant_spec = input_quant_spec
-
+        self.dynamic_mxfp4_quant = dynamic_mxfp4_quant
         self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp")
         self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp")
 
@@ -269,7 +277,13 @@ class QuarkOCP_MX(QuarkScheme):
                 layer.weight_scale.data, requires_grad=False
             )
         else:
-            if self.rocm_use_aiter_fp4_asm_gemm:
+            if self.dynamic_mxfp4_quant:
+                w_q, w_s = dynamic_mxfp4_quant(layer.weight)
+                layer.weight_scale = torch.nn.Parameter(
+                    w_s.T.contiguous(), requires_grad=False
+                )
+                layer.weight = torch.nn.Parameter(w_q, requires_grad=False)
+            elif self.rocm_use_aiter_fp4_asm_gemm:
                 # shuffle weight scale
                 weight_scale_shuffle = layer.weight_scale.data
                 sm, sn = weight_scale_shuffle.shape
@@ -302,36 +316,51 @@ class QuarkOCP_MX(QuarkScheme):
         weight_loader: Callable,
         **kwargs,
     ):
-        output_size_per_partition = sum(output_partition_sizes)
-        layer.logical_widths = output_partition_sizes
+        if self.dynamic_mxfp4_quant:
+            weight = ModelWeightParameter(
+                data=torch.empty(
+                    sum(output_partition_sizes),
+                    input_size_per_partition,
+                    dtype=params_dtype,
+                ),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
 
-        # WEIGHT
-        weight = PackedvLLMParameter(
-            data=torch.empty(
-                output_size_per_partition,
-                self.get_packed_dim(input_size_per_partition, self.weight_dtype),
-                dtype=torch.uint8,
-            ),
-            input_dim=1,
-            output_dim=0,
-            packed_dim=1,
-            packed_factor=self.packed_factor,
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight", weight)
+            layer.register_parameter("weight", weight)
+            set_weight_attrs(weight, kwargs)
+        else:
+            output_size_per_partition = sum(output_partition_sizes)
+            layer.logical_widths = output_partition_sizes
 
-        # WEIGHT SCALE
-        weight_scale = GroupQuantScaleParameter(
-            data=torch.empty(
-                output_size_per_partition,
-                input_size_per_partition // OCP_MX_BLOCK_SIZE,
-                dtype=torch.uint8,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight_scale", weight_scale)
+            # WEIGHT
+            weight = PackedvLLMParameter(
+                data=torch.empty(
+                    output_size_per_partition,
+                    self.get_packed_dim(input_size_per_partition, self.weight_dtype),
+                    dtype=torch.uint8,
+                ),
+                input_dim=1,
+                output_dim=0,
+                packed_dim=1,
+                packed_factor=self.packed_factor,
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter("weight", weight)
+
+            # WEIGHT SCALE
+            weight_scale = GroupQuantScaleParameter(
+                data=torch.empty(
+                    output_size_per_partition,
+                    input_size_per_partition // OCP_MX_BLOCK_SIZE,
+                    dtype=torch.uint8,
+                ),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter("weight_scale", weight_scale)
 
     def apply_weights(
         self,

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

@@ -20,26 +20,44 @@ SUPPORTED_OCP_MX_DTYPES = {"mxfp4", "mxfp6_e3m2", "mxfp6_e2m3"}
 
 
 class OCP_MX_Scheme(str, Enum):
+    w_mxfp4 = "w_mxfp4"
     w_mxfp4_a_mxfp4 = "w_mxfp4_a_mxfp4"
     w_mxfp4_a_mxfp6_e3m2 = "w_mxfp4_a_mxfp6_e3m2"
     w_mxfp4_a_mxfp6_e2m3 = "w_mxfp4_a_mxfp6_e2m3"
+    w_mxfp4_a_fp8 = "w_mxfp4_a_fp8"
+    w_mxfp6_e3m2 = "w_mxfp6_e3m2"
     w_mxfp6_e3m2_a_mxfp6_e3m2 = "w_mxfp6_e3m2_a_mxfp6_e3m2"
+    w_mxfp6_e3m2_a_fp8 = "w_mxfp6_e3m2_a_fp8"
+    w_mxfp6_e2m3 = "w_mxfp6_e2m3"
     w_mxfp6_e2m3_a_mxfp6_e2m3 = "w_mxfp6_e2m3_a_mxfp6_e2m3"
+    w_mxfp6_e2m3_a_fp8 = "w_mxfp6_e2m3_a_fp8"
 
     @classmethod
     def from_quant_dtype(cls, input_dtype: str | None, weight_dtype: str | None):
-        if input_dtype not in OCP_MX_DTYPES or weight_dtype not in OCP_MX_DTYPES:
+        if input_dtype not in OCP_MX_DTYPES and weight_dtype not in OCP_MX_DTYPES:
             return None
+        elif input_dtype is None and weight_dtype == "mxfp4":
+            return cls.w_mxfp4
+        elif input_dtype is None and weight_dtype == "mxfp6_e3m2":
+            return cls.w_mxfp6_e3m2
+        elif input_dtype is None and weight_dtype == "mxfp6_e2m3":
+            return cls.w_mxfp6_e2m3
         elif input_dtype == "mxfp4" and weight_dtype == "mxfp4":
             return cls.w_mxfp4_a_mxfp4
         elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp4":
             return cls.w_mxfp4_a_mxfp6_e3m2
         elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp4":
             return cls.w_mxfp4_a_mxfp6_e2m3
+        elif input_dtype == "fp8" and weight_dtype == "mxfp4":
+            return cls.w_mxfp4_a_fp8
         elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp6_e3m2":
             return cls.w_mxfp6_e3m2_a_mxfp6_e3m2
+        elif input_dtype == "fp8" and weight_dtype == "mxfp6_e3m2":
+            return cls.w_mxfp6_e3m2_a_fp8
         elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp6_e2m3":
             return cls.w_mxfp6_e2m3_a_mxfp6_e2m3
+        elif input_dtype == "fp8" and weight_dtype == "mxfp6_e2m3":
+            return cls.w_mxfp6_e2m3_a_fp8
         else:
             logger.warning(
                 "input_dtype='%s' and"

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -20,7 +20,7 @@ from vllm.v1.worker.workspace import current_workspace_manager
 if current_platform.is_cuda_alike():
     from vllm import _custom_ops as ops
 elif current_platform.is_xpu():
-    from vllm._ipex_ops import ipex_ops as ops
+    from vllm._xpu_ops import xpu_ops as ops
 
 logger = init_logger(__name__)
 

vllm/model_executor/models/gpt_oss.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Iterable
+import typing
+from collections.abc import Callable, Iterable
 
 import torch
 import torch.distributed as dist
@@ -25,13 +26,17 @@ from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.utils import rocm_unquantized_gemm
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
 from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
@@ -98,6 +103,7 @@ class OAIAttention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=self.num_attention_heads,
             total_num_kv_heads=self.num_key_value_heads,
+            bias=True,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
         )
@@ -105,6 +111,7 @@ class OAIAttention(nn.Module):
         self.o_proj = RowParallelLinear(
             input_size=self.num_attention_heads * self.head_dim,
             output_size=self.hidden_size,
+            bias=True,
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
@@ -306,6 +313,19 @@ class GptOssModel(nn.Module):
             return x, aux_hidden_states
         return x
 
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        # Params for weights, weight scales, activation scales
+        # (param_name, weight_name, expert_id, shard_id)
+        # NOTE: this is only used for quark.
+        return FusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="w1",
+            ckpt_down_proj_name="w2",
+            ckpt_up_proj_name="w3",
+            num_experts=self.config.num_local_experts,
+            num_redundant_experts=0,
+        )
+
     def _load_weights_mxfp4(
         self,
         ep_rank_end: int,
@@ -318,7 +338,6 @@ class GptOssModel(nn.Module):
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
-        mxfp4_block = 32
         use_ep = self.parallel_config.enable_expert_parallel
         num_experts = self.config.num_local_experts
 
@@ -333,9 +352,11 @@ class GptOssModel(nn.Module):
         )
 
         intermediate_size = self.config.intermediate_size
-        intermediate_size_block = intermediate_size // mxfp4_block
+        intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE
         per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size)
-        per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block
+        per_rank_intermediate_size = (
+            per_rank_intermediate_size_block * OCP_MX_BLOCK_SIZE
+        )
 
         # Calculate common slicing bounds for current rank
         tp_rank_start = tp_rank * per_rank_intermediate_size
@@ -370,7 +391,9 @@ class GptOssModel(nn.Module):
                     narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
                 else:
                     narrow_weight = weight[
-                        ..., tp_rank_start // mxfp4_block : tp_rank_end // mxfp4_block
+                        ...,
+                        tp_rank_start // OCP_MX_BLOCK_SIZE : tp_rank_end
+                        // OCP_MX_BLOCK_SIZE,
                     ]
 
                 param = params_dict[name]
@@ -495,6 +518,449 @@ class GptOssModel(nn.Module):
             loaded_params.add(name)
         return loaded_params
 
+    def _load_weights_quark(
+        self,
+        ep_rank_end: int,
+        ep_rank_start: int,
+        heads_per_rank: int,
+        head_start: int,
+        weights: Iterable[tuple[str, torch.Tensor]],
+        stacked_params_mapping: list[tuple[str, ...]],
+    ) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        use_ep = self.parallel_config.enable_expert_parallel
+        num_experts = self.config.num_local_experts
+
+        if use_ep:
+            tp_rank = get_tensor_model_parallel_rank()
+            tp_size = get_tensor_model_parallel_world_size()
+        else:
+            tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
+                tp_size=get_tensor_model_parallel_world_size(),
+                dp_size=get_dp_group().world_size,
+                dp_rank=get_dp_group().rank_in_group,
+                pcp_size=get_pcp_group().world_size,
+                pcp_rank=get_pcp_group().rank_in_group,
+            )
+
+        def _get_moe_weight_dtype(layer_id: int = 0) -> str | None:
+            """Helper function to get MoE quantization weight dtype.
+
+            Args:
+                layer_id: Layer index to check (default 0, as all layers should
+                        have the same quantization method)
+
+            Returns:
+                Weight dtype string (e.g., "mxfp4", "fp8") or None if not available
+            """
+            if hasattr(self.layers[layer_id].mlp.experts.quant_method, "weight_dtype"):
+                return self.layers[layer_id].mlp.experts.quant_method.weight_dtype
+            return None
+
+        intermediate_size = self.config.intermediate_size
+
+        moe_weight_dtype = _get_moe_weight_dtype(layer_id=0)
+
+        if moe_weight_dtype == "mxfp4":
+            # MXFP4 requires OCP_MX_BLOCK_SIZE alignment
+            intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE
+            per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size)
+            per_rank_intermediate_size = (
+                per_rank_intermediate_size_block * OCP_MX_BLOCK_SIZE
+            )
+        else:
+            # FP8 and other formats don't need alignment
+            per_rank_intermediate_size = cdiv(intermediate_size, tp_size)
+
+        tp_rank_start = tp_rank * per_rank_intermediate_size
+        tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size)
+        expert_params_mapping = self.get_expert_mapping()
+        for name, loaded_weight in weights:
+            if is_pp_missing_parameter(name, self):
+                continue
+
+            layer_id, expert_id, fused_name = None, None, None
+            moe_quant_method = None
+            if "experts" in name:
+                parts = name.split(".")
+                ids = [s for s in parts if s.isdigit()]
+
+                # for amd-quark format that each expert is seperated
+                # need to extract the parameter name with experts fused.
+                # example model: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8
+                if len(ids) == 2:
+                    layer_id, expert_id = int(ids[0]), int(ids[-1])
+                    parts.pop(len(parts) - 1 - parts[::-1].index(str(expert_id)))
+                    fused_name = ".".join(parts)
+
+                # for openai mxfp4 format that all experts are combined
+                # no need to extract the parameter name with experts fused.
+                # models: openai/gpt-oss-20b, openai/gpt-oss-120b
+                elif len(ids) == 1:
+                    layer_id, expert_id = int(ids[0]), None
+                    fused_name = name
+
+                else:
+                    raise NameError(
+                        f"Layer {name} contains more than 2 numeric indices. This is "
+                        "an unexpected condition. Please open an issue if encountered."
+                    )
+
+                moe_quant_method = _get_moe_weight_dtype(layer_id=layer_id)
+
+            def kv_cache_scale_loader(
+                quant_config: QuantizationConfig,
+                name: str,
+                params_dict: dict[str, typing.Any],
+                weight: torch.Tensor,
+                default_weight_loader: Callable[..., None],
+                loaded_params: set[str],
+            ) -> tuple[bool, set[str]]:
+                """
+                Load KV cache output scales.
+                Returns:
+                    Tuple of (bool, set):
+                    - bool: True if KV-cache scale was loaded into loaded_params
+                    - set: Updated set of loaded_params if True else the original set
+                """
+                # load explicit cached KV output scale from quant_config
+                if quant_config is not None and (
+                    scale_name := quant_config.get_cache_scale(name)
+                ):
+                    param = params_dict[scale_name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    if weight.numel() != 1:
+                        raise ValueError(
+                            f"KV cache scale '{scale_name}' is expected to be a "
+                            f"scalar, but got a tensor of shape {weight.shape}."
+                        )
+                    # Ensure weight is a scalar before passing to loader.
+                    weight_loader(param, weight.flatten()[0])
+                    loaded_params.add(scale_name)
+                    return True, loaded_params
+
+                return False, loaded_params
+
+            load_kv_cache_scale_completed, loaded_params = kv_cache_scale_loader(
+                self.quant_config,
+                name,
+                params_dict,
+                loaded_weight,
+                default_weight_loader,
+                loaded_params,
+            )
+            if load_kv_cache_scale_completed:
+                continue
+
+            if (
+                all(key in name for key in ["input_scale", "mlp.experts"])
+                and expert_id is not None
+            ):
+                assert loaded_weight.numel() == 1
+                expert_data = params_dict[fused_name].data[expert_id]
+                expert_data.copy_(loaded_weight)
+                loaded_params.add(fused_name)
+                continue
+
+            # Unified handler for mxfp4 weights and scales
+            elif moe_quant_method == "mxfp4" and any(
+                name.endswith(suffix)
+                for suffix in [
+                    ".w13_weight_scale",
+                    ".w2_weight_scale",
+                    ".w13_weight",
+                    ".w2_weight",
+                ]
+            ):
+                is_w13 = ".w13_" in name
+                is_scale = "_scale" in name
+
+                # Reshape weight for mxfp4 if needed (not for scales)
+                if not is_scale and expert_id is None:
+                    if is_w13:
+                        if loaded_weight.dim() < 3:
+                            raise ValueError(
+                                f"Expected w13_weight to have at least 3 "
+                                f"dimensions, got shape "
+                                f"{loaded_weight.shape}"
+                            )
+                        if loaded_weight.shape[0] != num_experts:
+                            raise ValueError(
+                                f"Expected w13_weight first dimension to be "
+                                f"{num_experts}, got "
+                                f"{loaded_weight.shape[0]}"
+                            )
+                        loaded_weight = loaded_weight.view(
+                            num_experts, 2 * intermediate_size, -1
+                        ).contiguous()
+                    else:
+                        if loaded_weight.dim() < 3:
+                            raise ValueError(
+                                f"Expected w2_weight to have at least 3 "
+                                f"dimensions, got shape "
+                                f"{loaded_weight.shape}"
+                            )
+                        if loaded_weight.shape[0] != num_experts:
+                            raise ValueError(
+                                f"Expected w2_weight first dimension to be "
+                                f"{num_experts}, got "
+                                f"{loaded_weight.shape[0]}"
+                            )
+                        loaded_weight = loaded_weight.view(
+                            num_experts, -1, intermediate_size // 2
+                        ).contiguous()
+
+                if use_ep:
+                    sliced_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
+                else:
+                    if is_w13:
+                        if expert_id is None:
+                            sliced_weight = loaded_weight[
+                                :, 2 * tp_rank_start : 2 * tp_rank_end, ...
+                            ]
+                        else:
+                            sliced_weight = loaded_weight[
+                                2 * tp_rank_start : 2 * tp_rank_end, ...
+                            ]
+                    else:
+                        if is_scale:
+                            sliced_weight = loaded_weight[
+                                ...,
+                                tp_rank_start // OCP_MX_BLOCK_SIZE : tp_rank_end
+                                // OCP_MX_BLOCK_SIZE,
+                            ]
+                        else:
+                            sliced_weight = loaded_weight[
+                                ..., tp_rank_start // 2 : tp_rank_end // 2
+                            ]
+
+                # NOTE(rob): because gpt-oss ckpt has "unique" structure with
+                # fused gate_up_proj fused on disk, we cannot use the existing
+                # weight loaders without added complexity, so just do the
+                # direct load here.
+                param = params_dict[fused_name]
+                expert_data = param.data[expert_id]
+                dim1 = sliced_weight.shape[0]
+                dim2 = sliced_weight.shape[1]
+                expert_data.data[:dim1, :dim2].copy_(sliced_weight)
+                loaded_params.add(fused_name)
+                continue
+
+            elif name.endswith(".w13_weight") and moe_quant_method == "fp8":
+                if use_ep:
+                    narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
+                else:
+                    if expert_id is None:
+                        narrow_weight = loaded_weight[
+                            :, 2 * tp_rank_start : 2 * tp_rank_end, :
+                        ]
+                    else:
+                        narrow_weight = loaded_weight[
+                            2 * tp_rank_start : 2 * tp_rank_end, :
+                        ]
+
+                assert fused_name is not None
+                param = params_dict[fused_name]
+
+                if expert_id is None:
+                    param.data.copy_(narrow_weight)
+                else:
+                    param.data[expert_id].copy_(narrow_weight)
+
+                loaded_params.add(fused_name)
+                continue
+
+            elif name.endswith(".w13_weight_scale") and moe_quant_method == "fp8":
+                assert fused_name is not None
+                param = params_dict[fused_name]
+
+                # Check if this is per-channel or per-tensor scale
+                if loaded_weight.numel() > 1 and loaded_weight.dim() == 1:
+                    if use_ep:
+                        narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
+                    else:
+                        narrow_weight = loaded_weight[
+                            2 * tp_rank_start : 2 * tp_rank_end
+                        ]
+                else:
+                    narrow_weight = loaded_weight
+
+                if expert_id is None:
+                    param.data.copy_(narrow_weight)
+                else:
+                    param.data[expert_id].copy_(narrow_weight)
+
+                loaded_params.add(fused_name)
+                continue
+
+            elif name.endswith(".w13_input_scale") and moe_quant_method == "fp8":
+                assert fused_name is not None
+                param = params_dict[fused_name]
+
+                if expert_id is None:
+                    param.data.copy_(loaded_weight)
+                else:
+                    param.data[expert_id].copy_(loaded_weight)
+
+                loaded_params.add(fused_name)
+                continue
+
+            elif name.endswith(".w2_weight") and moe_quant_method == "fp8":
+                if use_ep:
+                    narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
+                else:
+                    if expert_id is None:
+                        narrow_weight = loaded_weight[..., tp_rank_start:tp_rank_end]
+                    else:
+                        narrow_weight = loaded_weight[..., tp_rank_start:tp_rank_end]
+
+                assert fused_name is not None
+                param = params_dict[fused_name]
+
+                if expert_id is None:
+                    param.data.copy_(narrow_weight)
+                else:
+                    param.data[expert_id].copy_(narrow_weight)
+
+                loaded_params.add(fused_name)
+                continue
+
+            elif name.endswith(".w2_weight_scale") and moe_quant_method == "fp8":
+                assert fused_name is not None
+                param = params_dict[fused_name]
+
+                if use_ep:
+                    narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
+                else:
+                    narrow_weight = loaded_weight
+
+                if expert_id is None:
+                    param.data.copy_(narrow_weight)
+                else:
+                    param.data[expert_id].copy_(narrow_weight)
+
+                loaded_params.add(fused_name)
+                continue
+
+            # Unified handler for bias loading (w13_bias and w2_bias)
+            elif name.endswith(".w13_bias") or name.endswith(".w2_bias"):
+                is_w13_bias = name.endswith(".w13_bias")
+
+                if use_ep:
+                    sliced_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
+                else:
+                    if is_w13_bias:
+                        if expert_id is None:
+                            sliced_weight = loaded_weight[
+                                :, 2 * tp_rank_start : 2 * tp_rank_end
+                            ]
+                        else:
+                            sliced_weight = loaded_weight[
+                                2 * tp_rank_start : 2 * tp_rank_end
+                            ]
+                    else:
+                        sliced_weight = loaded_weight
+                        if tp_rank != 0:
+                            sliced_weight = sliced_weight.zero_()
+
+                # NOTE(rob): because gpt-oss ckpt has "unique" structure with
+                # fused gate_up_proj fused on disk, we cannot use the existing
+                # weight loaders without added complexity, so just do the
+                # direct load here.
+                assert fused_name is not None
+                param = params_dict[fused_name]
+                expert_data = param.data[expert_id]
+                dim1 = sliced_weight.shape[0]
+                expert_data.data[:dim1].copy_(sliced_weight)
+                loaded_params.add(fused_name)
+                continue
+
+            elif "sinks" in name:
+                # Handle attention sinks (distributed across ranks)
+                param = params_dict[name]
+                narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank)
+                param.data.copy_(narrow_weight)
+                loaded_params.add(name)
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                # Skip non-stacked layers and experts (experts handled below).
+                if weight_name not in name:
+                    continue
+                # We have mlp.experts[0].gate_proj in the checkpoint.
+                # Since we handle the experts below in expert_params_mapping,
+                # we need to skip here BEFORE we update the name, otherwise
+                # name will be updated to mlp.experts[0].gate_up_proj, which
+                # will then be updated below in expert_params_mapping
+                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
+                if ("mlp.experts." in name) and name not in params_dict:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                if name.endswith("scale"):
+                    # Remapping the name of FP8 kv-scale.
+                    name = maybe_remap_kv_scale_name(name, params_dict)
+                    if name is None:
+                        continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(name)
+                break
+            else:
+                for mapping in expert_params_mapping:
+                    # Anyway, this is an expert weight and should not be
+                    # attempted to load as other weights later
+                    param_name, weight_name, mapping_expert_id, shard_id = mapping
+                    weight_name = (
+                        weight_name[:-1] if weight_name.endswith(".") else weight_name
+                    )
+
+                    if weight_name not in name:
+                        continue
+
+                    param = params_dict[fused_name]
+                    # We should ask the weight loader to return success or not
+                    # here since otherwise we may skip experts with other
+                    # available replicas.
+                    weight_loader = typing.cast(
+                        Callable[..., bool], param.weight_loader
+                    )
+                    # Use checkpoint's expert_id for quark format (when expert_id
+                    # is extracted from weight name), otherwise use mapping's expert_id
+                    actual_expert_id = (
+                        expert_id if expert_id is not None else mapping_expert_id
+                    )
+                    success = weight_loader(
+                        param,
+                        loaded_weight,
+                        fused_name,
+                        shard_id=shard_id,
+                        expert_id=actual_expert_id,
+                        return_success=True,
+                    )
+                    if success:
+                        name = fused_name
+                        loaded_params.add(name)
+                        break
+                else:
+                    if name not in params_dict:
+                        continue
+                    param = params_dict[name]
+                    weight_loader = getattr(
+                        param, "weight_loader", default_weight_loader
+                    )
+                    weight_loader(param, loaded_weight)
+
+                loaded_params.add(name)
+        return loaded_params
+
     def _load_weights_other(
         self,
         ep_rank_end: int,
@@ -635,6 +1101,7 @@ class GptOssModel(nn.Module):
             if hasattr(self.config, "quantization_config")
             else None
         )
+
         if quant_method == "mxfp4":
             return self._load_weights_mxfp4(
                 ep_rank_end,
@@ -644,6 +1111,15 @@ class GptOssModel(nn.Module):
                 weights,
                 stacked_params_mapping,
             )
+        elif quant_method == "quark":
+            return self._load_weights_quark(
+                ep_rank_end,
+                ep_rank_start,
+                heads_per_rank,
+                head_start,
+                weights,
+                stacked_params_mapping,
+            )
         else:
             return self._load_weights_other(
                 ep_rank_end,
@@ -676,6 +1152,15 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
             # MoE Bias
             ".gate_up_proj_bias": ".w13_bias",
             ".down_proj_bias": ".w2_bias",
+            # For quark format
+            ".gate_up_proj.weight": ".w13_weight",
+            ".gate_up_proj.weight_scale": ".w13_weight_scale",
+            ".gate_up_proj.bias": ".w13_bias",
+            ".gate_up_proj.input_scale": ".w13_input_scale",
+            ".down_proj.weight": ".w2_weight",
+            ".down_proj.weight_scale": ".w2_weight_scale",
+            ".down_proj.bias": ".w2_bias",
+            ".down_proj.input_scale": ".w2_input_scale",
         },
     )
 
@@ -725,18 +1210,6 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
         logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
 
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        # Params for weights, weight scales, activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        return FusedMoE.make_expert_params_mapping(
-            self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_local_experts,
-            num_redundant_experts=0,
-        )
-
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(
             self,

vllm/platforms/cpu.py

@@ -338,7 +338,6 @@ class CpuPlatform(Platform):
                 ld_preload_str += pytorch_libgomp_so
                 os.environ["LD_PRELOAD"] = ld_preload_str
 
-        # To hint IPEX uses shared memory based AllReduce
         os.environ["LOCAL_WORLD_SIZE"] = str(
             vllm_config.parallel_config.tensor_parallel_size
         )

vllm/platforms/rocm.py

@@ -360,6 +360,11 @@ class RocmPlatform(Platform):
                 vllm_config is not None
                 and vllm_config.attention_config.use_prefill_decode_attention
             ):
+                logger.warning_once(
+                    "use_prefill_decode_attention is deprecated and will be removed in "
+                    "future releases. "
+                    "Use --attention_config.backend to select the desired backend"
+                )
                 logger.info("Using Rocm Attention backend.")
                 return AttentionBackendEnum.ROCM_ATTN.get_path()
 
@@ -373,9 +378,9 @@ class RocmPlatform(Platform):
                 logger.info("Using Aiter Flash Attention backend.")
                 return AttentionBackendEnum.ROCM_AITER_FA.get_path()
 
-            # Default: Triton Unified Attention
-            logger.info("Using Triton Attention backend.")
-            return AttentionBackendEnum.TRITON_ATTN.get_path()
+            # Default: ROCm split Attention
+            logger.info("Using ROCm Attention backend.")
+            return AttentionBackendEnum.ROCM_ATTN.get_path()
 
         raise RuntimeError(
             f"Attention backend {selected_backend.name} is not supported on "

vllm/utils/print_utils.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def print_embeddings(embeds: list[float]):
+    embeds_trimmed = (str(embeds[:4])[:-1] + ", ...]") if len(embeds) > 4 else embeds
+    print(f"Embeddings: {embeds_trimmed} (size={len(embeds)})")

vllm/v1/attention/backend.py

@@ -723,6 +723,33 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
         """
         return False
 
+    def fused_rope_kvcache_supported(self):
+        """
+        Does this attention implementation support RoPE+KVCache fusion.
+        This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
+        with the KV cache update for implementations that support it.
+        """
+        return False
+
+    def do_rope_and_kv_cache_update(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        positions: torch.Tensor,
+        cos_sin_cache: torch.Tensor,
+        is_neox: bool,
+        kv_cache: torch.Tensor,
+        layer_slot_mapping: torch.Tensor,
+    ):
+        """
+        If `fused_rope_kvcache_supported` returns True, this method will be called
+        by torch.ops.vllm.fused_rope_and_unified_kv_cache_update
+        to perform the inplace RoPE and KV cache update.
+        """
+        raise NotImplementedError
+
 
 class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
     """MLA attention implementation with forward_mqa and forward_mha methods."""

vllm/v1/attention/backends/fa_utils.py

@@ -23,12 +23,11 @@ if current_platform.is_cuda():
 
 elif current_platform.is_xpu():
     from vllm import _custom_ops as ops
+    from vllm._xpu_ops import xpu_ops
 
     reshape_and_cache_flash = ops.reshape_and_cache_flash
-    from vllm._ipex_ops import ipex_ops
-
-    flash_attn_varlen_func = ipex_ops.flash_attn_varlen_func  # type: ignore[assignment]
-    get_scheduler_metadata = ipex_ops.get_scheduler_metadata  # type: ignore[assignment]
+    flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func  # type: ignore[assignment]
+    get_scheduler_metadata = xpu_ops.get_scheduler_metadata  # type: ignore[assignment]
 elif current_platform.is_rocm():
     try:
         from flash_attn import flash_attn_varlen_func  # type: ignore[no-redef]
@@ -153,7 +152,7 @@ def is_flash_attn_varlen_func_available() -> bool:
 
     Platform-specific sources:
     - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func
-    - XPU: ipex_ops.flash_attn_varlen_func
+    - XPU: xpu_ops.flash_attn_varlen_func
     - ROCm: upstream flash_attn.flash_attn_varlen_func (if available)
 
     Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py)