[Bugfix][ROCm] RDNA3 W4A16 MoE: handle the N-first CT weight layout (#60388)

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

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

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

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)