[Refactor][ROCm] Migrate the RDNA3 W4A16 MoE to the oracle/experts pa… (#55522)
Signed-off-by: JartX <sagformas@epdcenter.es>
This commit is contained in:
JartX
@@ -34,6 +34,13 @@ from vllm.platforms import current_platform
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
|
||||
kInt4Static,
|
||||
kInt4Static32,
|
||||
kInt4Static32Asym,
|
||||
kInt4StaticAsym,
|
||||
kInt8Static,
|
||||
)
|
||||
from vllm.platforms.rocm import on_gfx1100 # noqa: E402
|
||||
|
||||
gfx1100_only = pytest.mark.skipif(
|
||||
@@ -166,15 +173,13 @@ def test_op_absent_on_non_gfx1100(op_name):
|
||||
|
||||
@not_gfx1100
|
||||
def test_rocm_moe_not_supported_on_non_gfx1100():
|
||||
"""rocm_moe_rdna.is_supported() must return False on non-gfx1100 hardware."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
"""The RDNA3 MoE experts must not be selectable on non-gfx1100 hardware."""
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
|
||||
wq = type("WQ", (), {"num_bits": 4})()
|
||||
assert rocm_moe_rdna.is_supported(wq) is False, (
|
||||
"rocm_moe_rdna.is_supported() returned True on non-gfx1100 — "
|
||||
"dispatch guard is broken"
|
||||
assert Rdna3WNA16Experts._supports_current_device() is False, (
|
||||
"Rdna3WNA16Experts reported support on non-gfx1100 — dispatch guard is broken"
|
||||
)
|
||||
|
||||
|
||||
@@ -359,54 +364,58 @@ class TestCustomOpsGuards:
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class _FakeWeightQuant:
|
||||
"""Minimal stand-in for a weight quantization config."""
|
||||
|
||||
def __init__(self, num_bits):
|
||||
self.num_bits = num_bits
|
||||
|
||||
|
||||
class TestMoEDispatchMocked:
|
||||
"""Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
|
||||
|
||||
def test_is_supported_false_when_mocked_cdna(self):
|
||||
"""rocm_moe_rdna.is_supported() must return False when not on gfx1100."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
def test_kernel_unavailable_when_mocked_cdna(self):
|
||||
"""The device gate must reject when not on gfx1100."""
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
rdna3_moe_kernel_available,
|
||||
)
|
||||
|
||||
with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
|
||||
assert rdna3_moe_kernel_available() is False
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [2, 3, 8, 16])
|
||||
def test_is_supported_rejects_non_w4(self, num_bits):
|
||||
"""is_supported() rejects non-4-bit even before checking arch."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
@pytest.mark.parametrize(
|
||||
"weight_key",
|
||||
[kInt8Static, kInt4StaticAsym, kInt4Static32Asym, None],
|
||||
)
|
||||
def test_quant_scheme_rejects_non_symmetric_int4(self, weight_key):
|
||||
"""Only symmetric int4 weight-only schemes reach the kernel."""
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
|
||||
assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is False
|
||||
|
||||
def test_is_supported_false_when_op_missing(self):
|
||||
"""is_supported() returns False when the C++ op doesn't exist."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
@pytest.mark.parametrize("weight_key", [kInt4Static, kInt4Static32])
|
||||
def test_quant_scheme_accepts_symmetric_int4(self, weight_key):
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
|
||||
assert Rdna3WNA16Experts._supports_quant_scheme(weight_key, None) is True
|
||||
|
||||
def test_kernel_unavailable_when_op_missing(self):
|
||||
"""The device gate returns False when the C++ op doesn't exist."""
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
rdna3_moe_kernel_available,
|
||||
)
|
||||
|
||||
fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
|
||||
with patch.object(torch, "ops", create=True) as mock_ops:
|
||||
mock_ops._rocm_C = fake_rocm_c
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
|
||||
assert rdna3_moe_kernel_available() is False
|
||||
|
||||
def test_is_supported_false_when_rocm_c_absent(self):
|
||||
"""is_supported() returns False when _rocm_C doesn't exist at all."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
def test_kernel_unavailable_when_rocm_c_absent(self):
|
||||
"""The device gate returns False when _rocm_C doesn't exist at all."""
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
rdna3_moe_kernel_available,
|
||||
)
|
||||
|
||||
fake_ops = type("FakeOps", (), {})()
|
||||
with patch.object(torch, "ops", fake_ops):
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
|
||||
assert rdna3_moe_kernel_available() is False
|
||||
|
||||
|
||||
class TestDenseKernelSelectionMocked:
|
||||
@@ -471,30 +480,30 @@ class TestDenseKernelSelectionMocked:
|
||||
)
|
||||
|
||||
|
||||
class TestCompressedTensorsMoEDispatchGuard:
|
||||
"""Verify compressed_tensors_moe.py only enters rocm_moe_rdna under is_rocm()."""
|
||||
class TestWNA16OracleWiring:
|
||||
"""The RDNA3 backend must reach the kernel only through the oracle."""
|
||||
|
||||
def test_rocm_guard_in_dispatch_source(self):
|
||||
"""The rocm_moe_rdna import and call must be inside an is_rocm() check."""
|
||||
src = _read_pkg_source_or_skip(
|
||||
"model_executor",
|
||||
"layers",
|
||||
"quantization",
|
||||
"compressed_tensors",
|
||||
"compressed_tensors_moe",
|
||||
"compressed_tensors_moe.py",
|
||||
def test_backend_maps_to_rdna3_experts(self):
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
WNA16MoEBackend,
|
||||
backend_to_kernel_cls,
|
||||
map_wna16_backend,
|
||||
)
|
||||
lines = src.splitlines()
|
||||
|
||||
for i, line in enumerate(lines, 1):
|
||||
stripped = line.strip()
|
||||
if "rocm_moe" in stripped and not stripped.startswith("#"):
|
||||
found_guard = False
|
||||
for j in range(i - 1, max(0, i - 15), -1):
|
||||
if "is_rocm()" in lines[j - 1]:
|
||||
found_guard = True
|
||||
break
|
||||
assert found_guard, (
|
||||
f"L{i}: rocm_moe_rdna reference not protected by "
|
||||
f"is_rocm() guard: {stripped}"
|
||||
)
|
||||
assert backend_to_kernel_cls(WNA16MoEBackend.RDNA3) == [Rdna3WNA16Experts]
|
||||
assert map_wna16_backend("rdna3") == WNA16MoEBackend.RDNA3
|
||||
|
||||
def test_backend_is_offered_before_the_triton_fallback(self):
|
||||
"""Priority order: the native kernel outranks Triton when supported."""
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
WNA16MoEBackend,
|
||||
_get_priority_backends,
|
||||
)
|
||||
|
||||
backends = _get_priority_backends()
|
||||
assert backends.index(WNA16MoEBackend.RDNA3) < backends.index(
|
||||
WNA16MoEBackend.TRITON
|
||||
)
|
||||
|
||||
@@ -93,6 +93,40 @@ def test_moe_wna16_accepts_channelwise_gptq_activation_order():
|
||||
False,
|
||||
"MoeWNA16 checkpoint layout",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.RDNA3,
|
||||
AutoGPTQConfig(4, 128, False, True, False, {}, {}),
|
||||
False,
|
||||
False,
|
||||
"compressed-tensors",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.RDNA3,
|
||||
QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.GROUP,
|
||||
symmetric=False,
|
||||
dynamic=False,
|
||||
group_size=128,
|
||||
),
|
||||
True,
|
||||
False,
|
||||
"asymmetric",
|
||||
),
|
||||
(
|
||||
WNA16MoEBackend.RDNA3,
|
||||
QuantizationArgs(
|
||||
num_bits=4,
|
||||
type=QuantizationType.INT,
|
||||
strategy=QuantizationStrategy.CHANNEL,
|
||||
symmetric=True,
|
||||
dynamic=False,
|
||||
),
|
||||
False,
|
||||
False,
|
||||
"group-wise scales",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_wna16_oracle_rejects_incompatible_quant_structures(
|
||||
|
||||
@@ -141,6 +141,7 @@ MoEBackend = Literal[
|
||||
"flydsl",
|
||||
"hpc",
|
||||
"emulation",
|
||||
"rdna3",
|
||||
]
|
||||
|
||||
# Backends that run the mega-MoE model path through the flashinfer moe_ep
|
||||
@@ -262,6 +263,7 @@ class KernelConfig:
|
||||
- "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16
|
||||
(moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250)
|
||||
- "flydsl": Use AMD FlyDSL kernels (ROCm only)
|
||||
- "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only)
|
||||
- "hpc": Use HPC kernels (FP8 and Hopper only)
|
||||
- "emulation": use BF16/FP16 GEMM, dequantizing weights and
|
||||
running QDQ on activations.
|
||||
|
||||
@@ -0,0 +1,235 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Fused MoE W4A16 experts on the RDNA3 (gfx1100) HIP kernel.
|
||||
|
||||
``moe_gptq_gemm_rdna3`` is a single HIP kernel launch per GEMM that handles
|
||||
expert routing + W4A16 dequant + dot product with atomic output accumulation.
|
||||
|
||||
Weight format (per expert, same as the dense RDNA3 W4A16 kernel):
|
||||
- Packed int32 ``[E, K/8, N]`` with exllama shuffle
|
||||
- Scales ``[E, groups, N]`` in activation dtype
|
||||
- Zero points ``[E, groups, N/8]`` packed int32 (synthesized, symmetric only)
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.activation import (
|
||||
MoEActivation,
|
||||
apply_moe_activation_supported,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
|
||||
moe_align_block_size,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
|
||||
TopKWeightAndReduceNoOP,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
kInt4Static,
|
||||
kInt4Static32,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def rdna3_moe_kernel_available() -> bool:
|
||||
"""Whether the fused RDNA3 MoE HIP kernel is built into this binary."""
|
||||
if not current_platform.is_rocm():
|
||||
return False
|
||||
|
||||
from vllm.platforms.rocm import on_gfx1100
|
||||
|
||||
return (
|
||||
on_gfx1100()
|
||||
and hasattr(torch.ops, "_rocm_C")
|
||||
and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
|
||||
)
|
||||
|
||||
|
||||
class Rdna3WNA16Experts(mk.FusedMoEExpertsModular):
|
||||
"""W4A16 experts backed by ``moe_gptq_gemm_rdna3`` (gfx1100).
|
||||
|
||||
Both GEMMs accumulate atomically, so their destinations are zeroed first.
|
||||
The second GEMM is given ``output_topk``, which makes it reduce over the
|
||||
top-k dimension while it accumulates — the ``moe_sum`` launch and the
|
||||
``[M * top_k, K]`` intermediate are fused away, hence
|
||||
``TopKWeightAndReduceNoOP``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
moe_config: FusedMoEConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
max_num_tokens: int | None = None,
|
||||
num_dispatchers: int | None = None,
|
||||
):
|
||||
assert quant_config.use_int4_w4a16, "Supports only int4_w4a16"
|
||||
super().__init__(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
max_num_tokens=max_num_tokens,
|
||||
num_dispatchers=num_dispatchers,
|
||||
)
|
||||
self._empty_topk_weights: torch.Tensor | None = None
|
||||
|
||||
@staticmethod
|
||||
def activation_format() -> mk.FusedMoEActivationFormat:
|
||||
return mk.FusedMoEActivationFormat.Standard
|
||||
|
||||
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
|
||||
return TopKWeightAndReduceNoOP()
|
||||
|
||||
@staticmethod
|
||||
def _supports_current_device() -> bool:
|
||||
return rdna3_moe_kernel_available()
|
||||
|
||||
@staticmethod
|
||||
def _supports_no_act_and_mul() -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _supports_quant_scheme(
|
||||
weight_key: QuantKey | None,
|
||||
activation_key: QuantKey | None,
|
||||
) -> bool:
|
||||
# Symmetric int4 weights only: the kernel consumes synthesized zero
|
||||
# points and has no path for checkpoint zero points.
|
||||
return activation_key is None and weight_key in (kInt4Static, kInt4Static32)
|
||||
|
||||
@staticmethod
|
||||
def _supports_activation(activation: MoEActivation) -> bool:
|
||||
return apply_moe_activation_supported(activation)
|
||||
|
||||
@staticmethod
|
||||
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
|
||||
return not (
|
||||
moe_parallel_config.use_fi_nvl_two_sided_kernels
|
||||
or moe_parallel_config.use_fi_nvl_one_sided_kernels
|
||||
)
|
||||
|
||||
def moe_problem_size(
|
||||
self,
|
||||
a1: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
) -> tuple[int, int, int, int, int]:
|
||||
# Weights are packed K-first: w1 is [E, K // 8, w13_shards * N] and
|
||||
# w2 is [E, N // 8, K], so neither N nor K can be read off a trailing
|
||||
# dimension the way the base implementation does.
|
||||
assert w1.dim() == 3 and w2.dim() == 3
|
||||
assert a1.dim() == 2
|
||||
assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"
|
||||
return w1.size(0), a1.size(0), w2.size(1) * 8, a1.size(-1), topk_ids.size(1)
|
||||
|
||||
def workspace_shapes(
|
||||
self,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
topk: int,
|
||||
global_num_experts: int,
|
||||
local_num_experts: int,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
activation: MoEActivation,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
# The modular kernel provisions the output buffer out of workspace13,
|
||||
# so both GEMM-1 buffers have to live in workspace2 — the second GEMM
|
||||
# reads the activation output while accumulating into the output.
|
||||
gate_up = N if not activation.is_gated else 2 * N
|
||||
return ((M, K), (M * topk * (gate_up + N),), (M, K))
|
||||
|
||||
def apply(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
a1q_scale: torch.Tensor | None,
|
||||
a2_scale: torch.Tensor | None,
|
||||
workspace13: torch.Tensor,
|
||||
workspace2: torch.Tensor,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
apply_router_weight_on_input: bool,
|
||||
) -> None:
|
||||
assert self.w1_scale is not None and self.w2_scale is not None
|
||||
assert self.w1_zp is not None and self.w2_zp is not None
|
||||
|
||||
E, M, N, K, top_k = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
|
||||
if global_num_experts == -1:
|
||||
global_num_experts = E
|
||||
|
||||
gate_up = w1.size(2)
|
||||
act_n = self.adjust_N_for_activation(gate_up, activation)
|
||||
rows = M * top_k
|
||||
|
||||
scratch = workspace2.view(-1)
|
||||
gate_up_out = scratch[: rows * gate_up].view(rows, gate_up)
|
||||
act_out = scratch[rows * gate_up : rows * (gate_up + act_n)].view(rows, act_n)
|
||||
|
||||
# BLOCK_SIZE_M=1 for decode (no padding waste), 4 for prefill.
|
||||
block_size_m = 1 if M <= 4 else 4
|
||||
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
|
||||
topk_ids,
|
||||
block_size_m,
|
||||
global_num_experts,
|
||||
expert_map,
|
||||
)
|
||||
|
||||
if (
|
||||
self._empty_topk_weights is None
|
||||
or self._empty_topk_weights.device != hidden_states.device
|
||||
):
|
||||
self._empty_topk_weights = torch.empty(
|
||||
0, dtype=torch.float32, device=hidden_states.device
|
||||
)
|
||||
topk_weights_f32 = topk_weights.reshape(-1).float()
|
||||
no_topk_weights = self._empty_topk_weights
|
||||
|
||||
gate_up_out.zero_()
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
hidden_states,
|
||||
gate_up_out,
|
||||
w1,
|
||||
self.w1_scale,
|
||||
self.w1_zp,
|
||||
topk_weights_f32 if apply_router_weight_on_input else no_topk_weights,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
top_k,
|
||||
block_size_m,
|
||||
apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
self.activation(activation, act_out, gate_up_out)
|
||||
|
||||
# output_topk=top_k makes the kernel accumulate into out[token_id],
|
||||
# fusing the top-k reduction into the atomic write-back.
|
||||
output.zero_()
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
act_out,
|
||||
output,
|
||||
w2,
|
||||
self.w2_scale,
|
||||
self.w2_zp,
|
||||
no_topk_weights if apply_router_weight_on_input else topk_weights_f32,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
1,
|
||||
block_size_m,
|
||||
not apply_router_weight_on_input,
|
||||
output_topk=top_k,
|
||||
)
|
||||
@@ -7,6 +7,7 @@ from typing import Any
|
||||
import torch
|
||||
from compressed_tensors.quantization import (
|
||||
QuantizationArgs,
|
||||
QuantizationStrategy,
|
||||
)
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
@@ -41,8 +42,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
pack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -56,6 +59,7 @@ class WNA16MoEBackend(Enum):
|
||||
TRITON = "TRITON"
|
||||
XPU = "XPU"
|
||||
EMULATION = "EMULATION"
|
||||
RDNA3 = "RDNA3"
|
||||
|
||||
|
||||
def backend_to_kernel_cls(
|
||||
@@ -100,6 +104,12 @@ def backend_to_kernel_cls(
|
||||
)
|
||||
|
||||
return [Int4EmulationTritonExperts]
|
||||
elif backend == WNA16MoEBackend.RDNA3:
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
|
||||
return [Rdna3WNA16Experts]
|
||||
else:
|
||||
raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
|
||||
|
||||
@@ -114,6 +124,8 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
|
||||
return [WNA16MoEBackend.XPU]
|
||||
|
||||
return [
|
||||
# Native HIP kernel, gated on gfx1100 by _supports_current_device().
|
||||
WNA16MoEBackend.RDNA3,
|
||||
WNA16MoEBackend.FLASHINFER_TRTLLM,
|
||||
WNA16MoEBackend.MARLIN,
|
||||
WNA16MoEBackend.BATCHED_MARLIN,
|
||||
@@ -134,6 +146,18 @@ def _backend_incompatibility_reason(
|
||||
if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias):
|
||||
return "zero points and bias are not supported"
|
||||
|
||||
if backend == WNA16MoEBackend.RDNA3:
|
||||
if not isinstance(quant_config, QuantizationArgs):
|
||||
return "only compressed-tensors checkpoints are supported"
|
||||
if may_have_zp:
|
||||
return "asymmetric checkpoints are not supported"
|
||||
if may_have_bias:
|
||||
return "expert bias is not supported"
|
||||
if quant_config.actorder == "group":
|
||||
return "group activation ordering is not supported"
|
||||
if quant_config.strategy != QuantizationStrategy.GROUP:
|
||||
return "only group-wise scales are supported"
|
||||
|
||||
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
|
||||
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
|
||||
from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
|
||||
@@ -190,6 +214,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
|
||||
"humming": WNA16MoEBackend.HUMMING,
|
||||
"flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM,
|
||||
"emulation": WNA16MoEBackend.EMULATION,
|
||||
"rdna3": WNA16MoEBackend.RDNA3,
|
||||
}
|
||||
if backend := mapping.get(runner_backend):
|
||||
return backend
|
||||
@@ -372,6 +397,9 @@ def make_wna16_moe_kernel(
|
||||
from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
|
||||
Int4EmulationTritonExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
|
||||
XPUExpertsWNA16,
|
||||
)
|
||||
@@ -387,6 +415,7 @@ def make_wna16_moe_kernel(
|
||||
XPUExpertsWNA16,
|
||||
CPUExpertsInt4,
|
||||
Int4EmulationTritonExperts,
|
||||
Rdna3WNA16Experts,
|
||||
)
|
||||
if backend == WNA16MoEBackend.HUMMING:
|
||||
allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING))
|
||||
@@ -811,6 +840,78 @@ def _process_awq_weights_marlin(
|
||||
)
|
||||
|
||||
|
||||
def _synthesize_rdna3_qzeros(
|
||||
groups: int, out_features: int, device: torch.device
|
||||
) -> torch.Tensor:
|
||||
"""Create the packed zero-point tensor for symmetric quantization.
|
||||
|
||||
GPTQv1 +1 quirk: the kernel adds 1 to the stored zeros, so encode
|
||||
(bias - 1) = 7 for uint4b8 (bias=8).
|
||||
"""
|
||||
zeros = torch.full(
|
||||
(groups, out_features),
|
||||
scalar_types.uint4b8.bias - 1,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
|
||||
|
||||
|
||||
def _process_weights_rdna3(
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
w13_scale: torch.Tensor,
|
||||
w2_scale: torch.Tensor,
|
||||
group_size: int,
|
||||
) -> tuple[
|
||||
torch.Tensor, # w13_qweight
|
||||
torch.Tensor, # w2_qweight
|
||||
torch.Tensor, # w13_scales
|
||||
torch.Tensor, # w2_scales
|
||||
torch.Tensor | None, # w13_qzeros
|
||||
torch.Tensor | None, # w2_qzeros
|
||||
torch.Tensor | None, # w13_input_global_scale
|
||||
torch.Tensor | None, # w2_input_global_scale
|
||||
torch.Tensor | None, # w13_bias
|
||||
torch.Tensor | None, # w2_bias
|
||||
]:
|
||||
"""RDNA3 (gfx1100) W4A16 weight post-processing.
|
||||
|
||||
Interleaves the packed nibbles per expert (the exllama shuffle the dense
|
||||
RDNA3 kernel also uses) and synthesizes the symmetric zero points that
|
||||
``moe_gptq_gemm_rdna3`` dequantizes with. The packed layout
|
||||
``[E, K // 8, N]`` and the ``[E, groups, N]`` scales are already what the
|
||||
kernel wants, so neither is repacked.
|
||||
"""
|
||||
device = w13.device
|
||||
num_experts = w13.size(0)
|
||||
|
||||
for e in range(num_experts):
|
||||
w13_e = w13[e].contiguous()
|
||||
ops.gptq_shuffle(w13_e, 4)
|
||||
w13[e] = w13_e
|
||||
w2_e = w2[e].contiguous()
|
||||
ops.gptq_shuffle(w2_e, 4)
|
||||
w2[e] = w2_e
|
||||
|
||||
def _qzeros(w: torch.Tensor) -> torch.Tensor:
|
||||
qz = _synthesize_rdna3_qzeros((w.size(1) * 8) // group_size, w.size(2), device)
|
||||
return qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous()
|
||||
|
||||
return (
|
||||
w13,
|
||||
w2,
|
||||
w13_scale.contiguous(),
|
||||
w2_scale.contiguous(),
|
||||
_qzeros(w13),
|
||||
_qzeros(w2),
|
||||
None, # w13_input_global_scale
|
||||
None, # w2_input_global_scale
|
||||
None, # w13_bias
|
||||
None, # w2_bias
|
||||
)
|
||||
|
||||
|
||||
def _process_weights_cpu(
|
||||
quant_config: QuantizationConfig | QuantizationArgs | None,
|
||||
w13: torch.Tensor,
|
||||
@@ -1432,6 +1533,15 @@ def convert_to_wna16_moe_kernel_format(
|
||||
w13_bias,
|
||||
w2_bias,
|
||||
)
|
||||
elif backend == WNA16MoEBackend.RDNA3:
|
||||
assert isinstance(quant_config, QuantizationArgs)
|
||||
return _process_weights_rdna3(
|
||||
w13,
|
||||
w2,
|
||||
w13_scale,
|
||||
w2_scale,
|
||||
quant_config.group_size,
|
||||
)
|
||||
elif backend == WNA16MoEBackend.CPU:
|
||||
return _process_weights_cpu(
|
||||
quant_config,
|
||||
|
||||
@@ -617,7 +617,6 @@ class RoutedExperts(PluggableLayer):
|
||||
# against known CompressionFormat enum values that have this quality
|
||||
if quant_method_name in (
|
||||
"CompressedTensorsWNA16MoEMethod",
|
||||
"CompressedTensorsWNA16RDNA3MoEMethod",
|
||||
"CompressedTensorsW4A16FlydslMoEMethod",
|
||||
):
|
||||
if is_transposed:
|
||||
|
||||
-7
@@ -88,14 +88,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
|
||||
f"and bits: {weight_quant.num_bits}"
|
||||
)
|
||||
|
||||
# Native ROCm HIP kernels (RDNA3, etc.)
|
||||
if current_platform.is_rocm():
|
||||
from . import rocm_moe_rdna
|
||||
|
||||
if rocm_moe_rdna.is_supported(weight_quant):
|
||||
return rocm_moe_rdna.make_method(
|
||||
weight_quant, input_quant, layer.moe_config
|
||||
)
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
|
||||
vllm_config = get_current_vllm_config()
|
||||
|
||||
+4
-2
@@ -443,10 +443,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
|
||||
replace_parameter(layer, "w13_weight_scale", w13_scales)
|
||||
replace_parameter(layer, "w2_weight_scale", w2_scales)
|
||||
|
||||
# CPU fused_experts_cpu requires zero points even for symmetric quant.
|
||||
# CPU fused_experts_cpu and the RDNA3 HIP kernel require zero points
|
||||
# even for symmetric quant (the oracle synthesizes them).
|
||||
# EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.
|
||||
if (
|
||||
not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU
|
||||
not self.symmetric
|
||||
or self.wna16_backend in (WNA16MoEBackend.CPU, WNA16MoEBackend.RDNA3)
|
||||
) and self.wna16_backend != WNA16MoEBackend.EMULATION:
|
||||
assert w13_qzeros is not None and w2_qzeros is not None
|
||||
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
|
||||
|
||||
-260
@@ -1,260 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""CompressedTensors MoE W4A16 using the fused RDNA3 (gfx1100) HIP kernel.
|
||||
|
||||
Uses ``moe_gptq_gemm_rdna3`` — a single HIP kernel launch per GEMM that
|
||||
handles expert routing + W4A16 dequant + dot product with atomic output.
|
||||
|
||||
Weight format (per expert, same as dense RDNA3 W4A16):
|
||||
- Packed int32 ``[E, K/8, N]`` with exllama shuffle
|
||||
- Scales ``[E, groups, N]`` in activation dtype
|
||||
- Zero points ``[E, groups, N/8]`` packed int32 (synthesized)
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fused_moe import (
|
||||
RoutedExperts,
|
||||
SharedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.activation import (
|
||||
MoEActivation,
|
||||
apply_moe_activation,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
|
||||
moe_align_block_size,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16 import ( # noqa: E501
|
||||
CompressedTensorsWNA16MoEMethod,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _synthesize_qzeros(
|
||||
groups: int, out_features: int, device: torch.device
|
||||
) -> torch.Tensor:
|
||||
"""Create packed zero-point tensor for symmetric quant.
|
||||
|
||||
GPTQv1 +1 quirk: kernel adds 1 to stored zeros, so encode
|
||||
(bias - 1) = 7 for uint4b8 (bias=8).
|
||||
"""
|
||||
zeros = torch.full(
|
||||
(groups, out_features),
|
||||
scalar_types.uint4b8.bias - 1,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)
|
||||
|
||||
|
||||
class CompressedTensorsWNA16RDNA3MoEMethod(CompressedTensorsWNA16MoEMethod):
|
||||
"""W4A16 MoE using the fused RDNA3 HIP kernel (moe_gptq_gemm_rdna3).
|
||||
|
||||
Weights are in RDNA3 format (shuffled int32 [E, K/8, N]),
|
||||
NOT Triton format (transposed uint8). apply() dispatches through
|
||||
the fused HIP kernel directly.
|
||||
"""
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
device = layer.w13_weight_packed.device
|
||||
num_experts = layer.w13_weight_packed.shape[0]
|
||||
|
||||
# Shuffle weights in-place per expert (exllama nibble interleave)
|
||||
for e in range(num_experts):
|
||||
w13_e = layer.w13_weight_packed.data[e].contiguous()
|
||||
ops.gptq_shuffle(w13_e, 4)
|
||||
layer.w13_weight_packed.data[e] = w13_e
|
||||
w2_e = layer.w2_weight_packed.data[e].contiguous()
|
||||
ops.gptq_shuffle(w2_e, 4)
|
||||
layer.w2_weight_packed.data[e] = w2_e
|
||||
|
||||
# Keep scales as [E, groups, N] in activation dtype
|
||||
act_dtype = layer.w13_weight_scale.dtype
|
||||
layer.w13_weight_scale = torch.nn.Parameter(
|
||||
layer.w13_weight_scale.to(dtype=act_dtype).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_weight_scale = torch.nn.Parameter(
|
||||
layer.w2_weight_scale.to(dtype=act_dtype).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# Synthesize packed zero points: [E, groups, N/8] int32
|
||||
w13_groups = (layer.w13_weight_packed.shape[1] * 8) // self.group_size
|
||||
w13_N = layer.w13_weight_packed.shape[2]
|
||||
w2_groups = (layer.w2_weight_packed.shape[1] * 8) // self.group_size
|
||||
w2_N = layer.w2_weight_packed.shape[2]
|
||||
|
||||
w13_qz = _synthesize_qzeros(w13_groups, w13_N, device)
|
||||
w2_qz = _synthesize_qzeros(w2_groups, w2_N, device)
|
||||
layer.w13_qzeros = torch.nn.Parameter(
|
||||
w13_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_qzeros = torch.nn.Parameter(
|
||||
w2_qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous(),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
# Pre-allocate reusable buffers for decode (sizes based on top_k=8)
|
||||
N_gate_up = w13_N
|
||||
hidden_size = w2_N
|
||||
intermediate = N_gate_up // 2 # gated activation
|
||||
# Max tokens we expect in decode; prefill will re-allocate if needed
|
||||
max_decode_tokens = 16
|
||||
top_k = 8 # conservative default
|
||||
buf_size = max_decode_tokens * top_k
|
||||
layer.rdna3_w1_buf = torch.zeros(
|
||||
buf_size, N_gate_up, dtype=act_dtype, device=device
|
||||
)
|
||||
layer.rdna3_act_buf = torch.empty(
|
||||
buf_size, intermediate, dtype=act_dtype, device=device
|
||||
)
|
||||
layer.rdna3_out_buf = torch.zeros(
|
||||
max_decode_tokens, hidden_size, dtype=act_dtype, device=device
|
||||
)
|
||||
layer.rdna3_empty_tw = torch.empty(0, device=device)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: RoutedExperts,
|
||||
x: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
shared_experts: SharedExperts | None,
|
||||
shared_experts_input: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
activation = (
|
||||
layer.activation
|
||||
if isinstance(layer.activation, MoEActivation)
|
||||
else MoEActivation.from_str(layer.activation)
|
||||
)
|
||||
return _rdna3_fused_moe(
|
||||
x,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
layer=layer,
|
||||
activation=activation,
|
||||
apply_router_weight_on_input=(layer.apply_router_weight_on_input),
|
||||
global_num_experts=layer.global_num_experts,
|
||||
expert_map=layer.expert_map,
|
||||
)
|
||||
|
||||
|
||||
def _rdna3_fused_moe(
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
layer: RoutedExperts,
|
||||
activation: MoEActivation,
|
||||
apply_router_weight_on_input: bool,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""Fused MoE forward using the RDNA3 W4A16 HIP kernel.
|
||||
|
||||
Optimizations vs naive dispatch:
|
||||
- BLOCK_SIZE_M=1 for decode (no padding waste, bf16 fast path)
|
||||
- Pre-allocated buffers (no torch.zeros per call)
|
||||
- Inline token sorting for small M (skip moe_align_block_size)
|
||||
- moe_sum fused into output accumulation
|
||||
"""
|
||||
num_tokens = hidden_states.shape[0]
|
||||
top_k = topk_ids.shape[1]
|
||||
total_tokens = num_tokens * top_k
|
||||
N_gate_up = layer.w13_weight_packed.shape[2]
|
||||
hidden_size = layer.w2_weight_packed.shape[2]
|
||||
dtype = hidden_states.dtype
|
||||
device = hidden_states.device
|
||||
|
||||
intermediate_size = N_gate_up // 2 if activation.is_gated else N_gate_up
|
||||
|
||||
if global_num_experts <= 0:
|
||||
global_num_experts = layer.w13_weight_packed.shape[0]
|
||||
|
||||
# BLOCK_SIZE_M=1 for decode (small M), 4 for prefill
|
||||
block_size_m = 1 if num_tokens <= 4 else 4
|
||||
|
||||
# --- Token routing ---
|
||||
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
|
||||
topk_ids,
|
||||
block_size_m,
|
||||
global_num_experts,
|
||||
expert_map,
|
||||
)
|
||||
|
||||
# --- Reuse pre-allocated buffers when possible ---
|
||||
if total_tokens <= layer.rdna3_w1_buf.shape[0]:
|
||||
w1_out = layer.rdna3_w1_buf[:total_tokens]
|
||||
w1_out.zero_()
|
||||
act_out = layer.rdna3_act_buf[:total_tokens]
|
||||
else:
|
||||
w1_out = torch.zeros(
|
||||
total_tokens,
|
||||
N_gate_up,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
act_out = torch.empty(
|
||||
total_tokens,
|
||||
intermediate_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# --- topk weights (pre-cast to float32 for kernel) ---
|
||||
topk_w_float = topk_weights.view(-1).float()
|
||||
empty_tw = layer.rdna3_empty_tw
|
||||
|
||||
# --- w1 GEMM: [M, K] -> [M*top_k, N_gate_up] ---
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
hidden_states,
|
||||
w1_out,
|
||||
layer.w13_weight_packed,
|
||||
layer.w13_weight_scale,
|
||||
layer.w13_qzeros,
|
||||
topk_w_float if apply_router_weight_on_input else empty_tw,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
top_k,
|
||||
block_size_m,
|
||||
apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
# --- Activation (silu_and_mul etc.) ---
|
||||
apply_moe_activation(activation, act_out, w1_out)
|
||||
|
||||
# --- w2 GEMM: [M*top_k, intermediate] -> [M, hidden] (fused reduce) ---
|
||||
# output_topk=top_k: kernel writes to out[token_id / top_k] directly,
|
||||
# fusing moe_sum into the atomic accumulation — saves one kernel launch
|
||||
# and the w2_out intermediate buffer.
|
||||
out = torch.zeros(
|
||||
num_tokens,
|
||||
hidden_size,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
act_out,
|
||||
out,
|
||||
layer.w2_weight_packed,
|
||||
layer.w2_weight_scale,
|
||||
layer.w2_qzeros,
|
||||
topk_w_float if not apply_router_weight_on_input else empty_tw,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
num_tokens_post_padded,
|
||||
1,
|
||||
block_size_m,
|
||||
not apply_router_weight_on_input,
|
||||
output_topk=top_k,
|
||||
)
|
||||
return out
|
||||
-48
@@ -1,48 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""ROCm MoE kernel dispatcher.
|
||||
|
||||
Selects architecture-specific native HIP MoE kernels in priority order.
|
||||
Falls back to the Triton WNA16 path when no native kernel is available.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def is_supported(weight_quant) -> bool:
|
||||
"""Check if a native ROCm MoE kernel is available for this config."""
|
||||
if weight_quant.num_bits != 4:
|
||||
return False
|
||||
|
||||
from vllm.platforms.rocm import on_gfx1100
|
||||
|
||||
# RDNA3 (gfx1100). Future: add RDNA4 (gfx12x), CDNA (gfx94x), etc.
|
||||
return (
|
||||
on_gfx1100()
|
||||
and hasattr(torch.ops, "_rocm_C")
|
||||
and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
|
||||
)
|
||||
|
||||
|
||||
def make_method(weight_quant, input_quant, moe_config):
|
||||
"""Create the native ROCm MoE method. Call only after is_supported()."""
|
||||
from vllm.platforms.rocm import on_gfx1100
|
||||
|
||||
if on_gfx1100():
|
||||
from .compressed_tensors_moe_wna16_rdna3 import (
|
||||
CompressedTensorsWNA16RDNA3MoEMethod,
|
||||
)
|
||||
|
||||
logger.info_once(
|
||||
"Using CompressedTensorsWNA16RDNA3MoEMethod (native RDNA3 HIP kernel)"
|
||||
)
|
||||
return CompressedTensorsWNA16RDNA3MoEMethod(
|
||||
weight_quant, input_quant, moe_config
|
||||
)
|
||||
|
||||
# Future: RDNA4, CDNA, etc.
|
||||
raise RuntimeError("is_supported() returned True but no kernel matched")
|
||||
Reference in New Issue
Block a user