@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 986e2f870a

[Refactor][ROCm] Migrate the RDNA3 W4A16 MoE to the oracle/experts pa… (#55522)

Signed-off-by: JartX <sagformas@epdcenter.es>

JartX committed Sep 12, 2026 at 15:20 UTC 986e2f870abc8f87a1830af7d335961b3e4234b2
10 files changed +453 -377
tests/kernels/quantization/test_rdna3_compile_guards.py
+68 -59
@@ -34,6 +34,13 @@ from vllm.platforms import current_platform
34 if not current_platform.is_rocm():
35 pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True)
36
37 +from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
38 + kInt4Static,
39 + kInt4Static32,
40 + kInt4Static32Asym,
41 + kInt4StaticAsym,
42 + kInt8Static,
43 +)
44 from vllm.platforms.rocm import on_gfx1100 # noqa: E402
45
46 gfx1100_only = pytest.mark.skipif(
@@ -166,15 +173,13 @@ def test_op_absent_on_non_gfx1100(op_name):
173
174 @not_gfx1100
175 def test_rocm_moe_not_supported_on_non_gfx1100():
169 - """rocm_moe_rdna.is_supported() must return False on non-gfx1100 hardware."""
170 - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
171 - rocm_moe_rdna,
176 + """The RDNA3 MoE experts must not be selectable on non-gfx1100 hardware."""
177 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
178 + Rdna3WNA16Experts,
179 )
180
174 - wq = type("WQ", (), {"num_bits": 4})()
175 - assert rocm_moe_rdna.is_supported(wq) is False, (
176 - "rocm_moe_rdna.is_supported() returned True on non-gfx1100 — "
177 - "dispatch guard is broken"
181 + assert Rdna3WNA16Experts._supports_current_device() is False, (
182 + "Rdna3WNA16Experts reported support on non-gfx1100 — dispatch guard is broken"
183 )
184
185
@@ -359,54 +364,58 @@ class TestCustomOpsGuards:
364 # ============================================================================
365
366
362 -class _FakeWeightQuant:
363 - """Minimal stand-in for a weight quantization config."""
364 -
365 - def __init__(self, num_bits):
366 - self.num_bits = num_bits
367 -
368 -
367 class TestMoEDispatchMocked:
368 """Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
369
372 - def test_is_supported_false_when_mocked_cdna(self):
373 - """rocm_moe_rdna.is_supported() must return False when not on gfx1100."""
374 - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
375 - rocm_moe_rdna,
370 + def test_kernel_unavailable_when_mocked_cdna(self):
371 + """The device gate must reject when not on gfx1100."""
372 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
373 + rdna3_moe_kernel_available,
374 )
375
376 with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
379 - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
377 + assert rdna3_moe_kernel_available() is False
378
381 - @pytest.mark.parametrize("num_bits", [2, 3, 8, 16])
382 - def test_is_supported_rejects_non_w4(self, num_bits):
383 - """is_supported() rejects non-4-bit even before checking arch."""
384 - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
385 - rocm_moe_rdna,
379 + @pytest.mark.parametrize(
380 + "weight_key",
381 + [kInt8Static, kInt4StaticAsym, kInt4Static32Asym, None],
382 + )
383 + def test_quant_scheme_rejects_non_symmetric_int4(self, weight_key):
384 + """Only symmetric int4 weight-only schemes reach the kernel."""
385 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
386 + Rdna3WNA16Experts,
387 )
388
388 - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
389 + assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is False
390
390 - def test_is_supported_false_when_op_missing(self):
391 - """is_supported() returns False when the C++ op doesn't exist."""
392 - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
393 - rocm_moe_rdna,
391 + @pytest.mark.parametrize("weight_key", [kInt4Static, kInt4Static32])
392 + def test_quant_scheme_accepts_symmetric_int4(self, weight_key):
393 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
394 + Rdna3WNA16Experts,
395 + )
396 +
397 + assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is True
398 +
399 + def test_kernel_unavailable_when_op_missing(self):
400 + """The device gate returns False when the C++ op doesn't exist."""
401 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
402 + rdna3_moe_kernel_available,
403 )
404
405 fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
406 with patch.object(torch, "ops", create=True) as mock_ops:
407 mock_ops._rocm_C = fake_rocm_c
399 - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
408 + assert rdna3_moe_kernel_available() is False
409
401 - def test_is_supported_false_when_rocm_c_absent(self):
402 - """is_supported() returns False when _rocm_C doesn't exist at all."""
403 - from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
404 - rocm_moe_rdna,
410 + def test_kernel_unavailable_when_rocm_c_absent(self):
411 + """The device gate returns False when _rocm_C doesn't exist at all."""
412 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
413 + rdna3_moe_kernel_available,
414 )
415
416 fake_ops = type("FakeOps", (), {})()
417 with patch.object(torch, "ops", fake_ops):
409 - assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
418 + assert rdna3_moe_kernel_available() is False
419
420
421 class TestDenseKernelSelectionMocked:
@@ -471,30 +480,30 @@ class TestDenseKernelSelectionMocked:
480 )
481
482
474 -class TestCompressedTensorsMoEDispatchGuard:
475 - """Verify compressed_tensors_moe.py only enters rocm_moe_rdna under is_rocm()."""
483 +class TestWNA16OracleWiring:
484 + """The RDNA3 backend must reach the kernel only through the oracle."""
485
477 - def test_rocm_guard_in_dispatch_source(self):
478 - """The rocm_moe_rdna import and call must be inside an is_rocm() check."""
479 - src = _read_pkg_source_or_skip(
480 - "model_executor",
481 - "layers",
482 - "quantization",
483 - "compressed_tensors",
484 - "compressed_tensors_moe",
485 - "compressed_tensors_moe.py",
486 + def test_backend_maps_to_rdna3_experts(self):
487 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
488 + Rdna3WNA16Experts,
489 + )
490 + from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
491 + WNA16MoEBackend,
492 + backend_to_kernel_cls,
493 + map_wna16_backend,
494 )
487 - lines = src.splitlines()
495
489 - for i, line in enumerate(lines, 1):
490 - stripped = line.strip()
491 - if "rocm_moe" in stripped and not stripped.startswith("#"):
492 - found_guard = False
493 - for j in range(i - 1, max(0, i - 15), -1):
494 - if "is_rocm()" in lines[j - 1]:
495 - found_guard = True
496 - break
497 - assert found_guard, (
498 - f"L{i}: rocm_moe_rdna reference not protected by "
499 - f"is_rocm() guard: {stripped}"
500 - )
496 + assert backend_to_kernel_cls(WNA16MoEBackend.RDNA3) == [Rdna3WNA16Experts]
497 + assert map_wna16_backend("rdna3") == WNA16MoEBackend.RDNA3
498 +
499 + def test_backend_is_offered_before_the_triton_fallback(self):
500 + """Priority order: the native kernel outranks Triton when supported."""
501 + from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
502 + WNA16MoEBackend,
503 + _get_priority_backends,
504 + )
505 +
506 + backends = _get_priority_backends()
507 + assert backends.index(WNA16MoEBackend.RDNA3) < backends.index(
508 + WNA16MoEBackend.TRITON
509 + )
tests/quantization/test_moe_wna16.py
+34
@@ -93,6 +93,40 @@ def test_moe_wna16_accepts_channelwise_gptq_activation_order():
93 False,
94 "MoeWNA16 checkpoint layout",
95 ),
96 + (
97 + WNA16MoEBackend.RDNA3,
98 + AutoGPTQConfig(4, 128, False, True, False, {}, {}),
99 + False,
100 + False,
101 + "compressed-tensors",
102 + ),
103 + (
104 + WNA16MoEBackend.RDNA3,
105 + QuantizationArgs(
106 + num_bits=4,
107 + type=QuantizationType.INT,
108 + strategy=QuantizationStrategy.GROUP,
109 + symmetric=False,
110 + dynamic=False,
111 + group_size=128,
112 + ),
113 + True,
114 + False,
115 + "asymmetric",
116 + ),
117 + (
118 + WNA16MoEBackend.RDNA3,
119 + QuantizationArgs(
120 + num_bits=4,
121 + type=QuantizationType.INT,
122 + strategy=QuantizationStrategy.CHANNEL,
123 + symmetric=True,
124 + dynamic=False,
125 + ),
126 + False,
127 + False,
128 + "group-wise scales",
129 + ),
130 ],
131 )
132 def test_wna16_oracle_rejects_incompatible_quant_structures(
vllm/config/kernel.py
+2
@@ -141,6 +141,7 @@ MoEBackend = Literal[
141 "flydsl",
142 "hpc",
143 "emulation",
144 + "rdna3",
145 ]
146
147 # Backends that run the mega-MoE model path through the flashinfer moe_ep
@@ -262,6 +263,7 @@ class KernelConfig:
263 - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16
264 (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250)
265 - "flydsl": Use AMD FlyDSL kernels (ROCm only)
266 + - "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only)
267 - "hpc": Use HPC kernels (FP8 and Hopper only)
268 - "emulation": use BF16/FP16 GEMM, dequantizing weights and
269 running QDQ on activations.
vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py new
+235
@@ -0,0 +1,235 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused MoE W4A16 experts on the RDNA3 (gfx1100) HIP kernel.
4 +
5 +``moe_gptq_gemm_rdna3`` is a single HIP kernel launch per GEMM that handles
6 +expert routing + W4A16 dequant + dot product with atomic output accumulation.
7 +
8 +Weight format (per expert, same as the dense RDNA3 W4A16 kernel):
9 + - Packed int32 ``[E, K/8, N]`` with exllama shuffle
10 + - Scales ``[E, groups, N]`` in activation dtype
11 + - Zero points ``[E, groups, N/8]`` packed int32 (synthesized, symmetric only)
12 +"""
13 +
14 +import torch
15 +
16 +import vllm._custom_ops as ops
17 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
18 +from vllm.model_executor.layers.fused_moe.activation import (
19 + MoEActivation,
20 + apply_moe_activation_supported,
21 +)
22 +from vllm.model_executor.layers.fused_moe.config import (
23 + FusedMoEConfig,
24 + FusedMoEParallelConfig,
25 + FusedMoEQuantConfig,
26 +)
27 +from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
28 + moe_align_block_size,
29 +)
30 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
31 + TopKWeightAndReduceNoOP,
32 +)
33 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
34 + QuantKey,
35 + kInt4Static,
36 + kInt4Static32,
37 +)
38 +from vllm.platforms import current_platform
39 +
40 +
41 +def rdna3_moe_kernel_available() -> bool:
42 + """Whether the fused RDNA3 MoE HIP kernel is built into this binary."""
43 + if not current_platform.is_rocm():
44 + return False
45 +
46 + from vllm.platforms.rocm import on_gfx1100
47 +
48 + return (
49 + on_gfx1100()
50 + and hasattr(torch.ops, "_rocm_C")
51 + and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
52 + )
53 +
54 +
55 +class Rdna3WNA16Experts(mk.FusedMoEExpertsModular):
56 + """W4A16 experts backed by ``moe_gptq_gemm_rdna3`` (gfx1100).
57 +
58 + Both GEMMs accumulate atomically, so their destinations are zeroed first.
59 + The second GEMM is given ``output_topk``, which makes it reduce over the
60 + top-k dimension while it accumulates — the ``moe_sum`` launch and the
61 + ``[M * top_k, K]`` intermediate are fused away, hence
62 + ``TopKWeightAndReduceNoOP``.
63 + """
64 +
65 + def __init__(
66 + self,
67 + moe_config: FusedMoEConfig,
68 + quant_config: FusedMoEQuantConfig,
69 + max_num_tokens: int | None = None,
70 + num_dispatchers: int | None = None,
71 + ):
72 + assert quant_config.use_int4_w4a16, "Supports only int4_w4a16"
73 + super().__init__(
74 + moe_config=moe_config,
75 + quant_config=quant_config,
76 + max_num_tokens=max_num_tokens,
77 + num_dispatchers=num_dispatchers,
78 + )
79 + self._empty_topk_weights: torch.Tensor | None = None
80 +
81 + @staticmethod
82 + def activation_format() -> mk.FusedMoEActivationFormat:
83 + return mk.FusedMoEActivationFormat.Standard
84 +
85 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
86 + return TopKWeightAndReduceNoOP()
87 +
88 + @staticmethod
89 + def _supports_current_device() -> bool:
90 + return rdna3_moe_kernel_available()
91 +
92 + @staticmethod
93 + def _supports_no_act_and_mul() -> bool:
94 + return True
95 +
96 + @staticmethod
97 + def _supports_quant_scheme(
98 + weight_key: QuantKey | None,
99 + activation_key: QuantKey | None,
100 + ) -> bool:
101 + # Symmetric int4 weights only: the kernel consumes synthesized zero
102 + # points and has no path for checkpoint zero points.
103 + return activation_key is None and weight_key in (kInt4Static, kInt4Static32)
104 +
105 + @staticmethod
106 + def _supports_activation(activation: MoEActivation) -> bool:
107 + return apply_moe_activation_supported(activation)
108 +
109 + @staticmethod
110 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
111 + return not (
112 + moe_parallel_config.use_fi_nvl_two_sided_kernels
113 + or moe_parallel_config.use_fi_nvl_one_sided_kernels
114 + )
115 +
116 + def moe_problem_size(
117 + self,
118 + a1: torch.Tensor,
119 + w1: torch.Tensor,
120 + w2: torch.Tensor,
121 + topk_ids: torch.Tensor,
122 + ) -> tuple[int, int, int, int, int]:
123 + # Weights are packed K-first: w1 is [E, K // 8, w13_shards * N] and
124 + # w2 is [E, N // 8, K], so neither N nor K can be read off a trailing
125 + # dimension the way the base implementation does.
126 + assert w1.dim() == 3 and w2.dim() == 3
127 + assert a1.dim() == 2
128 + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
129 + return w1.size(0), a1.size(0), w2.size(1) * 8, a1.size(-1), topk_ids.size(1)
130 +
131 + def workspace_shapes(
132 + self,
133 + M: int,
134 + N: int,
135 + K: int,
136 + topk: int,
137 + global_num_experts: int,
138 + local_num_experts: int,
139 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
140 + activation: MoEActivation,
141 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
142 + # The modular kernel provisions the output buffer out of workspace13,
143 + # so both GEMM-1 buffers have to live in workspace2 — the second GEMM
144 + # reads the activation output while accumulating into the output.
145 + gate_up = N if not activation.is_gated else 2 * N
146 + return ((M, K), (M * topk * (gate_up + N),), (M, K))
147 +
148 + def apply(
149 + self,
150 + output: torch.Tensor,
151 + hidden_states: torch.Tensor,
152 + w1: torch.Tensor,
153 + w2: torch.Tensor,
154 + topk_weights: torch.Tensor,
155 + topk_ids: torch.Tensor,
156 + activation: MoEActivation,
157 + global_num_experts: int,
158 + expert_map: torch.Tensor | None,
159 + a1q_scale: torch.Tensor | None,
160 + a2_scale: torch.Tensor | None,
161 + workspace13: torch.Tensor,
162 + workspace2: torch.Tensor,
163 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
164 + apply_router_weight_on_input: bool,
165 + ) -> None:
166 + assert self.w1_scale is not None and self.w2_scale is not None
167 + assert self.w1_zp is not None and self.w2_zp is not None
168 +
169 + E, M, N, K, top_k = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
170 + if global_num_experts == -1:
171 + global_num_experts = E
172 +
173 + gate_up = w1.size(2)
174 + act_n = self.adjust_N_for_activation(gate_up, activation)
175 + rows = M * top_k
176 +
177 + scratch = workspace2.view(-1)
178 + gate_up_out = scratch[: rows * gate_up].view(rows, gate_up)
179 + act_out = scratch[rows * gate_up : rows * (gate_up + act_n)].view(rows, act_n)
180 +
181 + # BLOCK_SIZE_M=1 for decode (no padding waste), 4 for prefill.
182 + block_size_m = 1 if M <= 4 else 4
183 + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
184 + topk_ids,
185 + block_size_m,
186 + global_num_experts,
187 + expert_map,
188 + )
189 +
190 + if (
191 + self._empty_topk_weights is None
192 + or self._empty_topk_weights.device != hidden_states.device
193 + ):
194 + self._empty_topk_weights = torch.empty(
195 + 0, dtype=torch.float32, device=hidden_states.device
196 + )
197 + topk_weights_f32 = topk_weights.reshape(-1).float()
198 + no_topk_weights = self._empty_topk_weights
199 +
200 + gate_up_out.zero_()
201 + ops.moe_gptq_gemm_rdna3(
202 + hidden_states,
203 + gate_up_out,
204 + w1,
205 + self.w1_scale,
206 + self.w1_zp,
207 + topk_weights_f32 if apply_router_weight_on_input else no_topk_weights,
208 + sorted_token_ids,
209 + expert_ids,
210 + num_tokens_post_padded,
211 + top_k,
212 + block_size_m,
213 + apply_router_weight_on_input,
214 + )
215 +
216 + self.activation(activation, act_out, gate_up_out)
217 +
218 + # output_topk=top_k makes the kernel accumulate into out[token_id],
219 + # fusing the top-k reduction into the atomic write-back.
220 + output.zero_()
221 + ops.moe_gptq_gemm_rdna3(
222 + act_out,
223 + output,
224 + w2,
225 + self.w2_scale,
226 + self.w2_zp,
227 + no_topk_weights if apply_router_weight_on_input else topk_weights_f32,
228 + sorted_token_ids,
229 + expert_ids,
230 + num_tokens_post_padded,
231 + 1,
232 + block_size_m,
233 + not apply_router_weight_on_input,
234 + output_topk=top_k,
235 + )
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+110
@@ -7,6 +7,7 @@ from typing import Any
7 import torch
8 from compressed_tensors.quantization import (
9 QuantizationArgs,
10 + QuantizationStrategy,
11 )
12
13 import vllm._custom_ops as ops
@@ -41,8 +42,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
42 )
43 from vllm.model_executor.layers.quantization.utils.quant_utils import (
44 QuantKey,
45 + pack_quantized_values_into_int32,
46 )
47 from vllm.platforms import current_platform
48 +from vllm.scalar_type import scalar_types
49
50 logger = init_logger(__name__)
51
@@ -56,6 +59,7 @@ class WNA16MoEBackend(Enum):
59 TRITON = "TRITON"
60 XPU = "XPU"
61 EMULATION = "EMULATION"
62 + RDNA3 = "RDNA3"
63
64
65 def backend_to_kernel_cls(
@@ -100,6 +104,12 @@ def backend_to_kernel_cls(
104 )
105
106 return [Int4EmulationTritonExperts]
107 + elif backend == WNA16MoEBackend.RDNA3:
108 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
109 + Rdna3WNA16Experts,
110 + )
111 +
112 + return [Rdna3WNA16Experts]
113 else:
114 raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
115
@@ -114,6 +124,8 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
124 return [WNA16MoEBackend.XPU]
125
126 return [
127 + # Native HIP kernel, gated on gfx1100 by _supports_current_device().
128 + WNA16MoEBackend.RDNA3,
129 WNA16MoEBackend.FLASHINFER_TRTLLM,
130 WNA16MoEBackend.MARLIN,
131 WNA16MoEBackend.BATCHED_MARLIN,
@@ -134,6 +146,18 @@ def _backend_incompatibility_reason(
146 if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias):
147 return "zero points and bias are not supported"
148
149 + if backend == WNA16MoEBackend.RDNA3:
150 + if not isinstance(quant_config, QuantizationArgs):
151 + return "only compressed-tensors checkpoints are supported"
152 + if may_have_zp:
153 + return "asymmetric checkpoints are not supported"
154 + if may_have_bias:
155 + return "expert bias is not supported"
156 + if quant_config.actorder == "group":
157 + return "group activation ordering is not supported"
158 + if quant_config.strategy != QuantizationStrategy.GROUP:
159 + return "only group-wise scales are supported"
160 +
161 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
162 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
163 from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
@@ -190,6 +214,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
214 "humming": WNA16MoEBackend.HUMMING,
215 "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM,
216 "emulation": WNA16MoEBackend.EMULATION,
217 + "rdna3": WNA16MoEBackend.RDNA3,
218 }
219 if backend := mapping.get(runner_backend):
220 return backend
@@ -372,6 +397,9 @@ def make_wna16_moe_kernel(
397 from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
398 Int4EmulationTritonExperts,
399 )
400 + from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
401 + Rdna3WNA16Experts,
402 + )
403 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
404 XPUExpertsWNA16,
405 )
@@ -387,6 +415,7 @@ def make_wna16_moe_kernel(
415 XPUExpertsWNA16,
416 CPUExpertsInt4,
417 Int4EmulationTritonExperts,
418 + Rdna3WNA16Experts,
419 )
420 if backend == WNA16MoEBackend.HUMMING:
421 allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING))
@@ -811,6 +840,78 @@ def _process_awq_weights_marlin(
840 )
841
842
843 +def _synthesize_rdna3_qzeros(
844 + groups: int, out_features: int, device: torch.device
845 +) -> torch.Tensor:
846 + """Create the packed zero-point tensor for symmetric quantization.
847 +
848 + GPTQv1 +1 quirk: the kernel adds 1 to the stored zeros, so encode
849 + (bias - 1) = 7 for uint4b8 (bias=8).
850 + """
851 + zeros = torch.full(
852 + (groups, out_features),
853 + scalar_types.uint4b8.bias - 1,
854 + dtype=torch.int32,
855 + device=device,
856 + )
857 + return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
858 +
859 +
860 +def _process_weights_rdna3(
861 + w13: torch.Tensor,
862 + w2: torch.Tensor,
863 + w13_scale: torch.Tensor,
864 + w2_scale: torch.Tensor,
865 + group_size: int,
866 +) -> tuple[
867 + torch.Tensor, # w13_qweight
868 + torch.Tensor, # w2_qweight
869 + torch.Tensor, # w13_scales
870 + torch.Tensor, # w2_scales
871 + torch.Tensor | None, # w13_qzeros
872 + torch.Tensor | None, # w2_qzeros
873 + torch.Tensor | None, # w13_input_global_scale
874 + torch.Tensor | None, # w2_input_global_scale
875 + torch.Tensor | None, # w13_bias
876 + torch.Tensor | None, # w2_bias
877 +]:
878 + """RDNA3 (gfx1100) W4A16 weight post-processing.
879 +
880 + Interleaves the packed nibbles per expert (the exllama shuffle the dense
881 + RDNA3 kernel also uses) and synthesizes the symmetric zero points that
882 + ``moe_gptq_gemm_rdna3`` dequantizes with. The packed layout
883 + ``[E, K // 8, N]`` and the ``[E, groups, N]`` scales are already what the
884 + kernel wants, so neither is repacked.
885 + """
886 + device = w13.device
887 + num_experts = w13.size(0)
888 +
889 + for e in range(num_experts):
890 + w13_e = w13[e].contiguous()
891 + ops.gptq_shuffle(w13_e, 4)
892 + w13[e] = w13_e
893 + w2_e = w2[e].contiguous()
894 + ops.gptq_shuffle(w2_e, 4)
895 + w2[e] = w2_e
896 +
897 + def _qzeros(w: torch.Tensor) -> torch.Tensor:
898 + qz = _synthesize_rdna3_qzeros((w.size(1) * 8) // group_size, w.size(2), device)
899 + return qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous()
900 +
901 + return (
902 + w13,
903 + w2,
904 + w13_scale.contiguous(),
905 + w2_scale.contiguous(),
906 + _qzeros(w13),
907 + _qzeros(w2),
908 + None, # w13_input_global_scale
909 + None, # w2_input_global_scale
910 + None, # w13_bias
911 + None, # w2_bias
912 + )
913 +
914 +
915 def _process_weights_cpu(
916 quant_config: QuantizationConfig | QuantizationArgs | None,
917 w13: torch.Tensor,
@@ -1432,6 +1533,15 @@ def convert_to_wna16_moe_kernel_format(
1533 w13_bias,
1534 w2_bias,
1535 )
1536 + elif backend == WNA16MoEBackend.RDNA3:
1537 + assert isinstance(quant_config, QuantizationArgs)
1538 + return _process_weights_rdna3(
1539 + w13,
1540 + w2,
1541 + w13_scale,
1542 + w2_scale,
1543 + quant_config.group_size,
1544 + )
1545 elif backend == WNA16MoEBackend.CPU:
1546 return _process_weights_cpu(
1547 quant_config,
vllm/model_executor/layers/fused_moe/routed_experts.py
-1
@@ -617,7 +617,6 @@ class RoutedExperts(PluggableLayer):
617 # against known CompressionFormat enum values that have this quality
618 if quant_method_name in (
619 "CompressedTensorsWNA16MoEMethod",
620 - "CompressedTensorsWNA16RDNA3MoEMethod",
620 "CompressedTensorsW4A16FlydslMoEMethod",
621 ):
622 if is_transposed:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
-7
@@ -88,14 +88,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
88 f"and bits: {weight_quant.num_bits}"
89 )
90
91 - # Native ROCm HIP kernels (RDNA3, etc.)
91 if current_platform.is_rocm():
93 - from . import rocm_moe_rdna
94 -
95 - if rocm_moe_rdna.is_supported(weight_quant):
96 - return rocm_moe_rdna.make_method(
97 - weight_quant, input_quant, layer.moe_config
98 - )
92 from vllm.platforms.rocm import on_gfx950
93
94 vllm_config = get_current_vllm_config()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+4 -2
@@ -443,10 +443,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
443 replace_parameter(layer, "w13_weight_scale", w13_scales)
444 replace_parameter(layer, "w2_weight_scale", w2_scales)
445
446 - # CPU fused_experts_cpu requires zero points even for symmetric quant.
446 + # CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points
447 + # even for symmetric quant (the oracle synthesizes them).
448 # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.
449 if (
449 - not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU
450 + not self.symmetric
451 + or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3)
452 ) and self.wna16_backend != WNA16MoEBackend.EMULATION:
453 assert w13_qzeros is not None and w2_qzeros is not None
454 replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py deleted
-260
@@ -1,260 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel.
4 -
5 -Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that
6 -handles expert routing + W4A16 dequant + dot product with atomic output.
7 -
8 -Weight format (per expert, same as dense RDNA3 W4A16):
9 - - Packed int32 ``[E, K/8, N]`` with exllama shuffle
10 - - Scales ``[E, groups, N]`` in activation dtype
11 - - Zero points ``[E, groups, N/8]`` packed int32 (synthesized)
12 -"""
13 -
14 -import torch
15 -
16 -from vllm import _custom_ops as ops
17 -from vllm.logger import init_logger
18 -from vllm.model_executor.layers.fused_moe import (
19 - RoutedExperts,
20 - SharedExperts,
21 -)
22 -from vllm.model_executor.layers.fused_moe.activation import (
23 - MoEActivation,
24 - apply_moe_activation,
25 -)
26 -from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
27 - moe_align_block_size,
28 -)
29 -from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
30 - CompressedTensorsWNA16MoEMethod,
31 -)
32 -from vllm.model_executor.layers.quantization.utils.quant_utils import (
33 - pack_quantized_values_into_int32,
34 -)
35 -from vllm.scalar_type import scalar_types
36 -
37 -logger = init_logger(__name__)
38 -
39 -
40 -def _synthesize_qzeros(
41 - groups: int, out_features: int, device: torch.device
42 -) -> torch.Tensor:
43 - """Create packed zero-point tensor for symmetric quant.
44 -
45 - GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode
46 - (bias - 1) = 7 for uint4b8 (bias=8).
47 - """
48 - zeros = torch.full(
49 - (groups, out_features),
50 - scalar_types.uint4b8.bias - 1,
51 - dtype=torch.int32,
52 - device=device,
53 - )
54 - return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
55 -
56 -
57 -class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod):
58 - """W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3).
59 -
60 - Weights are in RDNA3 format (shuffled int32 [E, K/8, N]),
61 - NOT Triton format (transposed uint8). apply() dispatches through
62 - the fused HIP kernel directly.
63 - """
64 -
65 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
66 - device = layer.w13_weight_packed.device
67 - num_experts = layer.w13_weight_packed.shape[0]
68 -
69 - # Shuffle weights in-place per expert (exllama nibble interleave)
70 - for e in range(num_experts):
71 - w13_e = layer.w13_weight_packed.data[e].contiguous()
72 - ops.gptq_shuffle(w13_e, 4)
73 - layer.w13_weight_packed.data[e] = w13_e
74 - w2_e = layer.w2_weight_packed.data[e].contiguous()
75 - ops.gptq_shuffle(w2_e, 4)
76 - layer.w2_weight_packed.data[e] = w2_e
77 -
78 - # Keep scales as [E, groups, N] in activation dtype
79 - act_dtype = layer.w13_weight_scale.dtype
80 - layer.w13_weight_scale = torch.nn.Parameter(
81 - layer.w13_weight_scale.to(dtype=act_dtype).contiguous(),
82 - requires_grad=False,
83 - )
84 - layer.w2_weight_scale = torch.nn.Parameter(
85 - layer.w2_weight_scale.to(dtype=act_dtype).contiguous(),
86 - requires_grad=False,
87 - )
88 -
89 - # Synthesize packed zero points: [E, groups, N/8] int32
90 - w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size
91 - w13_N = layer.w13_weight_packed.shape[2]
92 - w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size
93 - w2_N = layer.w2_weight_packed.shape[2]
94 -
95 - w13_qz = _synthesize_qzeros(w13_groups, w13_N, device)
96 - w2_qz = _synthesize_qzeros(w2_groups, w2_N, device)
97 - layer.w13_qzeros = torch.nn.Parameter(
98 - w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
99 - requires_grad=False,
100 - )
101 - layer.w2_qzeros = torch.nn.Parameter(
102 - w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
103 - requires_grad=False,
104 - )
105 -
106 - # Pre-allocate reusable buffers for decode (sizes based on top_k=8)
107 - N_gate_up = w13_N
108 - hidden_size = w2_N
109 - intermediate = N_gate_up // 2 # gated activation
110 - # Max tokens we expect in decode; prefill will re-allocate if needed
111 - max_decode_tokens = 16
112 - top_k = 8 # conservative default
113 - buf_size = max_decode_tokens * top_k
114 - layer.rdna3_w1_buf = torch.zeros(
115 - buf_size, N_gate_up, dtype=act_dtype, device=device
116 - )
117 - layer.rdna3_act_buf = torch.empty(
118 - buf_size, intermediate, dtype=act_dtype, device=device
119 - )
120 - layer.rdna3_out_buf = torch.zeros(
121 - max_decode_tokens, hidden_size, dtype=act_dtype, device=device
122 - )
123 - layer.rdna3_empty_tw = torch.empty(0, device=device)
124 -
125 - def apply(
126 - self,
127 - layer: RoutedExperts,
128 - x: torch.Tensor,
129 - topk_weights: torch.Tensor,
130 - topk_ids: torch.Tensor,
131 - shared_experts: SharedExperts | None,
132 - shared_experts_input: torch.Tensor | None,
133 - ) -> torch.Tensor:
134 - activation = (
135 - layer.activation
136 - if isinstance(layer.activation, MoEActivation)
137 - else MoEActivation.from_str(layer.activation)
138 - )
139 - return _rdna3_fused_moe(
140 - x,
141 - topk_weights,
142 - topk_ids,
143 - layer=layer,
144 - activation=activation,
145 - apply_router_weight_on_input=(layer.apply_router_weight_on_input),
146 - global_num_experts=layer.global_num_experts,
147 - expert_map=layer.expert_map,
148 - )
149 -
150 -
151 -def _rdna3_fused_moe(
152 - hidden_states: torch.Tensor,
153 - topk_weights: torch.Tensor,
154 - topk_ids: torch.Tensor,
155 - layer: RoutedExperts,
156 - activation: MoEActivation,
157 - apply_router_weight_on_input: bool,
158 - global_num_experts: int,
159 - expert_map: torch.Tensor | None,
160 -) -> torch.Tensor:
161 - """Fused MoE forward using the RDNA3 W4A16 HIP kernel.
162 -
163 - Optimizations vs naive dispatch:
164 - - BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path)
165 - - Pre-allocated buffers (no torch.zeros per call)
166 - - Inline token sorting for small M (skip moe_align_block_size)
167 - - moe_sum fused into output accumulation
168 - """
169 - num_tokens = hidden_states.shape[0]
170 - top_k = topk_ids.shape[1]
171 - total_tokens = num_tokens * top_k
172 - N_gate_up = layer.w13_weight_packed.shape[2]
173 - hidden_size = layer.w2_weight_packed.shape[2]
174 - dtype = hidden_states.dtype
175 - device = hidden_states.device
176 -
177 - intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up
178 -
179 - if global_num_experts <= 0:
180 - global_num_experts = layer.w13_weight_packed.shape[0]
181 -
182 - # BLOCK_SIZE_M=1 for decode (small M), 4 for prefill
183 - block_size_m = 1 if num_tokens <= 4 else 4
184 -
185 - # --- Token routing ---
186 - sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
187 - topk_ids,
188 - block_size_m,
189 - global_num_experts,
190 - expert_map,
191 - )
192 -
193 - # --- Reuse pre-allocated buffers when possible ---
194 - if total_tokens <= layer.rdna3_w1_buf.shape[0]:
195 - w1_out = layer.rdna3_w1_buf[:total_tokens]
196 - w1_out.zero_()
197 - act_out = layer.rdna3_act_buf[:total_tokens]
198 - else:
199 - w1_out = torch.zeros(
200 - total_tokens,
201 - N_gate_up,
202 - dtype=dtype,
203 - device=device,
204 - )
205 - act_out = torch.empty(
206 - total_tokens,
207 - intermediate_size,
208 - dtype=dtype,
209 - device=device,
210 - )
211 -
212 - # --- topk weights (pre-cast to float32 for kernel) ---
213 - topk_w_float = topk_weights.view(-1).float()
214 - empty_tw = layer.rdna3_empty_tw
215 -
216 - # --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] ---
217 - ops.moe_gptq_gemm_rdna3(
218 - hidden_states,
219 - w1_out,
220 - layer.w13_weight_packed,
221 - layer.w13_weight_scale,
222 - layer.w13_qzeros,
223 - topk_w_float if apply_router_weight_on_input else empty_tw,
224 - sorted_token_ids,
225 - expert_ids,
226 - num_tokens_post_padded,
227 - top_k,
228 - block_size_m,
229 - apply_router_weight_on_input,
230 - )
231 -
232 - # --- Activation (silu_and_mul etc.) ---
233 - apply_moe_activation(activation, act_out, w1_out)
234 -
235 - # --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) ---
236 - # output_topk=top_k: kernel writes to out[token_id / top_k] directly,
237 - # fusing moe_sum into the atomic accumulation — saves one kernel launch
238 - # and the w2_out intermediate buffer.
239 - out = torch.zeros(
240 - num_tokens,
241 - hidden_size,
242 - dtype=dtype,
243 - device=device,
244 - )
245 - ops.moe_gptq_gemm_rdna3(
246 - act_out,
247 - out,
248 - layer.w2_weight_packed,
249 - layer.w2_weight_scale,
250 - layer.w2_qzeros,
251 - topk_w_float if not apply_router_weight_on_input else empty_tw,
252 - sorted_token_ids,
253 - expert_ids,
254 - num_tokens_post_padded,
255 - 1,
256 - block_size_m,
257 - not apply_router_weight_on_input,
258 - output_topk=top_k,
259 - )
260 - return out
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py deleted
-48
@@ -1,48 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""ROCm MoE kernel dispatcher.
4 -
5 -Selects architecture-specific native HIP MoE kernels in priority order.
6 -Falls back to the Triton WNA16 path when no native kernel is available.
7 -"""
8 -
9 -import torch
10 -
11 -from vllm.logger import init_logger
12 -
13 -logger = init_logger(__name__)
14 -
15 -
16 -def is_supported(weight_quant) -> bool:
17 - """Check if a native ROCm MoE kernel is available for this config."""
18 - if weight_quant.num_bits != 4:
19 - return False
20 -
21 - from vllm.platforms.rocm import on_gfx1100
22 -
23 - # RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc.
24 - return (
25 - on_gfx1100()
26 - and hasattr(torch.ops, "_rocm_C")
27 - and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
28 - )
29 -
30 -
31 -def make_method(weight_quant, input_quant, moe_config):
32 - """Create the native ROCm MoE method. Call only after is_supported()."""
33 - from vllm.platforms.rocm import on_gfx1100
34 -
35 - if on_gfx1100():
36 - from .compressed_tensors_moe_wna16_rdna3 import (
37 - CompressedTensorsWNA16RDNA3MoEMethod,
38 - )
39 -
40 - logger.info_once(
41 - "Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)"
42 - )
43 - return CompressedTensorsWNA16RDNA3MoEMethod(
44 - weight_quant, input_quant, moe_config
45 - )
46 -
47 - # Future: RDNA4, CDNA, etc.
48 - raise RuntimeError("is_supported() returned True but no kernel matched")