diff --git a/tests/kernels/quantization/test_rdna3_moe_w4a16.py b/tests/kernels/quantization/test_rdna3_moe_w4a16.py index 6212e6f08c..58aa90aa07 100644 --- a/tests/kernels/quantization/test_rdna3_moe_w4a16.py +++ b/tests/kernels/quantization/test_rdna3_moe_w4a16.py @@ -18,6 +18,7 @@ Run `pytest tests/kernels/quantization/test_rdna3_moe_w4a16.py`. import pytest import torch +import torch.nn.functional as F from vllm.platforms import current_platform @@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402 moe_align_block_size, ) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( # noqa: E402 + _process_weights_rdna3, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402 pack_quantized_values_into_int32, ) @@ -95,6 +99,16 @@ def _make_qzeros(E, groups, N): return qz.unsqueeze(0).expand(E, -1, -1).contiguous() +def _make_ct_weights(E, N, K, group_size, dtype): + """Compressed-tensors canonical N-first layout: packed ``[E, N, K // 8]`` + and scales ``[E, N, K // group_size]``, plus the dequantized weights.""" + q = torch.randint(0, 16, (E, N, K), dtype=torch.int32, device=device) + s = torch.rand(E, N, K // group_size, dtype=dtype, device=device) * 0.1 + 0.01 + packed = pack_quantized_values_into_int32(q, scalar_types.uint4b8, packed_dim=2) + w = (q - 8).float() * s.float().repeat_interleave(group_size, dim=2) + return packed.contiguous(), s, w + + @gfx1100_only @pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS) @pytest.mark.parametrize("M", NUM_TOKENS) @@ -359,3 +373,70 @@ def test_expert_id_minus_one(): # Output should remain zero (expert skipped) assert torch.equal(out, torch.zeros_like(out)) + + +@gfx1100_only +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_process_weights_from_compressed_tensors_layout(dtype): + """Weights in the layout compressed-tensors hands to the oracle go through + the RDNA3 post-processing and a full MoE forward, against a dequantized + reference. The tests above build the kernel layout directly, so they do + not notice when the layout upstream of the kernel changes (#52798).""" + E, K, N_inter, top_k, group_size, M = 8, 512, 256, 2, 32, 16 + torch.manual_seed(5) + w13, w13_s, w13_ref = _make_ct_weights(E, 2 * N_inter, K, group_size, dtype) + w2, w2_s, w2_ref = _make_ct_weights(E, K, N_inter, group_size, dtype) + w13, w2, w13_s, w2_s, w13_z, w2_z, *_ = _process_weights_rdna3( + w13, w2, w13_s, w2_s, group_size + ) + assert w13_z is not None and w2_z is not None + + x = torch.randn(M, K, dtype=dtype, device=device) + topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32) + topk_w = torch.softmax(torch.randn(M, top_k, device=device), dim=-1).float() + si, ei, ntp = moe_align_block_size(topk_ids, 1, E) + + w1_out = torch.zeros(M * top_k, 2 * N_inter, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + x, + w1_out, + w13, + w13_s, + w13_z, + torch.empty(0, device=device), + si, + ei, + ntp, + top_k, + 1, + False, + 0, + ) + act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device) + apply_moe_activation(MoEActivation.SILU, act_out, w1_out) + out = torch.zeros(M, K, dtype=dtype, device=device) + ops.moe_gptq_gemm_rdna3( + act_out, + out, + w2, + w2_s, + w2_z, + topk_w.view(-1), + si, + ei, + ntp, + 1, + 1, + True, + top_k, + ) + + ref = torch.zeros(M, K, device=device) + for m in range(M): + for k in range(top_k): + e = int(topk_ids[m, k]) + h = w13_ref[e] @ x[m].float() + a = F.silu(h[:N_inter]) * h[N_inter:] + ref[m] += (w2_ref[e] @ a) * topk_w[m, k] + rel_l2 = (torch.norm(out.float() - ref) / torch.norm(ref)).item() + assert rel_l2 < (0.02 if dtype == torch.float16 else 0.05), rel_l2 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 3e91fd9e04..9b4b163379 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -926,15 +926,18 @@ def _process_weights_rdna3( ]: """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. + Transposes the canonical N-first inputs to the K-first layout of + ``moe_gptq_gemm_rdna3`` (weights ``[E, K // 8, N]``, scales + ``[E, groups, N]``), interleaves the packed nibbles per expert (the exllama + shuffle the dense RDNA3 kernel also uses) and synthesizes the symmetric + zero points the kernel dequantizes with. """ device = w13.device num_experts = w13.size(0) + w13 = w13.transpose(1, 2).contiguous() + w2 = w2.transpose(1, 2).contiguous() + for e in range(num_experts): w13_e = w13[e].contiguous() ops.gptq_shuffle(w13_e, 4) @@ -950,8 +953,8 @@ def _process_weights_rdna3( return ( w13, w2, - w13_scale.contiguous(), - w2_scale.contiguous(), + w13_scale.transpose(1, 2).contiguous(), + w2_scale.transpose(1, 2).contiguous(), _qzeros(w13), _qzeros(w2), None, # w13_input_global_scale 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 fdd37d1203..1e06f47413 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 @@ -450,10 +450,12 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): if w2_bias is not None: replace_parameter(layer, "w2_bias", w2_bias) - # 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)