[ROCm][Quantization][5/N] Refactor quark_moe w8a8-int8 w/ oracle (#46765)

Signed-off-by: amd-sourjya <amd-sourjya@users.noreply.github.com>
Co-authored-by: amd-sourjya <amd-sourjya@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml

@@ -0,0 +1,5 @@
+model_name: "amd/Qwen1.5-MoE-A2.7B-Chat-w-int8-a-int8-sym"
+accuracy_threshold: 0.50
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 4096"

tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt

@@ -1,6 +1,7 @@
 Qwen3-0.6B-FP8.yaml
 Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
 Qwen1.5-MoE-W4A16-CT.yaml
+Qwen1.5-MoE-A2.7B-Chat-INT8.yaml
 DeepSeek-V2-Lite-Instruct-FP8.yaml
 Qwen3-Next-FP8-EP2_MI355.yaml
 Qwen3-30B-A3B-Thinking-2507-FP8.yaml

tests/quantization/test_int8_moe_oracle.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for INT8 (W8A8) fused-MoE oracle backend selection.
+
+These exercise ``select_int8_moe_backend`` only (no kernels are launched), so
+they run on any platform where the Triton INT8 MoE kernel is available — CUDA
+(SM >= 7.5) or ROCm — not just gfx950.
+"""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    RoutingMethodType,
+)
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
+    Int8MoeBackend,
+    select_int8_moe_backend,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kInt8DynamicTensorSym,
+    kInt8DynamicTokenSym,
+    kInt8StaticChannelSym,
+    kInt8StaticTensorSym,
+)
+from vllm.platforms import current_platform
+
+# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on
+# ROCm CDNA GPUs. Gate on that rather than on a specific arch.
+INT8_MOE_SUPPORTED = (
+    current_platform.is_cuda() and current_platform.has_device_capability((7, 5))
+) or current_platform.is_rocm()
+
+requires_int8_moe = pytest.mark.skipif(
+    not INT8_MOE_SUPPORTED,
+    reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)",
+)
+
+
+def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
+    from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+
+    return FusedMoEConfig(
+        num_experts=8,
+        experts_per_token=2,
+        hidden_dim=256,
+        intermediate_size=256,
+        num_local_experts=8,
+        num_logical_experts=8,
+        moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
+        activation=MoEActivation.SILU,
+        in_dtype=torch.bfloat16,
+        device="cuda",
+        routing_method=RoutingMethodType.Renormalize,
+        moe_backend=moe_backend,
+    )
+
+
+@requires_int8_moe
+@pytest.mark.parametrize(
+    "weight_key,activation_key",
+    [
+        # per-channel weight + dynamic per-token activation
+        (kInt8StaticChannelSym, kInt8DynamicTokenSym),
+        # per-tensor weight + dynamic per-tensor activation
+        (kInt8StaticTensorSym, kInt8DynamicTensorSym),
+    ],
+)
+def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key):
+    """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor
+    weights) select the Triton backend."""
+    config = _make_int8_moe_config()
+    backend, experts_cls = select_int8_moe_backend(
+        config, weight_key=weight_key, activation_key=activation_key
+    )
+    assert backend == Int8MoeBackend.TRITON
+    assert experts_cls is not None
+
+
+@requires_int8_moe
+def test_int8_explicit_moe_backend_triton():
+    """An explicit --moe-backend triton selects the Triton INT8 backend."""
+    config = _make_int8_moe_config(moe_backend="triton")
+    backend, experts_cls = select_int8_moe_backend(
+        config,
+        weight_key=kInt8StaticChannelSym,
+        activation_key=kInt8DynamicTokenSym,
+    )
+    assert backend == Int8MoeBackend.TRITON
+    assert experts_cls is not None
+
+
+@requires_int8_moe
+def test_int8_unsupported_moe_backend_raises():
+    """An unsupported --moe-backend for INT8 MoE raises a clear error."""
+    config = _make_int8_moe_config(moe_backend="cutlass")
+    with pytest.raises(ValueError, match="not supported for Int8 MoE"):
+        select_int8_moe_backend(
+            config,
+            weight_key=kInt8StaticChannelSym,
+            activation_key=kInt8DynamicTokenSym,
+        )

tests/quantization/test_quark.py

@@ -150,7 +150,7 @@ def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
 @pytest.mark.parametrize("tp", [1])
 def test_quark_int8_w8a8_moe(vllm_runner, tp):
     """Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model."""
-    model_path = "nameistoken/tiny-qwen3-moe-w8a8-int8-quark"
+    model_path = "amd/tiny-qwen3-moe-w8a8-int8"
     with vllm_runner(
         model_path,
         enforce_eager=True,

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

@@ -45,9 +45,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticTensorSym,
     kInt4Static,
     kInt4Static32,
+    kInt8DynamicTensorSym,
     kInt8DynamicTokenSym,
     kInt8Static,
     kInt8StaticChannelSym,
+    kInt8StaticTensorSym,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl
@@ -103,15 +105,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        # INT8 requires at least 7.5 (Turing).
+        # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs
+        # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and
+        # the Triton int8_w8a8 fused MoE kernel handles them.
         device_supports_int8 = (
             current_platform.is_cuda()
             and current_platform.has_device_capability((7, 5))
-        )
+        ) or current_platform.is_rocm()
 
         supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
         if device_supports_int8:
-            supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym))
+            # Activations are consumed as float and quantized to int8
+            # dynamically inside the kernel, so only dynamic-activation int8
+            # schemes are supported (static-activation int8 is not).
+            supported += [
+                # per-channel weight + dynamic per-token activation
+                (kInt8StaticChannelSym, kInt8DynamicTokenSym),
+                # per-tensor weight + dynamic per-tensor activation
+                (kInt8StaticTensorSym, kInt8DynamicTensorSym),
+            ]
         if current_platform.supports_fp8():
             supported += [
                 (kFp8Static128BlockSym, kFp8Dynamic128Sym),

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

@@ -166,7 +166,9 @@ def select_int8_moe_backend(
                 logger.debug_once(_make_log_unsupported(backend, reason))
 
     raise NotImplementedError(
-        "No Int8 MoE backend supports the deployment configuration."
+        "No Int8 MoE backend supports the deployment configuration "
+        f"(weight_key={weight_key}, activation_key={activation_key}). "
+        "Set `VLLM_LOGGING_LEVEL=DEBUG` to see per-backend unsupported reasons."
     )
 
 

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

@@ -33,6 +33,13 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
     make_fp8_moe_quant_config,
     select_fp8_moe_backend,
 )
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
+    Int8MoeBackend,
+    convert_to_int8_moe_kernel_format,
+    make_int8_moe_kernel,
+    make_int8_moe_quant_config,
+    select_int8_moe_backend,
+)
 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
     TRITON_BACKENDS,
     Mxfp4MoeBackend,
@@ -59,6 +66,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8DynamicTokenSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
+    kInt8DynamicTensorSym,
+    kInt8DynamicTokenSym,
+    kInt8StaticChannelSym,
+    kInt8StaticTensorSym,
     kMxfp4Dynamic,
     kNvfp4Dynamic,
     kNvfp4Static,
@@ -502,6 +513,35 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
         self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor")
         self.static_input_scales = not self.input_quant.get("is_dynamic", False)
 
+        self.moe_quant_config: FusedMoEQuantConfig | None = None
+        self.moe_kernel: mk.FusedMoEKernel | None = None
+        self.int8_backend: Int8MoeBackend | None = None
+        self.experts_cls: type[mk.FusedMoEExperts] | None = None
+
+        # Dynamic-activation INT8 MoE goes through the oracle + modular kernel.
+        # The modular TritonExperts kernel consumes float activations and
+        # quantizes them to int8 itself, so it cannot apply a loaded static
+        # activation scale (this matches CompressedTensorsW8A8Int8MoEMethod).
+        # TODO: Static-activation INT8 therefore stays on the legacy fused_experts
+        # path (see apply()) for now, preserving pre-refactor behavior.
+        # Needs to be migrated to expert backend.
+        if not self.static_input_scales:
+            # Map the Quark weight scheme to oracle quant keys. Per-channel
+            # weights pair with dynamic per-token activations; per-tensor
+            # weights with dynamic per-tensor activations.
+            if self.weight_qscheme == "per_channel":
+                weight_key = kInt8StaticChannelSym
+                activation_key = kInt8DynamicTokenSym
+            else:
+                weight_key = kInt8StaticTensorSym
+                activation_key = kInt8DynamicTensorSym
+
+            self.int8_backend, self.experts_cls = select_int8_moe_backend(
+                config=moe,
+                weight_key=weight_key,
+                activation_key=activation_key,
+            )
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -563,7 +603,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
             set_weight_attrs(w13_weight_scale, extra_weight_attrs)
             set_weight_attrs(w2_weight_scale, extra_weight_attrs)
         else:
-            # per-tensor: one scalar per expert
+            # per-tensor: one scalar per expert (two for the fused w1/w3)
             w13_weight_scale = torch.nn.Parameter(
                 torch.ones(num_experts, 2, dtype=torch.float32),
                 requires_grad=False,
@@ -582,6 +622,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
 
         # INPUT_SCALES
         if self.static_input_scales:
+            # Static activations: the per-expert scales are loaded from the
+            # checkpoint (used by the legacy fused_experts path).
             w13_input_scale = torch.nn.Parameter(
                 torch.ones(num_experts, dtype=torch.float32),
                 requires_grad=False,
@@ -596,6 +638,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
             layer.register_parameter("w2_input_scale", w2_input_scale)
             set_weight_attrs(w2_input_scale, extra_weight_attrs)
         else:
+            # Dynamic activations are quantized in-kernel (no stored scale).
             layer.w13_input_scale = None
             layer.w2_input_scale = None
 
@@ -673,7 +716,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
             if hasattr(layer, attr):
                 delattr(layer, attr)
 
-        # For static input scales, collapse per-expert scales to single max
+        # For static input scales, collapse the per-expert scales to a single
+        # value (the legacy fused_experts path expects one scale per layer).
         if self.static_input_scales:
             if layer.w13_input_scale is None or layer.w2_input_scale is None:
                 raise ValueError(
@@ -709,7 +753,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
                         ),
                     )
 
-        # For per-tensor weights, merge w1/w3 scales into single per-expert
+        # For per-tensor weights, merge the w1/w3 scales into a single
+        # per-expert scale (dequant -> requant at the max scale).
         if self.weight_qscheme == "per_tensor":
             assert layer.w13_weight_scale is not None
             shard_size = layer.intermediate_size_per_partition
@@ -734,10 +779,42 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
                 max_w13_scales, requires_grad=False
             )
 
+        # Dynamic activations run through the oracle's modular kernel; static
+        # activations use the legacy fused_experts path in apply().
+        if not self.static_input_scales:
+            assert self.int8_backend is not None
+            assert self.experts_cls is not None
+            w13, w2 = convert_to_int8_moe_kernel_format(
+                int8_backend=self.int8_backend,
+                w13=layer.w13_weight,
+                w2=layer.w2_weight,
+                layer=layer,
+                w13_scale=layer.w13_weight_scale,
+            )
+            replace_parameter(layer, "w13_weight", w13)
+            replace_parameter(layer, "w2_weight", w2)
+
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+
+        if not self.static_input_scales:
+            assert self.int8_backend is not None
+            assert self.experts_cls is not None
+            self.moe_kernel = make_int8_moe_kernel(
+                int8_backend=self.int8_backend,
+                moe_quant_config=self.moe_quant_config,
+                moe_config=self.moe,
+                experts_cls=self.experts_cls,
+                routing_tables=layer._expert_routing_tables(),
+                layer=layer,
+            )
+
     def get_fused_moe_quant_config(
         self, layer: torch.nn.Module
     ) -> FusedMoEQuantConfig | None:
-        if self.weight_qscheme == "per_channel" and not self.static_input_scales:
+        # Static-activation INT8 has no oracle backend (it uses the legacy
+        # fused_experts path); build its config directly.
+        if self.int8_backend is None:
             return int8_w8a8_moe_quant_config(
                 w1_scale=layer.w13_weight_scale,
                 w2_scale=layer.w2_weight_scale,
@@ -745,21 +822,18 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
                 a2_scale=layer.w2_input_scale,
                 w1_bias=getattr(layer, "w13_bias", None),
                 w2_bias=getattr(layer, "w2_bias", None),
-                per_act_token_quant=True,
+                per_act_token_quant=False,
             )
-        is_dynamic = not self.static_input_scales
-        is_per_channel = self.weight_qscheme == "per_channel"
-        return FusedMoEQuantConfig.make(
-            torch.int8,
+        return make_int8_moe_quant_config(
+            int8_backend=self.int8_backend,
             w1_scale=layer.w13_weight_scale,
             w2_scale=layer.w2_weight_scale,
             a1_scale=layer.w13_input_scale,
             a2_scale=layer.w2_input_scale,
             w1_bias=getattr(layer, "w13_bias", None),
             w2_bias=getattr(layer, "w2_bias", None),
-            per_act_token_quant=is_dynamic,
-            per_out_ch_quant=is_per_channel,
-            block_shape=None,
+            per_act_token_quant=(self.weight_qscheme == "per_channel"),
+            layer=layer,
         )
 
     def apply(
@@ -771,6 +845,22 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
         shared_experts: SharedExperts | None,
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
+        if self.moe_kernel is not None:
+            return self.moe_kernel.apply(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=layer.activation,
+                global_num_experts=layer.global_num_experts,
+                apply_router_weight_on_input=layer.apply_router_weight_on_input,
+                expert_map=layer.expert_map,
+                shared_experts_input=shared_experts_input,
+            )
+
+        # Static-activation INT8 MoE: legacy monolithic path (the modular kernel
+        # quantizes activations dynamically and cannot apply a loaded scale).
         from vllm.model_executor.layers.fused_moe import fused_experts
 
         return fused_experts(

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

@@ -190,6 +190,8 @@ kInt4Static32Asym = QuantKey(
 
 kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
 kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
+kInt8StaticTensorSym = QuantKey(torch.int8, kStaticTensorScale, symmetric=True)
+kInt8DynamicTensorSym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=True)
 
 # INT4 W4A8 quantization keys