[MoE Refactor] Migrate MoeWNA16Method quantization method over to using the new MK oracle scheme.

tests/quantization/test_auto_gptq.py

@@ -5,13 +5,17 @@
 Run `pytest tests/quantization/test_auto_gptq.py -v -s`.
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
 from tests.quantization.utils import is_quant_method_supported
+from vllm.model_executor.layers.fused_moe import RoutedExperts
 from vllm.model_executor.layers.quantization.auto_gptq import (
     AutoGPTQConfig,
     AutoGPTQLinearMethod,
+    AutoGPTQMoEMethod,
 )
 
 PROMPT = "On the surface of Mars, we found"
@@ -54,3 +58,78 @@ def check_model(model):
 def test_auto_gptq_config_get_name():
     """Test that AutoGPTQConfig.get_name() returns 'auto_gptq'."""
     assert AutoGPTQConfig.get_name() == "auto_gptq"
+
+
+def test_auto_gptq_moe_creates_zero_initialized_expert_biases():
+    method = object.__new__(AutoGPTQMoEMethod)
+    method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})
+    method.input_dtype = None
+    method.experts_cls = None
+    layer = torch.nn.Module()
+
+    method.create_weights(
+        layer=layer,
+        num_experts=2,
+        hidden_size=8,
+        intermediate_size_per_partition=4,
+        params_dtype=torch.float16,
+        intermediate_size_full=4,
+        weight_loader=lambda *args, **kwargs: None,
+    )
+
+    assert layer.w13_bias.shape == (2, 8)
+    assert layer.w2_bias.shape == (2, 8)
+    assert torch.count_nonzero(layer.w13_bias) == 0
+    assert torch.count_nonzero(layer.w2_bias) == 0
+
+
+def test_routed_experts_loads_per_expert_biases():
+    class Loader:
+        quant_config = None
+        quant_method = object()
+        moe_config = SimpleNamespace(
+            is_act_and_mul=True,
+            tp_rank=0,
+            moe_parallel_config=SimpleNamespace(tp_size=1),
+        )
+        _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
+        _narrow_expert_data_for_padding = staticmethod(
+            RoutedExperts._narrow_expert_data_for_padding
+        )
+        _load_w13 = RoutedExperts._load_w13
+        _loaded_expert_biases = set()
+
+        @staticmethod
+        def _map_global_expert_id_to_local_expert_id(expert_id):
+            return expert_id
+
+    loader = Loader()
+    w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False)
+    w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False)
+
+    for shard_id, loaded in (
+        ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])),
+        ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])),
+    ):
+        assert RoutedExperts.weight_loader(
+            loader,
+            w13_bias,
+            loaded,
+            weight_name="model.layers.0.mlp.experts.w13_bias",
+            shard_id=shard_id,
+            expert_id=0,
+            return_success=True,
+        )
+
+    assert RoutedExperts.weight_loader(
+        loader,
+        w2_bias,
+        torch.tensor([9.0, 10.0, 11.0, 12.0]),
+        weight_name="model.layers.0.mlp.experts.w2_bias",
+        shard_id="w2",
+        expert_id=0,
+        return_success=True,
+    )
+    assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8))
+    assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4))
+    assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"}

tests/quantization/test_moe_wna16.py

@@ -5,45 +5,194 @@
 
 import pytest
 import torch
+from compressed_tensors.quantization import (
+    ActivationOrdering,
+    QuantizationArgs,
+    QuantizationStrategy,
+    QuantizationType,
+)
 
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
-from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method
-from vllm.platforms import current_platform
+from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
+    WNA16MoEBackend,
+    _backend_incompatibility_reason,
+    _convert_moe_wna16_humming_tensors,
+    convert_to_wna16_moe_kernel_format,
+    map_wna16_backend,
+)
+from vllm.model_executor.layers.quantization import moe_wna16
+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,
+    MoeWNA16Method,
+)
 
 
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
-def test_moe_wna16_apply_passes_layer_activation(monkeypatch):
-    captured_kwargs = {}
+def test_map_wna16_backend_supports_triton():
+    assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON
 
-    def fake_fused_experts(*args, **kwargs):
-        captured_kwargs.update(kwargs)
-        return torch.empty(1, 2)
 
-    monkeypatch.setattr(
-        "vllm.model_executor.layers.fused_moe.fused_experts",
-        fake_fused_experts,
+@pytest.mark.parametrize(
+    ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"),
+    [
+        (
+            WNA16MoEBackend.TRITON,
+            AutoAWQConfig(4, 128, True, False),
+            True,
+            False,
+            "AutoAWQ weight layout",
+        ),
+        (
+            WNA16MoEBackend.TRITON,
+            AutoGPTQConfig(4, 128, True, True, False, {}, {}),
+            False,
+            False,
+            "activation ordering",
+        ),
+        (
+            WNA16MoEBackend.TRITON,
+            QuantizationArgs(
+                num_bits=4,
+                type=QuantizationType.INT,
+                strategy=QuantizationStrategy.GROUP,
+                symmetric=True,
+                dynamic=False,
+                group_size=128,
+                actorder=ActivationOrdering.GROUP,
+            ),
+            False,
+            False,
+            "activation ordering",
+        ),
+        (
+            WNA16MoEBackend.TRITON,
+            AutoGPTQConfig(4, 128, False, True, False, {}, {}),
+            False,
+            True,
+            "bias",
+        ),
+        (
+            WNA16MoEBackend.MARLIN,
+            MoeWNA16Config(
+                linear_quant_method="gptq",
+                weight_bits=4,
+                group_size=128,
+                has_zp=False,
+                lm_head_quantized=False,
+                modules_to_not_convert=None,
+                full_config={},
+            ),
+            False,
+            False,
+            "MoeWNA16 checkpoint layout",
+        ),
+    ],
+)
+def test_wna16_oracle_rejects_incompatible_quant_structures(
+    backend, quant_config, may_have_zp, may_have_bias, expected
+):
+    reason = _backend_incompatibility_reason(
+        backend=backend,
+        quant_config=quant_config,
+        may_have_zp=may_have_zp,
+        may_have_bias=may_have_bias,
+    )
+
+    assert reason is not None
+    assert expected in reason
+
+
+def test_compressed_tensors_weights_are_transposed_for_triton():
+    quant_config = QuantizationArgs(
+        num_bits=4,
+        type=QuantizationType.INT,
+        strategy=QuantizationStrategy.GROUP,
+        symmetric=True,
+        dynamic=False,
+        group_size=32,
     )
+    w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8)
+    w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6)
+    w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8)
+    w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6)
 
+    converted = convert_to_wna16_moe_kernel_format(
+        backend=WNA16MoEBackend.TRITON,
+        layer=torch.nn.Module(),
+        quant_config=quant_config,
+        input_dtype=None,
+        w13=w13,
+        w2=w2,
+        w13_scale=w13_scale,
+        w2_scale=w2_scale,
+    )
+
+    assert converted is not None
+    assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8))
+    assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8))
+    assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous())
+    assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous())
+
+
+def test_moe_wna16_setup_forwards_selected_backend(monkeypatch):
     method = object.__new__(MoeWNA16Method)
-    method.moe = SimpleNamespace(disable_inplace=False)
-    method.moe_quant_config = object()
-    layer = SimpleNamespace(
-        w13_qweight=torch.empty(1, 2),
-        w2_qweight=torch.empty(1, 2),
-        activation=MoEActivation.GELU_TANH,
-        apply_router_weight_on_input=False,
-        global_num_experts=1,
-        expert_map=None,
+    method.experts_cls = object
+    method.wna16_backend = WNA16MoEBackend.HUMMING
+    method.moe = object()
+    quant_config = object()
+    method.get_fused_moe_quant_config = lambda layer: quant_config
+    layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None))
+    captured = {}
+    kernel = object()
+
+    def fake_make_wna16_moe_kernel(**kwargs):
+        captured.update(kwargs)
+        return kernel
+
+    monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel)
+
+    method._setup_kernel(layer)
+
+    assert method.moe_kernel is kernel
+    assert captured["backend"] == WNA16MoEBackend.HUMMING
+    assert captured["layer"] is layer
+
+
+def test_moe_wna16_humming_adapter_repacks_uint8_tensors():
+    qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8)
+    scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4)
+    qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2)
+
+    converted = _convert_moe_wna16_humming_tensors(
+        {"qweight": qweight, "scales": scales, "qzeros": qzeros},
+        has_zero_point=True,
     )
 
-    output = method.apply(
-        layer,
-        x=torch.empty(1, 2),
-        topk_weights=torch.empty(1, 1),
-        topk_ids=torch.empty(1, 1, dtype=torch.int32),
-        shared_experts=None,
-        shared_experts_input=None,
+    assert torch.equal(converted["weight"], qweight.view(torch.int32))
+    assert converted["weight"].shape == (1, 4, 2)
+    assert torch.equal(converted["weight_scale"], scales)
+    expected_qzeros = (
+        qzeros.transpose(-1, -2)
+        .contiguous()
+        .view(torch.int32)
+        .transpose(-1, -2)
+        .contiguous()
+    )
+    assert torch.equal(converted["zero_point"], expected_qzeros)
+    assert converted["zero_point"].shape == (1, 2, 2)
+
+
+def test_moe_wna16_uses_humming_quant_config(monkeypatch):
+    from vllm.model_executor.layers.quantization.utils import humming_utils
+
+    method = object.__new__(MoeWNA16Method)
+    method.wna16_backend = WNA16MoEBackend.HUMMING
+    layer = object()
+    quant_config = object()
+    monkeypatch.setattr(
+        humming_utils,
+        "get_humming_moe_quant_config",
+        lambda actual_layer: quant_config if actual_layer is layer else None,
     )
 
-    assert output.shape == (1, 2)
-    assert captured_kwargs["activation"] is MoEActivation.GELU_TANH
+    assert method.get_fused_moe_quant_config(layer) is quant_config