feat(quantization): adapt W8A8 inference to vLLM 0.24

Author: rdzhu225

tests/unit_tests/quantization/test_fused_moe_factory.py

@@ -0,0 +1,58 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+from vllm_fl.ops.fused_moe import layer as layer_module
+
+
+class _FakeRunner:
+    def __init__(self, quant_method):
+        self._quant_method = quant_method
+        self.moe_config = SimpleNamespace()
+        self.replacements = []
+
+    def _replace_quant_method(self, quant_method):
+        self.replacements.append(quant_method)
+        self._quant_method = quant_method
+
+
+def test_fused_moe_factory_preserves_quantized_method(monkeypatch):
+    quantized_method = object()
+    runner = _FakeRunner(quantized_method)
+    monkeypatch.setattr(layer_module, "_OrigFusedMoE", lambda *args, **kwargs: runner)
+    monkeypatch.setattr(layer_module, "replace_router_with_fl", lambda: None)
+
+    result = layer_module.FusedMoEFL()
+
+    assert result is runner
+    assert runner._quant_method is quantized_method
+    assert runner.replacements == []
+
+
+def test_fused_moe_factory_still_replaces_unquantized_method(monkeypatch):
+    upstream_method = object.__new__(layer_module.UnquantizedFusedMoEMethod)
+    fl_method = object()
+    runner = _FakeRunner(upstream_method)
+    monkeypatch.setattr(layer_module, "_OrigFusedMoE", lambda *args, **kwargs: runner)
+    monkeypatch.setattr(
+        layer_module,
+        "UnquantizedFusedMoEMethodFL",
+        lambda moe_config: fl_method,
+    )
+    monkeypatch.setattr(layer_module, "replace_router_with_fl", lambda: None)
+
+    layer_module.FusedMoEFL()
+
+    assert runner._quant_method is fl_method
+    assert runner.replacements == [fl_method]

tests/unit_tests/quantization/test_packed_int8.py

@@ -0,0 +1,117 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from compressed_tensors.quantization import (
+    QuantizationArgs,
+    QuantizationStrategy,
+    QuantizationType,
+)
+
+from vllm_fl.quantization.w8a8 import packed
+from vllm_fl.quantization.w8a8.int8_mode import (
+    INT8_MODE_ENV,
+    should_use_packed_w8a8,
+)
+
+
+def _weight_args(strategy: QuantizationStrategy) -> QuantizationArgs:
+    return QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=strategy,
+        symmetric=True,
+        dynamic=False,
+        group_size=128 if strategy == QuantizationStrategy.GROUP else None,
+    )
+
+
+def test_auto_mode_maps_channelwise_packed_int8_to_w8a8(monkeypatch):
+    monkeypatch.delenv(INT8_MODE_ENV, raising=False)
+    assert should_use_packed_w8a8(
+        _weight_args(QuantizationStrategy.CHANNEL),
+        None,
+        "pack-quantized",
+    )
+    assert not should_use_packed_w8a8(
+        _weight_args(QuantizationStrategy.GROUP),
+        None,
+        "pack-quantized",
+    )
+
+
+def test_w8a16_mode_keeps_channelwise_checkpoint_weight_only(monkeypatch):
+    monkeypatch.setenv(INT8_MODE_ENV, "w8a16")
+    assert not should_use_packed_w8a8(
+        _weight_args(QuantizationStrategy.CHANNEL),
+        None,
+        "pack-quantized",
+    )
+
+
+def test_w8a8_mode_rejects_groupwise_checkpoint(monkeypatch):
+    monkeypatch.setenv(INT8_MODE_ENV, "w8a8")
+    with pytest.raises(ValueError, match="--strategy channel"):
+        should_use_packed_w8a8(
+            _weight_args(QuantizationStrategy.GROUP),
+            None,
+            "pack-quantized",
+        )
+
+
+def test_packed_scheme_matches_vllm_024_layer_contract(monkeypatch):
+    class FakeKernel:
+        def process_weights_after_loading(self, layer):
+            assert layer.weight.dtype == torch.int8
+
+        def apply_weights(self, layer, x, bias):
+            raise AssertionError("not used")
+
+    monkeypatch.setattr(
+        packed,
+        "init_int8_linear_kernel",
+        lambda **kwargs: FakeKernel(),
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
+        lambda: 0,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
+        lambda: 1,
+    )
+
+    scheme = packed.FLPackedW8A8Scheme(layer_name="model.linear")
+    layer = torch.nn.Module()
+    scheme.create_weights(
+        layer,
+        output_partition_sizes=[4, 4],
+        input_size_per_partition=8,
+        params_dtype=torch.bfloat16,
+        weight_loader=lambda *args, **kwargs: None,
+    )
+
+    assert layer.logical_widths == [4, 4]
+    assert layer.weight_packed.shape == (8, 2)
+    assert layer.weight_scale.shape == (8, 1)
+    assert layer.weight_scale.dtype == torch.float32
+
+    values = torch.arange(-32, 32, dtype=torch.int8).reshape(8, 8)
+    layer.weight_packed.data.copy_(
+        (values.to(torch.int16) + 128).to(torch.uint8).contiguous().view(torch.int32)
+    )
+    scheme.process_weights_after_loading(layer)
+
+    assert not hasattr(layer, "weight_packed")
+    assert torch.equal(layer.weight, values)

tests/unit_tests/quantization/test_w8a8_moe.py

@@ -0,0 +1,131 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from types import SimpleNamespace
+
+import vllm.platforms as platforms
+
+from vllm_fl.quantization.w8a8 import moe as moe_adapter
+
+
+def _install_with_fake_modules(monkeypatch, upstream_selector, upstream_builder):
+    oracle = SimpleNamespace(
+        Int8MoeBackend=SimpleNamespace(TRITON="triton"),
+        select_int8_moe_backend=upstream_selector,
+    )
+    scheme = SimpleNamespace(
+        select_int8_moe_backend=upstream_selector,
+        make_int8_moe_quant_config=upstream_builder,
+    )
+    modules = {
+        moe_adapter._ORACLE_MODULE: oracle,
+        moe_adapter._SCHEME_MODULE: scheme,
+    }
+    monkeypatch.setattr(moe_adapter, "import_module", lambda name: modules[name])
+    return oracle, scheme
+
+
+def test_w8a8_moe_builder_preserves_dynamic_per_token_config(monkeypatch):
+    def upstream_selector(*args, **kwargs):
+        return "upstream"
+
+    def upstream_builder(*args, **kwargs):
+        return "upstream-config"
+
+    oracle, scheme = _install_with_fake_modules(
+        monkeypatch,
+        upstream_selector,
+        upstream_builder,
+    )
+    assert moe_adapter.install_fl_w8a8_moe_selector()
+    assert oracle.select_int8_moe_backend is scheme.select_int8_moe_backend
+
+    config_module = SimpleNamespace(
+        int8_w8a8_moe_quant_config=lambda **kwargs: kwargs,
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm.model_executor.layers.fused_moe.config",
+        config_module,
+    )
+    dynamic_config = scheme.make_int8_moe_quant_config(
+        w1_scale="w1",
+        w2_scale="w2",
+        a1_scale=None,
+        a2_scale=None,
+        w1_bias="b1",
+        w2_bias="b2",
+        per_act_token_quant=True,
+    )
+    assert dynamic_config == {
+        "w1_scale": "w1",
+        "w2_scale": "w2",
+        "a1_scale": None,
+        "a2_scale": None,
+        "w1_bias": "b1",
+        "w2_bias": "b2",
+        "per_act_token_quant": True,
+    }
+    assert (
+        scheme.make_int8_moe_quant_config(
+            w1_scale="w1",
+            w2_scale="w2",
+            per_act_token_quant=False,
+        )
+        == "upstream-config"
+    )
+
+
+def test_w8a8_moe_selector_uses_vllm_functional_experts(monkeypatch):
+    upstream_calls = []
+
+    def upstream_selector(*args, **kwargs):
+        upstream_calls.append((args, kwargs))
+        return "upstream"
+
+    oracle, _ = _install_with_fake_modules(
+        monkeypatch,
+        upstream_selector,
+        lambda **kwargs: kwargs,
+    )
+    monkeypatch.setattr(
+        type(platforms.current_platform),
+        "is_out_of_tree",
+        lambda self: True,
+    )
+
+    import vllm_fl.utils as fl_utils
+
+    monkeypatch.setattr(fl_utils, "is_oot_enabled", lambda: True)
+    moe_adapter.install_fl_w8a8_moe_selector()
+    config = SimpleNamespace(
+        is_lora_enabled=False,
+        moe_parallel_config=SimpleNamespace(
+            use_batched_activation_format=False,
+        ),
+    )
+
+    backend, experts_cls = oracle.select_int8_moe_backend(
+        config,
+        weight_key=None,
+        activation_key=None,
+    )
+
+    from vllm_fl.quantization.w8a8.moe_experts import (
+        VllmFunctionalW8A8Experts,
+    )
+
+    assert backend == "triton"
+    assert experts_cls is VllmFunctionalW8A8Experts
+    assert upstream_calls == []

tests/unit_tests/quantization/test_w8a8_moe_experts.py

@@ -0,0 +1,90 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+
+from vllm_fl.quantization.w8a8 import moe_experts
+
+
+def _apply_arguments():
+    hidden_states = torch.ones((2, 4), dtype=torch.bfloat16)
+    return {
+        "output": torch.empty_like(hidden_states),
+        "hidden_states": hidden_states,
+        "w1": torch.ones((2, 8, 4), dtype=torch.int8),
+        "w2": torch.ones((2, 4, 4), dtype=torch.int8),
+        "topk_weights": torch.ones((2, 1), dtype=torch.float32),
+        "topk_ids": torch.zeros((2, 1), dtype=torch.int64),
+        "activation": MoEActivation.SILU,
+        "global_num_experts": 2,
+        "expert_map": None,
+        "a1q_scale": None,
+        "a2_scale": None,
+        "workspace13": torch.empty(0),
+        "workspace2": torch.empty(0),
+        "expert_tokens_meta": None,
+        "apply_router_weight_on_input": False,
+    }
+
+
+def test_functional_experts_defer_activation_quantization():
+    instance = SimpleNamespace()
+    assert (
+        moe_experts.VllmFunctionalW8A8Experts.expects_unquantized_inputs.fget(instance)
+        is True
+    )
+
+
+def test_functional_experts_calls_vllm_024_with_float_input(monkeypatch):
+    calls = []
+
+    def fake_fused_experts(**kwargs):
+        calls.append(kwargs)
+        return torch.full_like(kwargs["hidden_states"], 3)
+
+    monkeypatch.setattr(moe_experts, "fused_experts", fake_fused_experts)
+    quant_config = SimpleNamespace(use_int8_w8a8=True)
+    instance = SimpleNamespace(quant_config=quant_config)
+    arguments = _apply_arguments()
+
+    moe_experts.VllmFunctionalW8A8Experts.apply(instance, **arguments)
+
+    assert calls[0]["hidden_states"].dtype == torch.bfloat16
+    assert calls[0]["quant_config"] is quant_config
+    assert "inplace" not in calls[0]
+    assert torch.equal(
+        arguments["output"],
+        torch.full_like(arguments["output"], 3),
+    )
+
+
+def test_functional_experts_rejects_prequantized_input(monkeypatch):
+    monkeypatch.setattr(
+        moe_experts,
+        "fused_experts",
+        lambda **kwargs: pytest.fail("fused_experts must not run"),
+    )
+    instance = SimpleNamespace(
+        quant_config=SimpleNamespace(use_int8_w8a8=True),
+    )
+    arguments = _apply_arguments()
+    arguments["hidden_states"] = torch.ones((2, 4), dtype=torch.int8)
+    arguments["a1q_scale"] = torch.ones((2, 1), dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="quantized before"):
+        moe_experts.VllmFunctionalW8A8Experts.apply(instance, **arguments)