fix(quantization): address INT8 review feedback

Author: rdzhu225

tests/unit_tests/quantization/test_packed_int8.py

@@ -19,12 +19,13 @@
     QuantizationType,
 )
 
-from vllm_fl.quantization.w8a8 import packed
-from vllm_fl.quantization.w8a8.int8_mode import (
-    INT8_MODE_ENV,
-    should_use_packed_w8a8,
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
+    CompressedTensorsW8A8Int8,
 )
 
+from vllm_fl.quantization.w8a8 import packed
+from vllm_fl.quantization.w8a8.int8_mode import should_use_packed_w8a8
+
 
 def _weight_args(strategy: QuantizationStrategy) -> QuantizationArgs:
     return QuantizationArgs(
@@ -37,50 +38,66 @@ def _weight_args(strategy: QuantizationStrategy) -> QuantizationArgs:
     )
 
 
-def test_auto_mode_maps_channelwise_packed_int8_to_w8a8(monkeypatch):
-    monkeypatch.delenv(INT8_MODE_ENV, raising=False)
+def _activation_args(
+    *,
+    strategy: QuantizationStrategy = QuantizationStrategy.TOKEN,
+    dynamic: bool = True,
+) -> QuantizationArgs:
+    return QuantizationArgs(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=strategy,
+        symmetric=True,
+        dynamic=dynamic,
+    )
+
+
+def test_model_config_maps_packed_channelwise_w8a8():
     assert should_use_packed_w8a8(
         _weight_args(QuantizationStrategy.CHANNEL),
-        None,
-        "pack-quantized",
-    )
-    assert not should_use_packed_w8a8(
-        _weight_args(QuantizationStrategy.GROUP),
-        None,
+        _activation_args(),
         "pack-quantized",
     )
 
 
-def test_w8a16_mode_keeps_channelwise_checkpoint_weight_only(monkeypatch):
-    monkeypatch.setenv(INT8_MODE_ENV, "w8a16")
+def test_missing_activation_config_keeps_weight_only_scheme():
     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"):
+def test_packed_w8a8_rejects_groupwise_weights():
+    with pytest.raises(ValueError, match="per-channel"):
         should_use_packed_w8a8(
             _weight_args(QuantizationStrategy.GROUP),
-            None,
+            _activation_args(),
             "pack-quantized",
         )
 
 
-def test_packed_scheme_uses_native_compatible_layer_contract(monkeypatch):
+def test_static_activation_config_does_not_select_dynamic_w8a8():
+    assert not should_use_packed_w8a8(
+        _weight_args(QuantizationStrategy.CHANNEL),
+        _activation_args(dynamic=False),
+        "pack-quantized",
+    )
+
+
+def test_packed_scheme_reuses_native_w8a8_execution(monkeypatch):
+    processed_layers = []
+
     class FakeKernel:
         def process_weights_after_loading(self, layer):
-            pass
+            processed_layers.append(layer)
 
         def apply_weights(self, layer, x, bias):
             raise AssertionError("not used")
 
     monkeypatch.setattr(
-        packed,
-        "create_w8a8_linear_kernel",
+        "vllm.model_executor.layers.quantization.compressed_tensors.schemes."
+        "compressed_tensors_w8a8_int8.init_int8_linear_kernel",
         lambda *args, **kwargs: FakeKernel(),
     )
     monkeypatch.setattr(
@@ -91,7 +108,10 @@ def apply_weights(self, layer, x, bias):
         "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
         lambda: 1,
     )
-    scheme = packed.FLPackedW8A8Scheme(layer_name="model.linear")
+
+    scheme = packed.CompressedTensorsPackedW8A8Int8(layer_name="model.linear")
+    assert isinstance(scheme, CompressedTensorsW8A8Int8)
+
     layer = torch.nn.Module()
     scheme.create_weights(
         layer,
@@ -102,6 +122,19 @@ def apply_weights(self, layer, x, bias):
     )
 
     assert layer.logical_widths == [4, 4]
+    assert not hasattr(layer, "weight")
     assert layer.weight_packed.shape == (8, 2)
     assert layer.weight_scale.shape == (8, 1)
     assert layer.weight_scale.dtype == torch.float32
+
+    unpacked = torch.ones((8, 8), dtype=torch.int8)
+    monkeypatch.setattr(
+        packed,
+        "unpack_uint8b128_int32",
+        lambda *args, **kwargs: unpacked,
+    )
+    scheme.process_weights_after_loading(layer)
+
+    assert not hasattr(layer, "weight_packed")
+    assert torch.equal(layer.weight, unpacked)
+    assert processed_layers == [layer]

tests/unit_tests/quantization/test_w8a8_linear.py

@@ -137,12 +137,10 @@ def upstream_selector(*args, **kwargs):
         activation_key=None,
     )
 
-    from vllm_fl.quantization.w8a8.moe_experts import (
-        VllmFunctionalW8A8Experts,
-    )
+    from vllm_fl.quantization.w8a8.moe_experts import TritonW8A8Experts
 
     assert backend == "triton"
-    assert experts_cls is VllmFunctionalW8A8Experts
+    assert experts_cls is TritonW8A8Experts
     assert upstream_calls == []
 
 

tests/unit_tests/quantization/test_w8a8_moe.py

@@ -46,8 +46,7 @@ def _apply_arguments():
 def test_functional_experts_defer_activation_quantization():
     instance = SimpleNamespace()
     assert (
-        moe_experts.VllmFunctionalW8A8Experts.expects_unquantized_inputs.fget(instance)
-        is True
+        moe_experts.TritonW8A8Experts.expects_unquantized_inputs.fget(instance) is True
     )
 
 
@@ -114,7 +113,7 @@ def fake_fused_experts(
     instance = SimpleNamespace(quant_config=quant_config)
     arguments = _apply_arguments()
 
-    moe_experts.VllmFunctionalW8A8Experts.apply(instance, **arguments)
+    moe_experts.TritonW8A8Experts.apply(instance, **arguments)
 
     assert calls[0]["hidden_states"].dtype == torch.bfloat16
     assert calls[0]["quant_config"] is quant_config
@@ -140,4 +139,4 @@ def test_functional_experts_rejects_prequantized_input(monkeypatch):
     arguments["a1q_scale"] = torch.ones((2, 1), dtype=torch.float32)
 
     with pytest.raises(ValueError, match="quantized before"):
-        moe_experts.VllmFunctionalW8A8Experts.apply(instance, **arguments)
+        moe_experts.TritonW8A8Experts.apply(instance, **arguments)

tests/unit_tests/quantization/test_w8a8_moe_experts.py

@@ -312,19 +312,17 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        # Dynamic W8A8 is handled by VllmFunctionalW8A8Experts so vLLM owns
+        # Dynamic W8A8 is handled by TritonW8A8Experts so vLLM owns
         # activation quantization. Do not allow it to fall back into the
         # FlagGems contract, which expects floating-point input here.
         if self.quant_config.use_int8_w8a8:
             raise RuntimeError(
-                "W8A8 MoE must use VllmFunctionalW8A8Experts, not TritonExpertsFL"
+                "W8A8 MoE must use TritonW8A8Experts, not TritonExpertsFL"
             )
 
-        # Fast path (no LoRA): let FlagGems own both expert GEMMs for
-        # unquantized and W8A16 inputs.
-        from vllm_fl.utils import use_flaggems_op
-
-        if self._lora_context is None and use_flaggems_op("fused_moe"):
+        # Fast path (no LoRA, NVIDIA only): let FlagGems own both expert GEMMs
+        # for unquantized and W8A16 inputs.
+        if self._lora_context is None and current_platform.is_cuda():
             import flag_gems
 
             output.copy_(