[ROCm][CI] Pin TorchCodec to v0.10.0 for ROCm compatibility (#34447 )

Signed-off-by: Andreas Karatzas <akaratza@amd.com> (cherry picked from commit 4c078fa546) (cherry picked from commit a976961fb7)
[ci] Use the right tag for CPU arm64 image (#34915 )
2026-02-24 20:30:22 -08:00 · 2026-02-24 20:30:13 -08:00 · 2026-02-24 20:30:02 -08:00 · 2026-02-24 20:28:53 -08:00 · 2026-02-24 20:28:45 -08:00 · 2026-02-12 02:37:52 -08:00
73 changed files with 650 additions and 3797 deletions
@@ -3,7 +3,6 @@ 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 $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
@@ -42,7 +41,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: []
@@ -14,7 +14,7 @@ BUILDKITE_COMMIT=$3
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
 # skip build if image already exists
-if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
+if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu) ]]; then
  echo "Image not found, proceeding with build..."
 else
  echo "Image found"
@@ -24,10 +24,10 @@ fi
 # build
 docker build --file docker/Dockerfile.cpu \
  --build-arg max_jobs=16 \
-  --build-arg buildkite_commit=$BUILDKITE_COMMIT \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
-  --tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
+  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu \
  --target vllm-test \
  --progress plain .
 # push
-docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64-cpu
@@ -248,8 +248,8 @@ steps:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"
-  - group: "Publish wheels"
+  - group: "Publish release artifacts"
-    key: "publish-wheels"
+    key: "publish-release-artifacts"
    steps:
      - block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
        key: block-upload-release-wheels
@@ -265,6 +265,27 @@ steps:
          queue: small_cpu_queue_postmerge
        commands:
          - "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
      - block: "Confirm update release images to DockerHub"
        key: block-update-release-images-dockerhub
        depends_on:
          - input-release-version
          - annotate-release-workflow
      - label: "Publish release images to DockerHub"
        depends_on:
          - block-update-release-images-dockerhub
        agents:
          queue: small_cpu_queue_postmerge
        commands:
          - "bash .buildkite/scripts/push-release-images-dockerhub.sh"
        plugins:
          - docker-login#v3.0.0:
              username: vllmbot
              password-env: DOCKERHUB_TOKEN
        env:
          DOCKER_BUILDKIT: "1"
          DOCKERHUB_USERNAME: "vllmbot"
  # =============================================================================
  # ROCm Release Pipeline (x86_64 only)
@@ -0,0 +1,98 @@
 #!/bin/bash
 set -ex
 RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null | sed 's/^v//')
 if [ -z "${RELEASE_VERSION}" ]; then
    echo "RELEASE_VERSION is not set"
    exit 1
 fi
 aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
 # Download images:
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
 docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
 docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
 docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
 # Tag and push images:
 ## CUDA
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
 docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
 docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 docker push vllm/vllm-openai:latest-x86_64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
 docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
 docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
 docker push vllm/vllm-openai:latest-x86_64-cu130
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
 docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker push vllm/vllm-openai:latest-aarch64
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
 docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
 docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
 docker push vllm/vllm-openai:latest-aarch64-cu130
 docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
 ## ROCm
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
 docker push vllm/vllm-openai-rocm:latest
 docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
 docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
 docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
 docker push vllm/vllm-openai-rocm:latest-base
 docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
 ## CPU
 docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
 docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
 docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
 docker push vllm/vllm-openai-cpu:latest-x86_64
 docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
 docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
 docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
 docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
 docker push vllm/vllm-openai-cpu:latest-arm64
 docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
 # Create multi-arch manifest:
 docker manifest rm vllm/vllm-openai:latest
 docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
 docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
 docker manifest push vllm/vllm-openai:latest
 docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
 docker manifest rm vllm/vllm-openai:latest-cu130
 docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
 docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
 docker manifest push vllm/vllm-openai:latest-cu130
 docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
 docker manifest rm vllm/vllm-openai-cpu:latest || true
 docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
 docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
 docker manifest push vllm/vllm-openai-cpu:latest
 docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
@@ -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.10.0")
+set(TORCH_SUPPORTED_VERSION_CUDA "2.9.1")
-set(TORCH_SUPPORTED_VERSION_ROCM "2.10.0")
+set(TORCH_SUPPORTED_VERSION_ROCM "2.9.1")
 #
 # 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}"
@@ -686,6 +686,7 @@ def get_model_params(config):
        "DeepseekV2ForCausalLM",
        "DeepseekV3ForCausalLM",
        "DeepseekV32ForCausalLM",
        "GlmMoeDsaForCausalLM",
        "Glm4MoeForCausalLM",
        "Glm4MoeLiteForCausalLM",
        "NemotronHForCausalLM",
@@ -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.6.0")
+set(DEFAULT_TRITON_KERNELS_TAG "v3.5.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")
@@ -1568,7 +1568,8 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
        {
  #endif
          unsigned int kOff = k + (thrd * A_CHUNK);
-          unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
+          unsigned int kOffcp =
              k_str + kOff;  // 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__(
@@ -1,7 +1,7 @@
 # default base image
 ARG REMOTE_VLLM="0"
 ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=rocm/vllm-dev:base_custom_releases_rocm_v0.16.0_20260211
+ARG BASE_IMAGE=rocm/vllm-dev:base
 # Sccache configuration (only used in release pipeline)
 ARG USE_SCCACHE
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2-complete
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.0-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="v0.1.10.post2"
+ARG AITER_BRANCH="6af8b687"
 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
 ARG MORI_BRANCH="2d02c6a9"
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
@@ -142,7 +142,6 @@ 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 \
@@ -240,7 +239,7 @@ RUN pip install pyyaml && cd aiter \
           export HIP_CLANG_PATH=/opt/sccache-wrappers \
           && sccache --show-stats; \
       fi \
-    && GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
    && ls /app/aiter/dist/*.whl
 RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
@@ -6,7 +6,7 @@ requires = [
    "packaging>=24.2",
    "setuptools>=77.0.3,<81.0.0",
    "setuptools-scm>=8.0",
-    "torch == 2.10.0",
+    "torch == 2.9.1",
    "wheel",
    "jinja2",
    "grpcio-tools==1.78.0",
@@ -4,7 +4,7 @@ ninja
 packaging>=24.2
 setuptools>=77.0.3,<81.0.0
 setuptools-scm>=8
-torch==2.10.0
+torch==2.9.1
 wheel
 jinja2>=3.1.6
 regex
@@ -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.10.0
+torch==2.9.1
-torchaudio==2.10.0
+torchaudio==2.9.1
 # These must be updated alongside torch
-torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
+torchvision==0.24.1 # 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
@@ -43,5 +43,5 @@ tritonclient>=2.51.0
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.1.10
+fastsafetensors>=0.2.2
 pydantic>=2.12 # 2.11 leads to error on python 3.13
@@ -1,11 +1,12 @@
 # Common dependencies
 -r common.txt
--extra-index-url https://download.pytorch.org/whl/test/rocm7.0
+--extra-index-url https://download.pytorch.org/whl/rocm6.4
-torch==2.10.0
+torch==2.9.1
-torchvision==0.25.0
+torchvision==0.24.1
-torchaudio==2.10.0
+torchaudio==2.9.1
-triton==3.6.0
+
 triton==3.5.1
 cmake>=3.26.1,<4
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
@@ -1,6 +1,11 @@
 # Common dependencies
 -r common.txt
 # The version of gRPC libraries should be consistent with each other
 grpcio==1.78.0
 grpcio-reflection==1.78.0
 grpcio-tools==1.78.0
 numba == 0.61.2 # Required for N-gram speculative decoding
 # Dependencies for AMD GPUs
@@ -14,5 +19,4 @@ setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
 runai-model-streamer[s3,gcs]==0.15.3
 conch-triton-kernels==1.2.1
-timm>=1.0.17
+timm>=1.0.17
 grpcio-tools==1.78.0 # Should match `build.txt`
@@ -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
+timm==1.0.17 # required for internvl and gemma3n-mm test
-torch==2.10.0
+torch==2.9.1
-torchaudio==2.10.0
+torchaudio==2.9.1
-torchvision==0.25.0
+torchvision==0.24.1
 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
@@ -48,12 +48,16 @@ buildkite-test-collector==0.1.9
 genai_perf>=0.0.8
 tritonclient>=2.51.0
-grpcio-tools==1.78.0 # Should match `build.txt`
+# The version of gRPC libraries should be consistent with each other
 grpcio==1.78.0
 grpcio-reflection==1.78.0
 grpcio-tools==1.78.0
 arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.1.10
+fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0
 terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test
@@ -155,10 +155,6 @@ 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
@@ -224,7 +220,7 @@ fastparquet==2024.11.0
    # via genai-perf
 fastrlock==0.8.2
    # via cupy-cuda12x
-fastsafetensors==0.1.10
+fastsafetensors==0.2.2
    # via -r requirements/test.in
 filelock==3.16.1
    # via
@@ -309,8 +305,13 @@ greenlet==3.2.3
    # via sqlalchemy
 grpcio==1.78.0
    # via
    #   -r requirements/test.in
    #   grpcio-reflection
    #   grpcio-tools
    #   ray
    #   tensorboard
 grpcio-reflection==1.78.0
    # via -r requirements/test.in
 grpcio-tools==1.78.0
    # via -r requirements/test.in
 gunicorn==23.0.0
@@ -635,7 +636,7 @@ nvidia-nvjitlink-cu12==12.9.86
    #   nvidia-cusolver-cu12
    #   nvidia-cusparse-cu12
    #   torch
-nvidia-nvshmem-cu12==3.4.5
+nvidia-nvshmem-cu12==3.3.20
    # via torch
 nvidia-nvtx-cu12==12.9.79
    # via torch
@@ -785,6 +786,7 @@ protobuf==6.33.2
    # via
    #   google-api-core
    #   googleapis-common-protos
    #   grpcio-reflection
    #   grpcio-tools
    #   mlflow-skinny
    #   opentelemetry-proto
@@ -1167,14 +1169,13 @@ tomli==2.2.1
    # via schemathesis
 tomli-w==1.2.0
    # via schemathesis
-torch==2.10.0+cu129
+torch==2.9.1+cu129
    # via
    #   -r requirements/test.in
    #   accelerate
    #   bitsandbytes
    #   efficientnet-pytorch
    #   encodec
    #   fastsafetensors
    #   kornia
    #   lightly
    #   lightning
@@ -1196,7 +1197,7 @@ torch==2.10.0+cu129
    #   torchvision
    #   vector-quantize-pytorch
    #   vocos
-torchaudio==2.10.0+cu129
+torchaudio==2.9.1+cu129
    # via
    #   -r requirements/test.in
    #   encodec
@@ -1209,7 +1210,7 @@ torchmetrics==1.7.4
    #   pytorch-lightning
    #   terratorch
    #   torchgeo
-torchvision==0.25.0+cu129
+torchvision==0.24.1+cu129
    # via
    #   -r requirements/test.in
    #   lightly
@@ -1251,7 +1252,7 @@ transformers==4.57.5
    #   transformers-stream-generator
 transformers-stream-generator==0.0.5
    # via -r requirements/test.in
-triton==3.6.0
+triton==3.5.1
    # via torch
 tritonclient==2.64.0
    # via -r requirements/test.in
@@ -1035,7 +1035,7 @@ setup(
    extras_require={
        "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy"],
        "tensorizer": ["tensorizer==2.10.1"],
-        "fastsafetensors": ["fastsafetensors >= 0.1.10"],
+        "fastsafetensors": ["fastsafetensors >= 0.2.2"],
        "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
        "audio": [
            "librosa",
@@ -1,11 +1,10 @@
 # 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 (
@@ -32,7 +31,6 @@ 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()
@@ -200,82 +198,23 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
        return [torch.ops.aten.slice_scatter.default]
-class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
+MODELS = [
-    OP_REGISTERED = False
+    TestSiluMul,
-
+    TestFusedAddRMSNorm,
-    def __init__(self):
+    TestRotaryEmbedding,
-        super().__init__()
+    TestRotaryEmbeddingSliceScatter,
-        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(
+@pytest.mark.parametrize("model_class", MODELS)
-    "model_class, do_fusion",
+@pytest.mark.parametrize("do_fusion", [True, False])
-    [
+@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
        (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),
@@ -307,14 +246,8 @@ def test_fix_functionalization(
        backend_no_func = TestBackend(*passes)
        model = model_class()
-        inputs_func = model.example_inputs()
+        torch.compile(model, backend=backend_func)(*model.example_inputs())
-        inputs_no_func = copy.deepcopy(inputs_func)
+        torch.compile(model, backend=backend_no_func)(*model.example_inputs())
        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):
@@ -332,8 +265,3 @@ 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)
@@ -1,325 +0,0 @@
 # 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,
        )
@@ -1,107 +0,0 @@
 # 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
@@ -90,7 +90,9 @@ def use_vllm_config(vllm_config: VllmConfig):
        yield
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
    with monkeypatch.context() as m:
        vllm_config = make_vllm_config()
@@ -114,7 +116,9 @@ 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"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), 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),)
@@ -128,7 +132,9 @@ 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"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
    with monkeypatch.context() as m:
        args = (torch.randn(10, 10),)
@@ -156,7 +162,9 @@ 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"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), 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.
@@ -215,7 +223,9 @@ def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
            )
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 def test_cache_load_returns_tuple_consistency_tuple_output(
    monkeypatch: pytest.MonkeyPatch,
 ):
@@ -284,7 +294,9 @@ def test_cache_load_returns_tuple_consistency_tuple_output(
            )
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 def test_shape_env(monkeypatch: pytest.MonkeyPatch):
    """
    Test that the shape environment is correctly serialized and preserved
@@ -321,7 +333,9 @@ 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"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 def test_partition_wrapper_applied_on_aot_load(
    monkeypatch: pytest.MonkeyPatch, vllm_tmp_cache: Path, mocker
 ):
@@ -412,7 +426,9 @@ def test_partition_wrapper_applied_on_aot_load(
        )
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
@create_new_process_for_each_test("spawn")
 def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
    """
@@ -476,7 +492,9 @@ 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"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 class TestStandaloneCompiledArtifacts:
    def test_init(self):
        cache = StandaloneCompiledArtifacts()
@@ -650,7 +668,9 @@ class TestStandaloneCompiledArtifacts:
        assert len(restored_cache.loaded_submodule_store) == 0
-@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 class TestStandaloneCompiledArtifactsIntegration:
    def test_add_pickle_unpickle(self):
        cache = StandaloneCompiledArtifacts()
@@ -39,7 +39,9 @@ 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"), reason="requires torch 2.10")
+@pytest.mark.skipif(
    not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
 )
 def test_dynamic_shapes_compilation(
    monkeypatch,
    model_name,
@@ -78,27 +78,3 @@ 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()
@@ -129,5 +129,5 @@ async def test_multi_chunk_streaming(
                " First words I spoke in the original phonograph."
                " A little piece of practical poetry. Mary had a little lamb,"
                " it sleeps with quite a flow, and everywhere that Mary went,"
-                " the lamb was sure to go"
+                " the lamb was sure to go."
            )
@@ -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 mxfp4_w4a16_moe_quant_config
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
    triton_kernel_moe_forward,
 )
@@ -298,18 +298,12 @@ 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)
-    if a_dtype == "bf16" and w_dtype == "mx4":
+    quant_config = FusedMoEQuantConfig.make(
-        quant_config = mxfp4_w4a16_moe_quant_config(
+        w1_bias=w1_bias_tri,
-            w1_scale=pc1,
+        w2_bias=w2_bias_tri,
-            w2_scale=pc2,
+        w1_scale=pc1,
-            w1_bias=w1_bias_tri,
+        w2_scale=pc2,
-            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,
@@ -14,7 +14,6 @@ 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):
@@ -61,10 +60,6 @@ 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,
@@ -1,110 +0,0 @@
 # 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}"
@@ -0,0 +1,80 @@
 # 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}"
@@ -275,6 +275,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
        "zai-org/GLM-4.7-Flash",
        min_transformers_version="5.0.0",
    ),
    "GlmMoeDsaForCausalLM": _HfExamplesInfo(
        "zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False
    ),
    "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
    "GPTBigCodeForCausalLM": _HfExamplesInfo(
        "bigcode/starcoder",
@@ -97,7 +97,7 @@ def can_initialize(
            "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
        )
-    if model_arch == "DeepseekV32ForCausalLM":
+    if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]:
        from vllm.platforms import current_platform
        capability = current_platform.get_device_capability()
@@ -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
@@ -7,7 +7,8 @@
 set -e
 TORCHCODEC_REPO="${TORCHCODEC_REPO:-https://github.com/pytorch/torchcodec.git}"
-TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-main}"
+# Pin to a specific release for reproducibility; update as needed.
 TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-v0.10.0}"
 echo "=== TorchCodec Installation Script ==="
@@ -87,10 +87,6 @@ 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
@@ -114,10 +110,6 @@ 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,
    )
@@ -315,28 +307,6 @@ 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
@@ -971,70 +941,6 @@ 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:
@@ -1164,14 +1070,6 @@ 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,
@@ -1393,10 +1291,6 @@ 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,
@@ -1414,10 +1308,6 @@ class rocm_aiter_ops:
            a2_scale,
            num_local_tokens,
            output_dtype,
            hidden_pad,
            intermediate_pad,
            bias1,
            bias2,
        )
    @staticmethod
@@ -1522,15 +1412,6 @@ 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,
@@ -1637,45 +1518,6 @@ 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,
@@ -1787,47 +1629,6 @@ 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)
@@ -233,7 +233,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
        from torch._inductor import standalone_compile
-        supports_aot = is_torch_equal_or_newer("2.10.0")
+        supports_aot = is_torch_equal_or_newer("2.10.0.dev")
        if not supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT:
            logger.error(
@@ -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"):
+                if is_torch_equal_or_newer("2.10.0.dev"):
                    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"):
+                        if is_torch_equal_or_newer("2.10.0.dev"):
                            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+
+        # assume_32bit_indexing is only available in torch 2.10.0.dev+
        inductor_config_patches = {}
-        if is_torch_equal_or_newer("2.10.0"):
+        if is_torch_equal_or_newer("2.10.0.dev"):
            inductor_config_patches["assume_32bit_indexing"] = (
                self.compilation_config.dynamic_shapes_config.assume_32_bit_indexing
            )
@@ -1,230 +0,0 @@
 # 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)
@@ -28,9 +28,7 @@ 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():
@@ -138,11 +136,6 @@ 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)]
@@ -162,24 +162,6 @@ 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
@@ -226,20 +208,13 @@ class FixFunctionalizationPass(VllmInductorPass):
        self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]
    ) -> None:
        """
-        Replace mutated getitem users of the auto-functionalized node with the
+        Replace all 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)
@@ -282,20 +257,10 @@ class FixFunctionalizationPass(VllmInductorPass):
        with graph.inserting_before(node):
            function = node.args[0]
            if args is None:
-                fn_node = graph.call_function(function, kwargs=node.kwargs)
+                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
                )
-                fn_node = graph.call_function(function, args=args)
+                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)
@@ -1,134 +0,0 @@
 # 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)
@@ -127,13 +127,6 @@ 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
@@ -206,7 +199,6 @@ class PassConfig:
        "fuse_gemm_comms",
        "fuse_allreduce_rms",
        "fuse_act_padding",
        "fuse_rope_kvcache",
        mode="wrap",
    )
    @classmethod
@@ -252,12 +244,6 @@ 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):
@@ -839,19 +825,6 @@ 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")
@@ -181,7 +181,7 @@ class SpeculativeConfig:
    @staticmethod
    def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
        initial_architecture = hf_config.architectures[0]
-        if hf_config.model_type in ("deepseek_v3", "deepseek_v32"):
+        if hf_config.model_type in ("deepseek_v3", "deepseek_v32", "glm_moe_dsa"):
            hf_config.model_type = "deepseek_mtp"
        if hf_config.model_type == "deepseek_mtp":
            n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
@@ -1365,20 +1365,6 @@ 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)
@@ -592,8 +592,6 @@ 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
@@ -1237,14 +1235,6 @@ 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
@@ -1401,8 +1391,6 @@ 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
@@ -48,7 +48,6 @@ class RealtimeConnection:
        self.generation_task: asyncio.Task | None = None
        self._is_connected = False
        self._is_input_finished = False
        self._is_model_validated = False
        self._max_audio_filesize_mb = envs.VLLM_MAX_AUDIO_CLIP_FILESIZE_MB
@@ -145,7 +144,7 @@ class RealtimeConnection:
            commit_event = InputAudioBufferCommit(**event)
            # final signals that the audio is finished
            if commit_event.final:
-                self._is_input_finished = True
+                self.audio_queue.put_nowait(None)
            else:
                await self.start_generation()
        else:
@@ -239,11 +238,6 @@ class RealtimeConnection:
                    # finish because websocket connection was killed
                    break
                if self.audio_queue.empty() and self._is_input_finished:
                    # finish because client signals that audio input
                    # is finished
                    break
            usage = UsageInfo(
                prompt_tokens=prompt_token_ids_len,
                completion_tokens=completion_tokens_len,
@@ -98,6 +98,7 @@ 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
@@ -258,14 +259,6 @@ 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")))
@@ -278,7 +271,7 @@ def use_aot_compile() -> bool:
    default_value = (
        "1"
-        if is_torch_equal_or_newer("2.11.0.dev") and not disable_compile_cache()
+        if is_torch_equal_or_newer("2.10.0.dev") and not disable_compile_cache()
        else "0"
    )
@@ -889,7 +882,14 @@ 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": use_aiter,
+    "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")
    ),
    # 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 disabled.
+    # By default is enabled.
    "VLLM_ROCM_USE_AITER_MHA": lambda: (
-        os.getenv("VLLM_ROCM_USE_AITER_MHA", "False").lower() in ("true", "1")
+        os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in ("true", "1")
    ),
    # Whether to use aiter fp4 gemm asm.
    # By default is disabled.
@@ -1606,15 +1606,6 @@ 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.
@@ -570,11 +570,11 @@ direct_register_custom_op(
 def get_attention_context(
    layer_name: str,
-) -> tuple[Any, "Attention | MLAAttention", torch.Tensor, torch.Tensor]:
+) -> tuple[Any, "Attention | MLAAttention", torch.Tensor]:
    """Extract attention context for a given layer.
    This helper function extracts the attention metadata, attention layer
-    instance, KV cache tensor, and slot mapping for a specific layer.
+    instance, and KV cache tensor for a specific layer.
    Args:
        layer_name: The name/identifier of the attention layer.
@@ -585,7 +585,6 @@ 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.
@@ -594,14 +593,9 @@ def get_attention_context(
    attn_metadata = forward_context.attn_metadata
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
-    attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
+    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
+    return attn_metadata, attn_layer, kv_cache
    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
@@ -611,7 +605,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
@@ -642,7 +636,15 @@ 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.
    """
-    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
+    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)
    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"
@@ -689,7 +691,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,
@@ -40,8 +40,8 @@ def maybe_transfer_kv_layer(func: Callable) -> Callable:
        layer_name: str = args[layer_name_index]
-        # Extract attention context (metadata, layer, kv_cache, layer_slot_mapping)
+        # Extract attention context (layer-specific metadata, layer, and kv_cache)
-        attn_metadata, _, kv_cache, _ = get_attention_context(layer_name)
+        attn_metadata, attn_layer, 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)
@@ -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,
@@ -974,7 +974,7 @@ def enable_batch_invariant_mode():
    )
    reduced_precision_val = (
-        (False, False) if is_torch_equal_or_newer("2.10.0") else False
+        (False, False) if is_torch_equal_or_newer("2.10.0.dev") else False
    )
    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
        reduced_precision_val
@@ -386,10 +386,6 @@ 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
@@ -536,8 +532,6 @@ 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,
@@ -555,8 +549,6 @@ 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,
@@ -572,8 +564,6 @@ 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:
    """
@@ -585,8 +575,6 @@ 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,
@@ -666,26 +654,6 @@ 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"],
@@ -723,8 +691,6 @@ 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.
@@ -733,8 +699,6 @@ 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,
@@ -823,32 +787,6 @@ 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,
@@ -38,6 +38,7 @@ 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,
@@ -1582,11 +1583,6 @@ 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
@@ -1621,10 +1617,17 @@ 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.startswith("w_mxfp4"):
+        if ocp_mx_scheme in {
            "w_mxfp4_a_mxfp4",
            "w_mxfp4_a_mxfp6_e3m2",
            "w_mxfp4_a_mxfp6_e2m3",
        }:
            # 16bit activation and fp4x2 packed weight
            assert hidden_states.size(1) == w1.size(2) * 2, "hidden size mismatch"
-        elif ocp_mx_scheme.startswith("w_mxfp6"):
+        elif ocp_mx_scheme in {
            "w_mxfp6_e3m2_a_mxfp6_e3m2",
            "w_mxfp6_e2m3_a_mxfp6_e2m3",
        }:
            assert hidden_states.size(1) == (w1.size(2) * 4) // 3, (
                "hidden size mismatch"
            )
@@ -1714,13 +1717,17 @@ 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.startswith("w_mxfp4"):
+        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,
        }:
            # 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.startswith("w_mxfp6_e3m2"):
+        elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2:
            w1 = dequant_mxfp6(
                w1, w1_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
            )
@@ -1729,7 +1736,7 @@ def fused_experts_impl(
                w2, w2_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
            )
            w2_scale = None
-        elif ocp_mx_scheme.startswith("w_mxfp6_e2m3"):
+        elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3:
            w1 = dequant_mxfp6(
                w1, w1_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
            )
@@ -1772,7 +1779,6 @@ 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
@@ -1840,7 +1846,6 @@ 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:
@@ -19,42 +19,17 @@ 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 (
+        from triton_kernels.matmul_ogs import FnSpecs, FusedActivation, matmul_ogs
-            FnSpecs,
+        from triton_kernels.routing import RoutingData, routing, routing_from_bitmatrix
-            FusedActivation,
+        from triton_kernels.tensor import Bitmatrix
            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 "
@@ -103,68 +78,6 @@ 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
@@ -178,7 +91,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 = legacy_routing(
+    routing_data, gather_idx, scatter_idx = routing(
        gating_output, topk, sm_first=not renormalize
    )
@@ -254,22 +167,10 @@ def triton_kernel_fused_experts(
    )
    output_tensor = _resize_cache(output_tensor, (batch_dim, M, K))
-    act = (
+    act = FusedActivation(
-        FusedActivation(
+        FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
-            FnSpecs(
+        (swiglu_alpha, swiglu_limit),
-                "swiglu",
+        2,
                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
@@ -299,182 +200,6 @@ 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,
@@ -506,22 +231,13 @@ def make_routing_data(
    bitmatrix_shape = [n_rows, bm_cols * 32]
    bitmatrix_shape_max = [n_rows, None]
-    bitmatrix = (
+    bitmatrix = Bitmatrix(
-        Bitmatrix(
+        bitmatrix, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max, scratchpad=None
            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 = legacy_routing_from_bitmatrix(
+    routing_data, gather_indx, scatter_indx = routing_from_bitmatrix(
        bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
    )
@@ -663,9 +379,6 @@ 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]
@@ -221,14 +221,12 @@ 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
@@ -237,12 +235,11 @@ 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.
@@ -257,7 +254,7 @@ def maybe_roundup_hidden_size(
    )
    # we are padding globally so EP buffer allocation works
-    if model_type == "gpt_oss" and is_mxfp4_quant:
+    if quant_config and quant_config.get_name() == "mxfp4":
        from vllm.model_executor.layers.quantization.mxfp4 import (
            Mxfp4Backend,
            get_mxfp4_backend,
@@ -401,6 +398,15 @@ 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:
@@ -502,6 +508,7 @@ 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
@@ -541,26 +548,6 @@ 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,
@@ -628,7 +615,6 @@ 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,
@@ -1161,43 +1147,16 @@ class FusedMoE(CustomOp):
        expert_id: int,
        return_success: bool = False,
    ) -> bool | None:
-        if self.quant_config is not None:
+        if self.quant_config and self.quant_config.get_name() == "mxfp4":
-            if self.quant_config.get_name() == "mxfp4":
+            # (FIXME) for gpt-oss all experts are combined
-                # (FIXME) for gpt-oss all experts are combined
+            if "bias" in weight_name:
-                if "bias" in weight_name:
+                dim1 = loaded_weight.shape[1]
-                    dim1 = loaded_weight.shape[1]
+                param.data[:, :dim1].copy_(loaded_weight)
-                    param.data[:, :dim1].copy_(loaded_weight)
+            else:
-                else:
+                dim1 = loaded_weight.shape[1]
-                    dim1 = loaded_weight.shape[1]
+                dim2 = loaded_weight.shape[2]
-                    dim2 = loaded_weight.shape[2]
+                param.data[:, :dim1, :dim2].copy_(loaded_weight)
-                    param.data[:, :dim1, :dim2].copy_(loaded_weight)
+            return True if return_success else None
                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
@@ -23,9 +23,6 @@ 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
@@ -244,27 +241,7 @@ 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:
@@ -168,19 +168,3 @@ 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
@@ -6,7 +6,6 @@ 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
@@ -50,10 +49,9 @@ 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 replace_parameter, set_weight_attrs
+from vllm.model_executor.utils import 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
@@ -77,8 +75,6 @@ class Mxfp4Backend(Enum):
    # Triton Backend
    TRITON = 6
    CK = 7
 def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
    """
@@ -166,13 +162,9 @@ def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
    elif current_platform.is_xpu():
        logger.info_once("Using xpu backend on XPU")
        return Mxfp4Backend.MARLIN
-    elif current_platform.is_rocm():
+    elif current_platform.is_rocm() and has_triton_kernels():
-        if rocm_aiter_ops.is_enabled() and on_gfx950():
+        logger.info_once("Using Triton backend")
-            logger.info_once("Using CK MXFP4 MoE backend (Aiter ROCm)")
+        return Mxfp4Backend.TRITON
            return Mxfp4Backend.CK
        elif has_triton_kernels():
            logger.info_once("Using Triton backend")
            return Mxfp4Backend.TRITON
    return Mxfp4Backend.NONE
@@ -237,15 +229,10 @@ 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
@@ -339,10 +326,6 @@ 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(
@@ -755,99 +738,46 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                layer.w2_weight_scale = torch.nn.Parameter(
                    w2_scales_interleaved, requires_grad=False
                )
-        elif (
+        elif self.mxfp4_backend == Mxfp4Backend.TRITON:
-            self.mxfp4_backend == Mxfp4Backend.TRITON
+            from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
-            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)
-            if self.mxfp4_backend == Mxfp4Backend.CK:
+            # Ideally we'd use FusedMoEModularKernel.prepare_finalize object
-                w13_aiter_weight = layer.w13_weight.contiguous()
+            # (stored in self.fused_experts) to determine if the MoE has a
-                w13_aiter_scale = layer.w13_weight_scale.contiguous()
+            # batched activation format. As self.fused_experts is not
-                w2_aiter_weight = layer.w2_weight.contiguous()
+            # initialized at this point, we resort to checking the MoE config
-                w2_aiter_scale = layer.w2_weight_scale.contiguous()
+            # directly.
-
+            is_batched_moe = self.moe.use_pplx_kernels or self.moe.use_deepep_ll_kernels
-                e, n, k = w13_aiter_weight.shape
+            if is_batched_moe:
-                w13_aiter_weight = (
+                num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
                    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:
-                from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
+                num_warps = 8
-                # Ideally we'd use FusedMoEModularKernel.prepare_finalize object
+            w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
-                # (stored in self.fused_experts) to determine if the MoE has a
+                layer.w13_weight, layer.w13_weight_scale, num_warps
-                # batched activation format. As self.fused_experts is not
+            )
-                # initialized at this point, we resort to checking the MoE config
+            w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
-                # directly.
+                layer.w2_weight, layer.w2_weight_scale, num_warps
-                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(
+            self.w13_precision_config = PrecisionConfig(
-                    weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
+                weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
-                )
+            )
-                self.w2_precision_config = PrecisionConfig(
+            self.w2_precision_config = PrecisionConfig(
-                    weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
+                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}: "
@@ -883,10 +813,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                w1_scale=layer.w13_weight_scale,
                w2_scale=layer.w2_weight_scale,
            )
-        elif (
+        elif self.mxfp4_backend in [Mxfp4Backend.SM100_FI_MXFP4_BF16]:
            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,
@@ -960,7 +887,6 @@ 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(
@@ -1149,27 +1075,6 @@ 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,
@@ -35,7 +35,6 @@ 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
@@ -60,35 +59,6 @@ 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)
@@ -138,20 +108,7 @@ class QuarkConfig(QuantizationConfig):
        if should_ignore_layer(
            prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping
        ):
-            if (
+            return UnquantizedLinearMethod()
                "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
@@ -363,83 +320,38 @@ class QuarkConfig(QuantizationConfig):
        # Only symmetric weight quantization supported.
        return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
-    def _is_mx_fp4(
+    def _is_ocp_mx(
-        self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
+        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 MX-FP4 format: expected input dtype "
+                "Quark model is not in OCP MX format: "
-                "to be one of {'fp4', 'fp6_e2m3', 'fp8_e4m3'}"
+                "weight_quant or input_quant not set"
            )
            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":
+        if (
-            logger.debug(
+            weight_quant.get("qscheme") != "per_group"
-                "Quark model's weight quantization is incompatible with OCP MX format: "
+            or input_quant.get("qscheme") != "per_group"
-                "weight is not per_group."
+        ):
-            )
+            logger.debug("Quark model is not in OCP MX format: not per_group")
            return False
        # Input and weight group size needs to be 32.
-        if weight_quant.get("group_size") != 32:
+        if weight_quant.get("group_size") != 32 or input_quant.get("group_size") != 32:
-            logger.debug(
+            logger.debug("Quark model is not in OCP MX format: not group_size=32")
                "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":
+        if (
-            logger.debug(
+            weight_quant.get("scale_format") != "e8m0"
-                "Quark model's weight quantization is incompatible with OCP MX format: "
+            or input_quant.get("scale_format") != "e8m0"
-                "scale_format of weight is not e8m0."
+        ):
-            )
+            logger.debug("Quark model is not in OCP MX format: not scale_format e8m0")
            return False
        # Input and weight dtypes need to be any of fp4,
@@ -448,31 +360,14 @@ class QuarkConfig(QuantizationConfig):
            "fp4",
            "fp6_e3m2",
            "fp6_e2m3",
-        }:
+        } or input_quant.get("dtype") not in {"fp4", "fp6_e3m2", "fp6_e2m3"}:
            logger.debug(
-                "Quark model's weight quantization is incompatible with OCP MX format: "
+                "Quark model is not in OCP MX format: dtype not fp4, fp6_e3m2, fp6_e2m3"
                "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]:
@@ -524,9 +419,7 @@ class QuarkConfig(QuantizationConfig):
            )
            return global_quant_config
-    def _get_scheme_from_config(
+    def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkScheme":
        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 "
@@ -548,10 +441,8 @@ class QuarkConfig(QuantizationConfig):
                is_static_input_scheme=True,
                input_symmetric=input_config.get("symmetric"),
            )
-        elif self._is_w_ocp_mx_a_x(weight_config, input_config):
+        elif self._is_ocp_mx(weight_config, input_config):
-            return QuarkOCP_MX(
+            return QuarkOCP_MX(weight_config, input_config)
                weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
            )
        raise NotImplementedError(
            "No quark compatible scheme was found. "
@@ -559,15 +450,11 @@ class QuarkConfig(QuantizationConfig):
            f"Input config: {input_config}"
        )
-    def get_scheme(
+    def get_scheme(self, layer: torch.nn.Module, layer_name: str) -> "QuarkScheme":
        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(
+        scheme = self._get_scheme_from_config(layer_quant_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())
@@ -24,12 +24,7 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
    OCP_MX_BLOCK_SIZE,
    OCP_MX_Scheme,
 )
-from vllm.model_executor.parameter import (
+from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
    GroupQuantScaleParameter,
    ModelWeightParameter,
    PackedvLLMParameter,
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from .quark_scheme import QuarkScheme
@@ -174,16 +169,13 @@ except (ImportError, AttributeError, RuntimeError):
 class QuarkOCP_MX(QuarkScheme):
    def __init__(
-        self,
+        self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any]
        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")
@@ -277,13 +269,7 @@ class QuarkOCP_MX(QuarkScheme):
                layer.weight_scale.data, requires_grad=False
            )
        else:
-            if self.dynamic_mxfp4_quant:
+            if self.rocm_use_aiter_fp4_asm_gemm:
                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
@@ -316,51 +302,36 @@ class QuarkOCP_MX(QuarkScheme):
        weight_loader: Callable,
        **kwargs,
    ):
-        if self.dynamic_mxfp4_quant:
+        output_size_per_partition = sum(output_partition_sizes)
-            weight = ModelWeightParameter(
+        layer.logical_widths = output_partition_sizes
                data=torch.empty(
                    sum(output_partition_sizes),
                    input_size_per_partition,
                    dtype=params_dtype,
                ),
                input_dim=1,
                output_dim=0,
                weight_loader=weight_loader,
            )
-            layer.register_parameter("weight", weight)
+        # WEIGHT
-            set_weight_attrs(weight, kwargs)
+        weight = PackedvLLMParameter(
-        else:
+            data=torch.empty(
-            output_size_per_partition = sum(output_partition_sizes)
+                output_size_per_partition,
-            layer.logical_widths = output_partition_sizes
+                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
+        # WEIGHT SCALE
-            weight = PackedvLLMParameter(
+        weight_scale = GroupQuantScaleParameter(
-                data=torch.empty(
+            data=torch.empty(
-                    output_size_per_partition,
+                output_size_per_partition,
-                    self.get_packed_dim(input_size_per_partition, self.weight_dtype),
+                input_size_per_partition // OCP_MX_BLOCK_SIZE,
-                    dtype=torch.uint8,
+                dtype=torch.uint8,
-                ),
+            ),
-                input_dim=1,
+            input_dim=1,
-                output_dim=0,
+            output_dim=0,
-                packed_dim=1,
+            weight_loader=weight_loader,
-                packed_factor=self.packed_factor,
+        )
-                weight_loader=weight_loader,
+        layer.register_parameter("weight_scale", weight_scale)
            )
            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,
@@ -20,44 +20,26 @@ 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 and weight_dtype not in OCP_MX_DTYPES:
+        if input_dtype not in OCP_MX_DTYPES or 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"
@@ -836,7 +836,7 @@ class DeepseekV2MLAAttention(nn.Module):
                qk_rope_head_dim,
                max_position=max_position_embeddings,
                rope_parameters=config.rope_parameters,
-                is_neox_style=True,
+                is_neox_style=not getattr(config, "indexer_rope_interleave", False),
            )
            self.indexer = Indexer(
                vllm_config,
@@ -1499,6 +1499,10 @@ class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
 class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM):
    pass
 # Compatibility with
 # https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
 def get_spec_layer_idx_from_weight_name(
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import typing
+from collections.abc import Iterable
 from collections.abc import Callable, Iterable
 import torch
 import torch.distributed as dist
@@ -26,17 +25,13 @@ 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 (
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
    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
@@ -103,7 +98,6 @@ 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",
        )
@@ -111,7 +105,6 @@ 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",
        )
@@ -313,19 +306,6 @@ 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,
@@ -338,6 +318,7 @@ 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
@@ -352,11 +333,9 @@ class GptOssModel(nn.Module):
        )
        intermediate_size = self.config.intermediate_size
-        intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE
+        intermediate_size_block = intermediate_size // mxfp4_block
        per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size)
-        per_rank_intermediate_size = (
+        per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block
            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
@@ -391,9 +370,7 @@ 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]
@@ -518,449 +495,6 @@ 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,
@@ -1101,7 +635,6 @@ class GptOssModel(nn.Module):
            if hasattr(self.config, "quantization_config")
            else None
        )
        if quant_method == "mxfp4":
            return self._load_weights_mxfp4(
                ep_rank_end,
@@ -1111,15 +644,6 @@ 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,
@@ -1152,15 +676,6 @@ 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",
        },
    )
@@ -1210,6 +725,18 @@ 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,
@@ -114,6 +114,7 @@ _TEXT_GENERATION_MODELS = {
    "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
    "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
    "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
    "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
    "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
    "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
    "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),
@@ -360,11 +360,6 @@ 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()
@@ -378,9 +373,9 @@ class RocmPlatform(Platform):
                logger.info("Using Aiter Flash Attention backend.")
                return AttentionBackendEnum.ROCM_AITER_FA.get_path()
-            # Default: ROCm split Attention
+            # Default: Triton Unified Attention
-            logger.info("Using ROCm Attention backend.")
+            logger.info("Using Triton Attention backend.")
-            return AttentionBackendEnum.ROCM_ATTN.get_path()
+            return AttentionBackendEnum.TRITON_ATTN.get_path()
        raise RuntimeError(
            f"Attention backend {selected_backend.name} is not supported on "
@@ -237,6 +237,7 @@ class ModelArchConfigConvertorBase:
            "deepseek_v3",
            "deepseek_v32",
            "deepseek_mtp",
            "glm_moe_dsa",
            "glm4_moe_lite",
            "glm4_moe_lite_mtp",
            "kimi_k2",
@@ -723,33 +723,6 @@ 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."""
@@ -687,8 +687,6 @@ class AiterFlashAttentionBackend(AttentionBackend):
    def get_supported_head_sizes(cls) -> list[int]:
        return [64, 128, 256]
    forward_includes_kv_cache_update: bool = False
    @staticmethod
    def get_name() -> str:
        return "FLASH_ATTN"
@@ -984,10 +982,49 @@ class AiterFlashAttentionImpl(AttentionImpl):
        # performance to make sure it does not introduce any overhead.
        num_actual_tokens = attn_metadata.num_actual_tokens
        key_cache, value_cache = kv_cache.unbind(0)
-
+        # key and value may be None in the case of cross attention. They are
        # calculated once based on the output from the encoder and then cached
        # in KV cache.
        if self.kv_cache_dtype.startswith("fp8"):
            key_cache = key_cache.view(current_platform.fp8_dtype())
            value_cache = value_cache.view(current_platform.fp8_dtype())
        if (
            self.kv_sharing_target_layer_name is None
            and key is not None
            and value is not None
        ):
            # Reshape the input keys and values and store them in the cache.
            # Skip this if sharing KV cache with an earlier attention layer.
            # NOTE(woosuk): Here, key and value are padded while slot_mapping
            # is not padded. However, we don't need to do
            # key[:num_actual_tokens] and value[:num_actual_tokens] because
            # the reshape_and_cache_flash op uses the slot_mapping's shape
            # to determine the number of actual tokens.
            if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
                # We may calculate per token quant scale in
                # reshape_and_cache_shuffle_triton which might differ from
                # vllm's style when shuffle layout is used.
                reshape_and_cache_shuffle_triton(
                    key,
                    value,
                    key_cache,
                    value_cache,
                    attn_metadata.slot_mapping,
                    self.kv_cache_dtype,
                    attn_metadata.k_scale,
                    attn_metadata.v_scale,
                )
            else:
                torch.ops._C_cache_ops.reshape_and_cache_flash(
                    key,
                    value,
                    key_cache,
                    value_cache,
                    attn_metadata.slot_mapping,
                    self.kv_cache_dtype,
                    layer._k_scale,
                    layer._v_scale,
                )
        # decode:extend:prefill
        query = query[:num_actual_tokens]
@@ -1178,57 +1215,3 @@ class AiterFlashAttentionImpl(AttentionImpl):
            )
        return output
    def do_kv_cache_update(
        self,
        layer: Attention,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
    ):
        key_cache, value_cache = kv_cache.unbind(0)
        # key and value may be None in the case of cross attention. They are
        # calculated once based on the output from the encoder and then cached
        # in KV cache.
        if self.kv_cache_dtype.startswith("fp8"):
            key_cache = key_cache.view(current_platform.fp8_dtype())
            value_cache = value_cache.view(current_platform.fp8_dtype())
        # Reshape the input keys and values and store them in the cache.
        # Skip this if sharing KV cache with an earlier attention layer.
        # NOTE(woosuk): Here, key and value are padded while slot_mapping
        # is not padded. However, we don't need to do
        # key[:num_actual_tokens] and value[:num_actual_tokens] because
        # the reshape_and_cache_flash op uses the slot_mapping's shape
        # to determine the number of actual tokens.
        if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
            # We may calculate per token quant scale in
            # reshape_and_cache_shuffle_triton which might differ from
            # vllm's style when shuffle layout is used.
            k_scale = layer._k_scale
            v_scale = layer._v_scale
            assert k_scale is not None and v_scale is not None, (
                "k_scale and v_scale are required for shuffled update"
            )
            reshape_and_cache_shuffle_triton(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                k_scale,
                v_scale,
            )
        else:
            torch.ops._C_cache_ops.reshape_and_cache_flash(
                key,
                value,
                key_cache,
                value_cache,
                slot_mapping,
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
            )
@@ -5,7 +5,6 @@
 import torch
 from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
    QuantKey,
@@ -208,42 +207,3 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
            layer._k_scale,
            layer._v_scale,
        )
    def fused_rope_kvcache_supported(self):
        return rocm_aiter_ops.is_enabled()
    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,
    ):
        key_cache, value_cache = kv_cache.unbind(0)
        flash_layout = True
        is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
        if is_fp8_kv_cache:
            key_cache = key_cache.view(self.fp8_dtype)
            value_cache = value_cache.view(self.fp8_dtype)
        rocm_aiter_ops.triton_rope_and_cache(
            query,
            key,
            value,
            positions,
            cos_sin_cache,
            is_neox,
            key_cache,
            value_cache,
            layer_slot_mapping,
            layer._k_scale,
            layer._v_scale,
            flash_layout,
            is_fp8_kv_cache,
        )
@@ -7,7 +7,6 @@ from typing import ClassVar
 import torch
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -416,46 +415,3 @@ class RocmAttentionImpl(AttentionImpl):
                layer._k_scale,
                layer._v_scale,
            )
    def fused_rope_kvcache_supported(self):
        return rocm_aiter_ops.is_enabled()
    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,
    ):
        key_cache, value_cache = PagedAttention.split_kv_cache(
            kv_cache,
            layer.num_kv_heads,  # type: ignore[attr-defined]
            layer.head_size,  # type: ignore[attr-defined]
        )
        flash_layout = False
        is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
        if is_fp8_kv_cache:
            key_cache = key_cache.view(self.fp8_dtype)
            value_cache = value_cache.view(self.fp8_dtype)
        rocm_aiter_ops.triton_rope_and_cache(
            query,
            key,
            value,
            positions,
            cos_sin_cache,
            is_neox,
            key_cache,
            value_cache,
            layer_slot_mapping,
            layer._k_scale,
            layer._v_scale,
            flash_layout,
            is_fp8_kv_cache,
        )
@@ -7,7 +7,6 @@ from typing import ClassVar
 import torch
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
@@ -597,42 +596,3 @@ class TritonAttentionImpl(AttentionImpl):
            layer._k_scale,
            layer._v_scale,
        )
    def fused_rope_kvcache_supported(self):
        return rocm_aiter_ops.is_enabled()
    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,
    ):
        key_cache, value_cache = kv_cache.unbind(1)
        flash_layout = True
        is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
        if is_fp8_kv_cache:
            key_cache = key_cache.view(self.fp8_dtype)
            value_cache = value_cache.view(self.fp8_dtype)
        rocm_aiter_ops.triton_rope_and_cache(
            query,
            key,
            value,
            positions,
            cos_sin_cache,
            is_neox,
            key_cache,
            value_cache,
            layer_slot_mapping,
            layer._k_scale,
            layer._v_scale,
            flash_layout,
            is_fp8_kv_cache,
        )
@@ -1503,6 +1503,24 @@ class SpecDecodeBaseProposer:
                del self.model.lm_head
            self.model.lm_head = target_language_model.lm_head
            # MTP models call compute_logits via shared_head.head (a
            # ParallelLMHead inside each MTP layer), not self.model.lm_head.
            # If the checkpoint omits a copy of the lm_head weights at the
            # MTP layer path, shared_head.head stays uninitialised and
            # produces NaN logits. Always share it explicitly.
            inner = getattr(self.model, "model", None)
            layers = getattr(inner, "layers", None) if inner else None
            if layers is not None:
                items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
                for layer in items:
                    sh = getattr(layer, "shared_head", None)
                    if sh is not None and hasattr(sh, "head"):
                        del sh.head
                        sh.head = target_language_model.lm_head
                        logger.info(
                            "Shared target model lm_head with MTP shared_head.head."
                        )
    @torch.inference_mode()
    def dummy_run(
        self,
Author	SHA1	Message	Date
Andreas Karatzasandkhluu	89a77b1084	[ROCm][CI] Pin TorchCodec to v0.10.0 for ROCm compatibility (#34447 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com> (cherry picked from commit `4c078fa546`) (cherry picked from commit `a976961fb7`)	2026-02-24 20:30:22 -08:00
Kevin H. Luu	d3c1513f5f	[ci] Use the right tag for CPU arm64 image (#34915 ) Signed-off-by: Kevin H. Luu <khluu000@gmail.com> (cherry picked from commit `a1a2d79442`) (cherry picked from commit `772f708391`)	2026-02-24 20:30:13 -08:00
Cyrus Leungandkhluu	5dbfbc967b	[CI/Build] Fix gRPC version mismatch (#35013 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> (cherry picked from commit `965fe45935`) (cherry picked from commit `9030895929`)	2026-02-24 20:30:02 -08:00
khluu	c86cdcbcd2	Revert "[Release 2.10] Update to Torch 2.10 - final release (#30525 )" This reverts commit `f97ca67176`.	2026-02-24 20:28:53 -08:00
khluu	3c9496f146	Revert "[Bugfix][ROCm][GPT-OSS] Use old triton_kernels implementation on ROCm if the new API is not available (#34153 )" This reverts commit `55a1baebc5`.	2026-02-24 20:28:45 -08:00
khluu	2d5be1dd5c	release script Signed-off-by: khluu <khluu000@gmail.com>	2026-02-12 02:37:52 -08:00
Michael Goinandkhluu	7a06e5b05b	[Bugfix] Fix MTP accuracy for GLM-5 (#34385 ) Signed-off-by: mgoin <mgoin64@gmail.com> (cherry picked from commit `ec12d39d44`)	2026-02-11 20:54:27 -08:00
Junseo Parkandkhluu	946b2f106c	[Bugfix] send None sentinel on final commit so server properly sends transcription.done (#33963 ) Signed-off-by: pjs102793 <pjs102793@naver.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> (cherry picked from commit `5458eb835d`)	2026-02-11 20:54:14 -08:00
Nick Hillandkhluu	5e8adb0c49	[Misc] Bump `fastsafetensors` version for latest fixes (#34273 ) Signed-off-by: Nick Hill <nickhill123@gmail.com> (cherry picked from commit `79504027ef`)	2026-02-11 20:54:00 -08:00
Xinyu Dongandkhluu	9be1ff2d3a	[Bugfix] fix default is_neox_style is True for deepseek (#34353 ) Signed-off-by: dongxinyu03 <dongxinyu03@baidu.com> (cherry picked from commit `be7f3d5d20`)	2026-02-11 20:53:40 -08:00
Jee Jee Liandkhluu	b3ee90f961	[Model] GLM adaptation (#34124 ) (cherry picked from commit `978a37c823`)	2026-02-11 20:53:11 -08:00