feat(quantization): add compressed-tensors INT8 inference (#335)

### PR Category
Core

### PR Type
New Features, Bug Fixes, Improvements, Test Case

### Description
Add compressed-tensors INT8 inference support for channel-wise INT8
checkpoints, including canonical dynamic per-token W8A8 and packed INT8
formats. The W8A8 execution path uses native vLLM kernels instead of
FlagGems so linear layers avoid graph-breaking logging and routed MoE
follows the vLLM activation-quantization contract.

### Related Issues
N/A

### Changes
- Validate and register canonical compressed-tensors channel-wise W8A8
schemes.
- Run W8A8 linear layers through native vLLM Cutlass/Triton scaled-mm
kernels.
- Add a functional Triton MoE adapter that keeps activations
floating-point until vLLM performs dynamic per-token INT8 quantization.
- Support packed channel-wise INT8 checkpoints as W8A8 while preserving
weight-only handling for group-wise INT8.
- Add portable Triton WNA16 fallback and backend-selection compatibility
updates.
- Add focused reference, linear, MoE, packed INT8, and WNA16 tests.

### Testing
- ruff check on all changed Python modules and tests: passed.
- ruff format --check on all changed Python modules and tests: passed.
- git diff --check: passed.
- vLLM 0.20.2/0.24.0 fused_experts signature compatibility exercised
with a local stub: passed.
- Full vLLM/GPU unit and TP=4 model validation is pending because the
local macOS environment does not have vLLM or CUDA installed.

### Checklist
- [x] I have run the existing tests and they pass
- [x] I have added tests for my changes (if applicable)
- [x] I have updated the documentation (not applicable; repository
README intentionally unchanged)

---------

Co-authored-by: ceci3 <ceci3@users.noreply.github.com>

Author: 2 people

tests/unit_tests/ops/test_c_ops_registry.py

@@ -0,0 +1,101 @@
+# 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 vllm_fl
+from vllm_fl.ops import _C_ops_registry as registry
+
+
+def test_loads_legacy_and_stable_extensions_on_every_platform(monkeypatch):
+    calls = []
+    monkeypatch.setattr(registry, "_import_extension", calls.append)
+
+    assert registry.load_vllm_native_extensions()
+    assert calls == [
+        registry._LEGACY_C_EXTENSION,
+        registry._STABLE_C_EXTENSION,
+    ]
+
+
+def test_loads_stable_extension_without_legacy(monkeypatch):
+    calls = []
+
+    def import_extension(module_name):
+        calls.append(module_name)
+        if module_name == registry._LEGACY_C_EXTENSION:
+            raise ModuleNotFoundError(module_name)
+
+    monkeypatch.setattr(registry, "_import_extension", import_extension)
+
+    assert registry.load_vllm_native_extensions()
+    assert calls == [
+        registry._LEGACY_C_EXTENSION,
+        registry._STABLE_C_EXTENSION,
+    ]
+
+
+def test_loads_legacy_extension_without_stable(monkeypatch):
+    calls = []
+
+    def import_extension(module_name):
+        calls.append(module_name)
+        if module_name == registry._STABLE_C_EXTENSION:
+            raise ModuleNotFoundError(module_name)
+
+    monkeypatch.setattr(registry, "_import_extension", import_extension)
+
+    assert registry.load_vllm_native_extensions()
+    assert calls == [
+        registry._LEGACY_C_EXTENSION,
+        registry._STABLE_C_EXTENSION,
+    ]
+
+
+def test_fallback_schema_registration_skipped_when_native_extension_loaded(
+    monkeypatch,
+):
+    monkeypatch.delattr(registry.register_op_schemas, "_lib", raising=False)
+    monkeypatch.setattr(
+        registry,
+        "load_vllm_native_extensions",
+        lambda: True,
+    )
+
+    def fail_if_fallback_library_is_created(*args, **kwargs):
+        raise AssertionError("fallback schemas must not precede native CUDA ops")
+
+    monkeypatch.setattr(
+        registry.torch.library,
+        "Library",
+        fail_if_fallback_library_is_created,
+    )
+
+    registry.register_op_schemas()
+
+
+def test_plugin_initialization_loads_native_ops_before_fallback(monkeypatch):
+    monkeypatch.setattr(
+        registry,
+        "load_vllm_native_extensions",
+        lambda: True,
+    )
+
+    def fail_if_fallback_is_registered():
+        raise AssertionError("native ops must bypass fallback schema registration")
+
+    monkeypatch.setattr(
+        registry,
+        "register_op_schemas",
+        fail_if_fallback_is_registered,
+    )
+
+    vllm_fl._patch_custom_ops()

tests/unit_tests/quantization/test_compressed_tensors.py

@@ -4,16 +4,20 @@
 
 import pytest
 
+from vllm_fl import utils as fl_utils
 from vllm_fl.quantization import compressed_tensors
 from vllm_fl.quantization.compressed_tensors import (
     CompatibilityReport,
+    W8A8DynamicTokenScheme,
     WNA16Scheme,
     inspect_vllm_compressed_tensors_api,
     register_compressed_tensors_oot,
+    validate_compressed_tensors_w8a8_config,
     validate_compressed_tensors_wna16_config,
 )
 from vllm_fl.quantization.marlin import is_marlin_moe_platform
-from vllm_fl.quantization.wna16 import moe as moe_adapter
+from vllm_fl.quantization.w8a8 import moe as w8a8_moe_adapter, packed as packed_w8a8
+from vllm_fl.quantization.wna16 import kernels as wna16_kernels, moe as moe_adapter
 
 
 def _config():
@@ -38,6 +42,110 @@ def _config():
     }
 
 
+def _w8a8_config():
+    return {
+        "quant_method": "compressed-tensors",
+        "format": "int-quantized",
+        "quantization_status": "compressed",
+        "config_groups": {
+            "w8a8": {
+                "targets": ["Linear"],
+                "weights": {
+                    "num_bits": 8,
+                    "type": "int",
+                    "strategy": "channel",
+                    "symmetric": True,
+                    "dynamic": False,
+                },
+                "input_activations": {
+                    "num_bits": 8,
+                    "type": "int",
+                    "strategy": "token",
+                    "symmetric": True,
+                    "dynamic": True,
+                },
+            }
+        },
+        "ignore": [],
+    }
+
+
+def test_accepts_standard_dynamic_token_w8a8_config():
+    schemes = validate_compressed_tensors_w8a8_config(_w8a8_config())
+    assert schemes == [
+        W8A8DynamicTokenScheme(
+            weight_num_bits=8,
+            weight_type="int",
+            weight_strategy="channel",
+            weight_symmetric=True,
+            weight_dynamic=False,
+            weight_group_size=None,
+            input_num_bits=8,
+            input_type="int",
+            input_strategy="token",
+            input_symmetric=True,
+            input_dynamic=True,
+        )
+    ]
+
+
+@pytest.mark.parametrize(
+    ("section", "field", "value", "message"),
+    [
+        ("weights", "strategy", "group", "per-channel"),
+        ("weights", "group_size", 128, "must not set group_size"),
+        ("input_activations", "strategy", "tensor", "per-token"),
+        ("input_activations", "dynamic", False, "must be dynamic"),
+        ("input_activations", "symmetric", False, "symmetric"),
+    ],
+)
+def test_rejects_noncanonical_w8a8_config(section, field, value, message):
+    config = _w8a8_config()
+    config["config_groups"]["w8a8"][section][field] = value
+    with pytest.raises(ValueError, match=message):
+        validate_compressed_tensors_w8a8_config(config)
+
+
+def test_rejects_packed_format_for_canonical_w8a8():
+    config = _w8a8_config()
+    config["format"] = "pack-quantized"
+    with pytest.raises(ValueError, match="int-quantized"):
+        validate_compressed_tensors_w8a8_config(config)
+
+
+def test_w8a8_registration_does_not_depend_on_wna16_kernel(monkeypatch):
+    calls = []
+    report = CompatibilityReport(
+        vllm_version="0.20.2",
+        linear_wna16=True,
+        moe_wna16=True,
+    )
+    monkeypatch.setattr(
+        compressed_tensors,
+        "inspect_vllm_compressed_tensors_api",
+        lambda: report,
+    )
+    monkeypatch.setattr(fl_utils, "is_oot_enabled", lambda: True)
+    monkeypatch.setattr(
+        w8a8_moe_adapter,
+        "install_fl_w8a8_moe_selector",
+        lambda: calls.append("w8a8"),
+    )
+    monkeypatch.setattr(
+        packed_w8a8,
+        "install_packed_w8a8_scheme",
+        lambda: calls.append("packed-w8a8"),
+    )
+    monkeypatch.setattr(
+        wna16_kernels,
+        "is_wna16_moe_available",
+        lambda: False,
+    )
+
+    assert compressed_tensors.register_compressed_tensors_oot() is report
+    assert calls == ["packed-w8a8", "w8a8"]
+
+
 def test_accepts_standard_w4a16_group_config():
     schemes = validate_compressed_tensors_wna16_config(_config())
     assert schemes == [
@@ -83,6 +191,11 @@ def test_local_moe_adapter_is_not_installed_without_kernel(monkeypatch):
         "is_wna16_moe_available",
         lambda: False,
     )
+    monkeypatch.setattr(
+        moe_adapter,
+        "is_flaggems_wna16_moe_available",
+        lambda: False,
+    )
     assert moe_adapter.install_fl_wna16_moe_method() is False
 
 

tests/unit_tests/quantization/test_packed_int8.py

@@ -0,0 +1,143 @@
+# 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.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(
+        num_bits=8,
+        type=QuantizationType.INT,
+        strategy=strategy,
+        symmetric=True,
+        dynamic=False,
+        group_size=128 if strategy == QuantizationStrategy.GROUP else None,
+    )
+
+
+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),
+        _activation_args(),
+        "pack-quantized",
+    )
+
+
+def test_missing_activation_config_keeps_weight_only_scheme():
+    assert not should_use_packed_w8a8(
+        _weight_args(QuantizationStrategy.CHANNEL),
+        None,
+        "pack-quantized",
+    )
+
+
+def test_packed_w8a8_rejects_groupwise_weights():
+    with pytest.raises(ValueError, match="per-channel"):
+        should_use_packed_w8a8(
+            _weight_args(QuantizationStrategy.GROUP),
+            _activation_args(),
+            "pack-quantized",
+        )
+
+
+def test_static_tensor_activation_config_does_not_select_dynamic_w8a8():
+    assert not should_use_packed_w8a8(
+        _weight_args(QuantizationStrategy.CHANNEL),
+        _activation_args(
+            strategy=QuantizationStrategy.TENSOR,
+            dynamic=False,
+        ),
+        "pack-quantized",
+    )
+
+
+def test_packed_scheme_reuses_native_w8a8_execution(monkeypatch):
+    processed_layers = []
+
+    class FakeKernel:
+        def process_weights_after_loading(self, layer):
+            processed_layers.append(layer)
+
+        def apply_weights(self, layer, x, bias):
+            raise AssertionError("not used")
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.compressed_tensors.schemes."
+        "compressed_tensors_w8a8_int8.init_int8_linear_kernel",
+        lambda *args, **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.CompressedTensorsPackedW8A8Int8(layer_name="model.linear")
+    assert isinstance(scheme, CompressedTensorsW8A8Int8)
+
+    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 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]