[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>
amd-sourjya committed
Jul 27, 2026 at 14:01 UTC
1053e248f02f453390fbaadbecd2b94beff2fbb4
8 files changed
+234
-17
tests/evals/gsm8k/configs/Qwen1.5-MoE-A2.7B-Chat-INT8.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "amd/Qwen1.5-MoE-A2.7B-Chat-w-int8-a-int8-sym"
2
+accuracy_threshold: 0.50
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt
+1
@@ -1,6 +1,7 @@
1
Qwen3-0.6B-FP8.yaml
2
Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
3
Qwen1.5-MoE-W4A16-CT.yaml
4
+Qwen1.5-MoE-A2.7B-Chat-INT8.yaml
5
DeepSeek-V2-Lite-Instruct-FP8.yaml
6
Qwen3-Next-FP8-EP2_MI355.yaml
7
Qwen3-30B-A3B-Thinking-2507-FP8.yaml
tests/quantization/test_int8_moe_oracle.py
new
+105
@@ -0,0 +1,105 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for INT8 (W8A8) fused-MoE oracle backend selection.
5
+
6
+These exercise ``select_int8_moe_backend`` only (no kernels are launched), so
7
+they run on any platform where the Triton INT8 MoE kernel is available — CUDA
8
+(SM >= 7.5) or ROCm — not just gfx950.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+
14
+from vllm.model_executor.layers.fused_moe.config import (
15
+ FusedMoEConfig,
16
+ FusedMoEParallelConfig,
17
+ RoutingMethodType,
18
+)
19
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
20
+ Int8MoeBackend,
21
+ select_int8_moe_backend,
22
+)
23
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
+ kInt8DynamicTensorSym,
25
+ kInt8DynamicTokenSym,
26
+ kInt8StaticChannelSym,
27
+ kInt8StaticTensorSym,
28
+)
29
+from vllm.platforms import current_platform
30
+
31
+# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on
32
+# ROCm CDNA GPUs. Gate on that rather than on a specific arch.
33
+INT8_MOE_SUPPORTED = (
34
+ current_platform.is_cuda() and current_platform.has_device_capability((7, 5))
35
+) or current_platform.is_rocm()
36
+
37
+requires_int8_moe = pytest.mark.skipif(
38
+ not INT8_MOE_SUPPORTED,
39
+ reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)",
40
+)
41
+
42
+
43
+def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
44
+ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
45
+
46
+ return FusedMoEConfig(
47
+ num_experts=8,
48
+ experts_per_token=2,
49
+ hidden_dim=256,
50
+ intermediate_size=256,
51
+ num_local_experts=8,
52
+ num_logical_experts=8,
53
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
54
+ activation=MoEActivation.SILU,
55
+ in_dtype=torch.bfloat16,
56
+ device="cuda",
57
+ routing_method=RoutingMethodType.Renormalize,
58
+ moe_backend=moe_backend,
59
+ )
60
+
61
+
62
+@requires_int8_moe
63
+@pytest.mark.parametrize(
64
+ "weight_key,activation_key",
65
+ [
66
+ # per-channel weight + dynamic per-token activation
67
+ (kInt8StaticChannelSym, kInt8DynamicTokenSym),
68
+ # per-tensor weight + dynamic per-tensor activation
69
+ (kInt8StaticTensorSym, kInt8DynamicTensorSym),
70
+ ],
71
+)
72
+def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key):
73
+ """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor
74
+ weights) select the Triton backend."""
75
+ config = _make_int8_moe_config()
76
+ backend, experts_cls = select_int8_moe_backend(
77
+ config, weight_key=weight_key, activation_key=activation_key
78
+ )
79
+ assert backend == Int8MoeBackend.TRITON
80
+ assert experts_cls is not None
81
+
82
+
83
+@requires_int8_moe
84
+def test_int8_explicit_moe_backend_triton():
85
+ """An explicit --moe-backend triton selects the Triton INT8 backend."""
86
+ config = _make_int8_moe_config(moe_backend="triton")
87
+ backend, experts_cls = select_int8_moe_backend(
88
+ config,
89
+ weight_key=kInt8StaticChannelSym,
90
+ activation_key=kInt8DynamicTokenSym,
91
+ )
92
+ assert backend == Int8MoeBackend.TRITON
93
+ assert experts_cls is not None
94
+
95
+
96
+@requires_int8_moe
97
+def test_int8_unsupported_moe_backend_raises():
98
+ """An unsupported --moe-backend for INT8 MoE raises a clear error."""
99
+ config = _make_int8_moe_config(moe_backend="cutlass")
100
+ with pytest.raises(ValueError, match="not supported for Int8 MoE"):
101
+ select_int8_moe_backend(
102
+ config,
103
+ weight_key=kInt8StaticChannelSym,
104
+ activation_key=kInt8DynamicTokenSym,
105
+ )
tests/quantization/test_quark.py
+1
-1
@@ -150,7 +150,7 @@ def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
150
@pytest.mark.parametrize("tp", [1])
151
def test_quark_int8_w8a8_moe(vllm_runner, tp):
152
"""Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model."""
153
- model_path = "nameistoken/tiny-qwen3-moe-w8a8-int8-quark"
153
+ model_path = "amd/tiny-qwen3-moe-w8a8-int8"
154
with vllm_runner(
155
model_path,
156
enforce_eager=True,
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+15
-3
@@ -45,9 +45,11 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
kFp8StaticTensorSym,
46
kInt4Static,
47
kInt4Static32,
48
+ kInt8DynamicTensorSym,
49
kInt8DynamicTokenSym,
50
kInt8Static,
51
kInt8StaticChannelSym,
52
+ kInt8StaticTensorSym,
53
)
54
from vllm.platforms import current_platform
55
from vllm.triton_utils import tl
@@ -103,15 +105,25 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
105
weight_key: QuantKey | None,
106
activation_key: QuantKey | None,
107
) -> bool:
106
- # INT8 requires at least 7.5 (Turing).
108
+ # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs
109
+ # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and
110
+ # the Triton int8_w8a8 fused MoE kernel handles them.
111
device_supports_int8 = (
112
current_platform.is_cuda()
113
and current_platform.has_device_capability((7, 5))
110
- )
114
+ ) or current_platform.is_rocm()
115
116
supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
117
if device_supports_int8:
114
- supported.append((kInt8StaticChannelSym, kInt8DynamicTokenSym))
118
+ # Activations are consumed as float and quantized to int8
119
+ # dynamically inside the kernel, so only dynamic-activation int8
120
+ # schemes are supported (static-activation int8 is not).
121
+ supported += [
122
+ # per-channel weight + dynamic per-token activation
123
+ (kInt8StaticChannelSym, kInt8DynamicTokenSym),
124
+ # per-tensor weight + dynamic per-tensor activation
125
+ (kInt8StaticTensorSym, kInt8DynamicTensorSym),
126
+ ]
127
if current_platform.supports_fp8():
128
supported += [
129
(kFp8Static128BlockSym, kFp8Dynamic128Sym),
vllm/model_executor/layers/fused_moe/oracle/int8.py
+3
-1
@@ -166,7 +166,9 @@ def select_int8_moe_backend(
166
logger.debug_once(_make_log_unsupported(backend, reason))
167
168
raise NotImplementedError(
169
- "No Int8 MoE backend supports the deployment configuration."
169
+ "No Int8 MoE backend supports the deployment configuration "
170
+ f"(weight_key={weight_key}, activation_key={activation_key}). "
171
+ "Set `VLLM_LOGGING_LEVEL=DEBUG` to see per-backend unsupported reasons."
172
)
173
174
vllm/model_executor/layers/quantization/quark/quark_moe.py
+102
-12
@@ -33,6 +33,13 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
33
make_fp8_moe_quant_config,
34
select_fp8_moe_backend,
35
)
36
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
37
+ Int8MoeBackend,
38
+ convert_to_int8_moe_kernel_format,
39
+ make_int8_moe_kernel,
40
+ make_int8_moe_quant_config,
41
+ select_int8_moe_backend,
42
+)
43
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
44
TRITON_BACKENDS,
45
Mxfp4MoeBackend,
@@ -59,6 +66,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
66
kFp8DynamicTokenSym,
67
kFp8StaticChannelSym,
68
kFp8StaticTensorSym,
69
+ kInt8DynamicTensorSym,
70
+ kInt8DynamicTokenSym,
71
+ kInt8StaticChannelSym,
72
+ kInt8StaticTensorSym,
73
kMxfp4Dynamic,
74
kNvfp4Dynamic,
75
kNvfp4Static,
@@ -502,6 +513,35 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
513
self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor")
514
self.static_input_scales = not self.input_quant.get("is_dynamic", False)
515
516
+ self.moe_quant_config: FusedMoEQuantConfig | None = None
517
+ self.moe_kernel: mk.FusedMoEKernel | None = None
518
+ self.int8_backend: Int8MoeBackend | None = None
519
+ self.experts_cls: type[mk.FusedMoEExperts] | None = None
520
+
521
+ # Dynamic-activation INT8 MoE goes through the oracle + modular kernel.
522
+ # The modular TritonExperts kernel consumes float activations and
523
+ # quantizes them to int8 itself, so it cannot apply a loaded static
524
+ # activation scale (this matches CompressedTensorsW8A8Int8MoEMethod).
525
+ # TODO: Static-activation INT8 therefore stays on the legacy fused_experts
526
+ # path (see apply()) for now, preserving pre-refactor behavior.
527
+ # Needs to be migrated to expert backend.
528
+ if not self.static_input_scales:
529
+ # Map the Quark weight scheme to oracle quant keys. Per-channel
530
+ # weights pair with dynamic per-token activations; per-tensor
531
+ # weights with dynamic per-tensor activations.
532
+ if self.weight_qscheme == "per_channel":
533
+ weight_key = kInt8StaticChannelSym
534
+ activation_key = kInt8DynamicTokenSym
535
+ else:
536
+ weight_key = kInt8StaticTensorSym
537
+ activation_key = kInt8DynamicTensorSym
538
+
539
+ self.int8_backend, self.experts_cls = select_int8_moe_backend(
540
+ config=moe,
541
+ weight_key=weight_key,
542
+ activation_key=activation_key,
543
+ )
544
+
545
def create_weights(
546
self,
547
layer: torch.nn.Module,
@@ -563,7 +603,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
603
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
604
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
605
else:
566
- # per-tensor: one scalar per expert
606
+ # per-tensor: one scalar per expert (two for the fused w1/w3)
607
w13_weight_scale = torch.nn.Parameter(
608
torch.ones(num_experts, 2, dtype=torch.float32),
609
requires_grad=False,
@@ -582,6 +622,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
622
623
# INPUT_SCALES
624
if self.static_input_scales:
625
+ # Static activations: the per-expert scales are loaded from the
626
+ # checkpoint (used by the legacy fused_experts path).
627
w13_input_scale = torch.nn.Parameter(
628
torch.ones(num_experts, dtype=torch.float32),
629
requires_grad=False,
@@ -596,6 +638,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
638
layer.register_parameter("w2_input_scale", w2_input_scale)
639
set_weight_attrs(w2_input_scale, extra_weight_attrs)
640
else:
641
+ # Dynamic activations are quantized in-kernel (no stored scale).
642
layer.w13_input_scale = None
643
layer.w2_input_scale = None
644
@@ -673,7 +716,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
716
if hasattr(layer, attr):
717
delattr(layer, attr)
718
676
- # For static input scales, collapse per-expert scales to single max
719
+ # For static input scales, collapse the per-expert scales to a single
720
+ # value (the legacy fused_experts path expects one scale per layer).
721
if self.static_input_scales:
722
if layer.w13_input_scale is None or layer.w2_input_scale is None:
723
raise ValueError(
@@ -709,7 +753,8 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
753
),
754
)
755
712
- # For per-tensor weights, merge w1/w3 scales into single per-expert
756
+ # For per-tensor weights, merge the w1/w3 scales into a single
757
+ # per-expert scale (dequant -> requant at the max scale).
758
if self.weight_qscheme == "per_tensor":
759
assert layer.w13_weight_scale is not None
760
shard_size = layer.intermediate_size_per_partition
@@ -734,10 +779,42 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
779
max_w13_scales, requires_grad=False
780
)
781
782
+ # Dynamic activations run through the oracle's modular kernel; static
783
+ # activations use the legacy fused_experts path in apply().
784
+ if not self.static_input_scales:
785
+ assert self.int8_backend is not None
786
+ assert self.experts_cls is not None
787
+ w13, w2 = convert_to_int8_moe_kernel_format(
788
+ int8_backend=self.int8_backend,
789
+ w13=layer.w13_weight,
790
+ w2=layer.w2_weight,
791
+ layer=layer,
792
+ w13_scale=layer.w13_weight_scale,
793
+ )
794
+ replace_parameter(layer, "w13_weight", w13)
795
+ replace_parameter(layer, "w2_weight", w2)
796
+
797
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
798
+ assert self.moe_quant_config is not None
799
+
800
+ if not self.static_input_scales:
801
+ assert self.int8_backend is not None
802
+ assert self.experts_cls is not None
803
+ self.moe_kernel = make_int8_moe_kernel(
804
+ int8_backend=self.int8_backend,
805
+ moe_quant_config=self.moe_quant_config,
806
+ moe_config=self.moe,
807
+ experts_cls=self.experts_cls,
808
+ routing_tables=layer._expert_routing_tables(),
809
+ layer=layer,
810
+ )
811
+
812
def get_fused_moe_quant_config(
813
self, layer: torch.nn.Module
814
) -> FusedMoEQuantConfig | None:
740
- if self.weight_qscheme == "per_channel" and not self.static_input_scales:
815
+ # Static-activation INT8 has no oracle backend (it uses the legacy
816
+ # fused_experts path); build its config directly.
817
+ if self.int8_backend is None:
818
return int8_w8a8_moe_quant_config(
819
w1_scale=layer.w13_weight_scale,
820
w2_scale=layer.w2_weight_scale,
@@ -745,21 +822,18 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
822
a2_scale=layer.w2_input_scale,
823
w1_bias=getattr(layer, "w13_bias", None),
824
w2_bias=getattr(layer, "w2_bias", None),
748
- per_act_token_quant=True,
825
+ per_act_token_quant=False,
826
)
750
- is_dynamic = not self.static_input_scales
751
- is_per_channel = self.weight_qscheme == "per_channel"
752
- return FusedMoEQuantConfig.make(
753
- torch.int8,
827
+ return make_int8_moe_quant_config(
828
+ int8_backend=self.int8_backend,
829
w1_scale=layer.w13_weight_scale,
830
w2_scale=layer.w2_weight_scale,
831
a1_scale=layer.w13_input_scale,
832
a2_scale=layer.w2_input_scale,
833
w1_bias=getattr(layer, "w13_bias", None),
834
w2_bias=getattr(layer, "w2_bias", None),
760
- per_act_token_quant=is_dynamic,
761
- per_out_ch_quant=is_per_channel,
762
- block_shape=None,
835
+ per_act_token_quant=(self.weight_qscheme == "per_channel"),
836
+ layer=layer,
837
)
838
839
def apply(
@@ -771,6 +845,22 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
845
shared_experts: SharedExperts | None,
846
shared_experts_input: torch.Tensor | None,
847
) -> torch.Tensor:
848
+ if self.moe_kernel is not None:
849
+ return self.moe_kernel.apply(
850
+ hidden_states=x,
851
+ w1=layer.w13_weight,
852
+ w2=layer.w2_weight,
853
+ topk_weights=topk_weights,
854
+ topk_ids=topk_ids,
855
+ activation=layer.activation,
856
+ global_num_experts=layer.global_num_experts,
857
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
858
+ expert_map=layer.expert_map,
859
+ shared_experts_input=shared_experts_input,
860
+ )
861
+
862
+ # Static-activation INT8 MoE: legacy monolithic path (the modular kernel
863
+ # quantizes activations dynamically and cannot apply a loaded scale).
864
from vllm.model_executor.layers.fused_moe import fused_experts
865
866
return fused_experts(
vllm/model_executor/layers/quantization/utils/quant_utils.py
+2
@@ -190,6 +190,8 @@ kInt4Static32Asym = QuantKey(
190
191
kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
192
kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
193
+kInt8StaticTensorSym = QuantKey(torch.int8, kStaticTensorScale, symmetric=True)
194
+kInt8DynamicTensorSym = QuantKey(torch.int8, kDynamicTensorScale, symmetric=True)
195
196
# INT4 W4A8 quantization keys
197