diff --git a/tests/kernels/quantization/test_rdna3_compile_guards.py b/tests/kernels/quantization/test_rdna3_compile_guards.py index 3c1e226d2c07..194441f0f909 100644 --- a/tests/kernels/quantization/test_rdna3_compile_guards.py +++ b/tests/kernels/quantization/test_rdna3_compile_guards.py @@ -34,6 +34,13 @@ 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 @@ def test_no_toplevel_rocm_c_import(self): # ============================================================================ -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 @@ def test_chooser_skips_rdna3_when_mocked_cdna(self): ) -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 + ) diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py index 475ca84f6a17..0fd5a10120b3 100644 --- a/tests/quantization/test_moe_wna16.py +++ b/tests/quantization/test_moe_wna16.py @@ -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( diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 1312cbb24157..42ff415d39cc 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -141,6 +141,7 @@ def with_default( "flydsl", "hpc", "emulation", + "rdna3", ] # Backends that run the mega-MoE model path through the flashinfer moe_ep @@ -261,6 +262,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. diff --git a/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py b/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py new file mode 100644 index 000000000000..dbbce6b0b182 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/rdna3_moe.py @@ -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, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 899f9bc3d94a..d47ad9c97e31 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -7,6 +7,7 @@ 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.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, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 5d284bcc80f3..2dafb6c493d6 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -617,7 +617,6 @@ def weight_loader( # against known CompressionFormat enum values that have this quality if quant_method_name in ( "CompressedTensorsWNA16MoEMethod", - "CompressedTensorsWNA16RDNA3MoEMethod", "CompressedTensorsW4A16FlydslMoEMethod", ): if is_transposed: diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index bcc930471da7..f2d2a9c9452c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -88,14 +88,7 @@ def get_moe_method( 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() diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index 6a3f73e5392d..7a5b2d949b9f 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -443,10 +443,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: 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) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py deleted file mode 100644 index 05db6ec3ddd8..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_rdna3.py +++ /dev/null @@ -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 diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py deleted file mode 100644 index e72caa0796a5..000000000000 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/rocm_moe_rdna.py +++ /dev/null @@ -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")