Merge ROCm:fused_aiter_rope_kvcache (#33443)

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

csrc/rocm/skinny_gemms.cu

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

docker/Dockerfile.rocm

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

docker/Dockerfile.rocm_base

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

tests/compile/passes/test_functionalization.py

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

tests/compile/passes/test_rope_kvcache_fusion.py

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

tests/compile/passes/test_scatter_split_replace.py

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

tests/config/test_config_generation.py

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

tests/kernels/moe/test_gpt_oss_triton_kernels.py

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

tests/models/quantization/test_gpt_oss.py

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

tests/models/quantization/test_gpt_oss_attn_quantization.py

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

tests/v1/attention/test_attention_backends.py

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

vllm/_aiter_ops.py

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

vllm/compilation/passes/fusion/rope_kvcache_fusion.py

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

vllm/compilation/passes/pass_manager.py

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

vllm/compilation/passes/utility/fix_functionalization.py

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

vllm/compilation/passes/utility/scatter_split_replace.py

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

vllm/config/compilation.py

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

vllm/config/vllm.py

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

vllm/engine/arg_utils.py

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

vllm/envs.py

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

vllm/model_executor/layers/attention/attention.py

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

vllm/model_executor/layers/attention/kv_transfer_utils.py

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

vllm/model_executor/layers/attention/mla_attention.py

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

vllm/model_executor/layers/fused_moe/config.py

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

vllm/model_executor/layers/fused_moe/fused_moe.py

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

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

@@ -299,6 +299,182 @@ def triton_kernel_fused_experts(
     return output_tensor
 
 
+def triton_kernel_moe_oss_forward(
+    hidden_states: torch.Tensor,
+    w1,  # Tensor or triton_kernels.Tensor
+    w2,  # Tensor or triton_kernels.Tensor
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    activation: str = "silu",
+    quant_config: FusedMoEQuantConfig | None = None,
+    apply_router_weight_on_input: bool = False,
+    global_num_experts: int = -1,
+    expert_map: torch.Tensor | None = None,
+    unpadded_N_w1=None,
+    unpadded_K_w1=None,
+    unpadded_N_w2=None,
+    unpadded_K_w2=None,
+) -> torch.Tensor:
+    assert quant_config is not None
+
+    if quant_config.use_mxfp4_w4a16:
+        from triton_kernels.routing import routing
+
+        routing_data, gather_idx, scatter_idx = routing(
+            gating_output, topk, sm_first=not renormalize
+        )
+    elif quant_config.use_mxfp4_w4a4:
+        from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
+
+        routing_data, gather_idx, scatter_idx = aiter_routing(
+            gating_output, topk, sm_first=not renormalize
+        )
+
+    return triton_kernel_fused_oss_experts(
+        None,
+        hidden_states,
+        w1,
+        w2,
+        routing_data,
+        gather_idx,
+        scatter_idx,
+        activation=activation,
+        quant_config=quant_config,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        global_num_experts=global_num_experts,
+        expert_map=expert_map,
+        unpadded_N_w1=unpadded_N_w1,
+        unpadded_K_w1=unpadded_K_w1,
+        unpadded_N_w2=unpadded_N_w2,
+        unpadded_K_w2=unpadded_K_w2,
+    )
+
+
+# This is a triton implementation of the fused_experts function
+def triton_kernel_fused_oss_experts(
+    output_tensor: torch.Tensor,
+    hidden_states: torch.Tensor,
+    w1,  # Tensor or triton_kernels.Tensor
+    w2,  # Tensor or triton_kernels.Tensor
+    routing_data,  # RoutingData
+    gather_indx,  # GatherIndx
+    scatter_indx,  # ScatterIndx
+    activation: str = "silu",
+    quant_config: FusedMoEQuantConfig | None = None,
+    swiglu_alpha: float = 1.702,
+    swiglu_limit: float = 7.0,
+    apply_router_weight_on_input: bool = False,
+    global_num_experts: int = -1,
+    expert_map: torch.Tensor | None = None,
+    a1q_scale: torch.Tensor | None = None,
+    unpadded_N_w1=None,
+    unpadded_K_w1=None,
+    unpadded_N_w2=None,
+    unpadded_K_w2=None,
+) -> torch.Tensor:
+    if quant_config is None:
+        quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
+
+    # type check, uint8 means mxfp4
+    assert hidden_states.dtype == torch.bfloat16
+    assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
+    assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
+
+    # Shape check, only check non-mxfp4
+    assert hidden_states.shape[-1] == w1.shape[-2]
+    assert w2.shape[-1] == w1.shape[1]
+
+    E, _, N = w1.shape
+
+    if global_num_experts == -1:
+        global_num_experts = E
+
+    gammas = routing_data.gate_scal if routing_data else None
+
+    if quant_config.use_mxfp4_w4a16:
+        act = FusedActivation(
+            FnSpecs("swiglu", triton_kernels.swiglu.swiglu_fn, ("alpha", "limit")),
+            (swiglu_alpha, swiglu_limit),
+            2,
+        )
+        intermediate_cache1 = matmul_ogs(
+            hidden_states,
+            w1,
+            quant_config.w1_bias,
+            routing_data,
+            gather_indx=gather_indx,
+            precision_config=quant_config.w1_precision,
+            gammas=gammas if apply_router_weight_on_input else None,
+            fused_activation=act,
+        )
+        intermediate_cache3 = matmul_ogs(
+            intermediate_cache1,
+            w2,
+            quant_config.w2_bias,
+            routing_data,
+            scatter_indx=scatter_indx,
+            precision_config=quant_config.w2_precision,
+            gammas=None if apply_router_weight_on_input else gammas,
+            y=output_tensor,
+        )
+
+    elif quant_config.use_mxfp4_w4a4:
+        from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
+        from aiter.ops.triton.quant_moe import downcast_to_static_fp8
+
+        assert quant_config.w1_precision is not None, (
+            "w1_precision in quant config can't be None"
+        )
+        assert quant_config.w2_precision is not None, (
+            "w2_precision in quant config can't be None"
+        )
+
+        hidden_states = downcast_to_static_fp8(
+            hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
+        )
+
+        intermediate_cache1 = moe_gemm_a8w4(
+            hidden_states,
+            w1.storage.data,
+            None,
+            quant_config.w1_precision.weight_scale.storage.data,
+            quant_config.w1_precision.flex_ctx.lhs_data.scale,
+            quant_config.w2_precision.flex_ctx.lhs_data.scale,
+            quant_config.w1_bias,
+            routing_data,
+            gather_indx=gather_indx,
+            gammas=gammas if apply_router_weight_on_input else None,
+            swizzle_mx_scale="CDNA4_SCALE",
+            out_dtype=torch.float8_e4m3fn,
+            apply_swiglu=True,
+            alpha=swiglu_alpha,
+            limit=swiglu_limit,
+            unpadded_N=unpadded_N_w1,
+            unpadded_K=unpadded_K_w1,
+        )
+
+        intermediate_cache3 = moe_gemm_a8w4(
+            intermediate_cache1,
+            w2.storage.data,
+            None,
+            quant_config.w2_precision.weight_scale.storage.data,
+            quant_config.w2_precision.flex_ctx.lhs_data.scale,
+            None,
+            quant_config.w2_bias,
+            routing_data,
+            scatter_indx=scatter_indx,
+            gammas=None if apply_router_weight_on_input else gammas,
+            swizzle_mx_scale="CDNA4_SCALE",
+            unpadded_N=unpadded_N_w2,
+            unpadded_K=unpadded_K_w2,
+        )
+    else:
+        raise AssertionError(f"Non supported {quant_config=} in fused MoE op")
+
+    return intermediate_cache3
+
+
 def make_routing_data(
     topk_ids: torch.Tensor,
     topk_weights: torch.Tensor,
@@ -487,6 +663,9 @@ class OAITritonExperts(BaseOAITritonExperts):
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
+        if self.quant_config is None:
+            self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG
+
         if expert_map is not None:
             topk_ids = expert_map[topk_ids]
 

vllm/model_executor/layers/fused_moe/layer.py

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

vllm/model_executor/layers/fused_moe/utils.py

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

vllm/model_executor/layers/quantization/base_config.py

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

vllm/model_executor/layers/quantization/mxfp4.py

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

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

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

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

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

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

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

vllm/model_executor/models/gpt_oss.py

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

vllm/platforms/rocm.py

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

vllm/v1/attention/backend.py

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

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -687,6 +687,8 @@ 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"
@@ -982,49 +984,10 @@ 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]
@@ -1215,3 +1178,57 @@ 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,
+            )

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -5,6 +5,7 @@
 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,
@@ -207,3 +208,42 @@ 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,
+        )

vllm/v1/attention/backends/rocm_attn.py

@@ -7,6 +7,7 @@ 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 (
@@ -415,3 +416,46 @@ 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,
+        )

vllm/v1/attention/backends/triton_attn.py

@@ -7,6 +7,7 @@ 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
@@ -596,3 +597,42 @@ 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,
+        )