Migrate FP8/TorchAO tests to test/registered/quant/ (#16453)

test/registered/quant/test_autoround.py

@@ -17,7 +17,6 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-# AutoRound quantization tests
 register_cuda_ci(est_time=77, suite="stage-b-test-small-1-gpu")
 
 

test/registered/quant/test_awq_dequant.py

@@ -19,7 +19,6 @@ from sglang.srt.layers.quantization.awq_triton import (
 from sglang.test.ci.ci_register import register_amd_ci
 from sglang.test.test_utils import CustomTestCase
 
-# AWQ dequantization tests (AMD only)
 register_amd_ci(est_time=2, suite="stage-a-test-1")
 
 device = "cuda"

test/registered/quant/test_block_int8.py

@@ -10,7 +10,6 @@ from sglang.srt.server_args import ServerArgs, set_global_server_args_for_schedu
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-# Block INT8 quantization kernel tests
 register_cuda_ci(est_time=44, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=22, suite="stage-a-test-1")
 

test/registered/quant/test_eval_fp8_accuracy.py

@@ -0,0 +1,118 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.utils import is_hip, kill_process_tree
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_MODEL_NAME_FOR_ACCURACY_TEST_FP8,
+    DEFAULT_MODEL_NAME_FOR_DYNAMIC_QUANT_ACCURACY_TEST_FP8,
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=250, suite="stage-b-test-small-1-gpu")
+register_amd_ci(est_time=303, suite="stage-b-test-small-1-gpu")
+
+
+class TestEvalFP8Accuracy(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_ACCURACY_TEST_FP8
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_mmlu(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmlu",
+            num_examples=64,
+            num_threads=32,
+            temperature=0.1,
+        )
+
+        metrics = run_eval(args)
+        if is_hip():
+            # Another threshold for AMD because fp8 dtype is difference
+            self.assertGreaterEqual(metrics["score"], 0.60)
+        else:
+            self.assertGreaterEqual(metrics["score"], 0.60)
+
+
+class TestEvalFP8DynamicQuantAccuracy(CustomTestCase):
+
+    def _run_test(self, model, other_args, expected_score):
+        base_url = DEFAULT_URL_FOR_TEST
+        other_args = other_args or []
+
+        process = popen_launch_server(
+            model,
+            base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+        )
+
+        try:
+            args = SimpleNamespace(
+                base_url=base_url,
+                model=model,
+                eval_name="mmlu",
+                num_examples=64,
+                num_threads=32,
+                temperature=0.1,
+            )
+
+            metrics = run_eval(args)
+            self.assertGreaterEqual(metrics["score"], expected_score)
+        finally:
+            kill_process_tree(process.pid)
+
+    def test_mmlu_offline_only(self):
+        """Test with offline quantization only."""
+        self._run_test(
+            model=DEFAULT_MODEL_NAME_FOR_DYNAMIC_QUANT_ACCURACY_TEST_FP8,
+            other_args=[],
+            expected_score=0.64,
+        )
+
+    def test_mmlu_offline_and_online_override(self):
+        """Test with both offline and online quantization."""
+        self._run_test(
+            model=DEFAULT_MODEL_NAME_FOR_DYNAMIC_QUANT_ACCURACY_TEST_FP8,
+            other_args=["--quantization", "w8a8_fp8"],
+            # inference will use sgl kernel w/ online quant override
+            # we observed that the accuracy is higher then offline only
+            expected_score=0.64,
+        )
+
+    def test_mmlu_online_only(self):
+        """Test with online quantization only."""
+        self._run_test(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            # inference will use sgl kernel w/ online quantization only
+            # we observed that the accuracy is higher then offline only
+            other_args=["--quantization", "w8a8_fp8"],
+            expected_score=0.64,
+        )
+
+    def test_mmlu_fp16_baseline(self):
+        """Test with unquantized fp16 baseline."""
+        self._run_test(
+            model=DEFAULT_MODEL_NAME_FOR_TEST,
+            other_args=[],
+            expected_score=0.64,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/quant/test_fp8_kernel.py

@@ -9,7 +9,6 @@ from sglang.srt.layers.quantization.fp8_kernel import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-# FP8 quantization kernel tests
 register_cuda_ci(est_time=10, suite="stage-b-test-small-1-gpu")
 
 

test/registered/quant/test_fp8_utils.py

@@ -0,0 +1,47 @@
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.fp8_utils import (
+    inverse_transform_scale_ue8m0,
+    quant_weight_ue8m0,
+    transform_scale_ue8m0,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
+
+
+class TestInverseTransformScaleUe8m0(CustomTestCase):
+    def test_round_trip(self):
+        for _ in range(100):
+            weight_bf16 = torch.randn(
+                # DeepSeek V3 kv_b_proj
+                (32768, 512),
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+
+            weight_block_size = [128, 128]
+
+            qweight, sf_fp32_original = quant_weight_ue8m0(
+                weight_bf16, weight_block_size=weight_block_size
+            )
+            mn = qweight.shape[-2]
+
+            sf_packed_original = transform_scale_ue8m0(sf_fp32_original, mn=mn)
+            sf_fp32_recreated = inverse_transform_scale_ue8m0(sf_packed_original, mn=mn)
+
+            sf_packed_recreated = transform_scale_ue8m0(sf_fp32_recreated, mn=mn)
+
+            assert torch.all(
+                sf_packed_original == sf_packed_recreated
+            ), f"{sf_packed_original=} {sf_packed_recreated}"
+            assert torch.all(
+                sf_fp32_original == sf_fp32_recreated
+            ), f"{sf_fp32_original=} {sf_fp32_recreated}"
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/quant/test_fused_rms_fp8_group_quant.py

@@ -7,7 +7,6 @@ import torch.nn.functional as F
 from sglang.test.ci.ci_register import register_amd_ci
 from sglang.test.test_utils import CustomTestCase
 
-# Fused RMS FP8 group quantization tests (AMD/ROCm only)
 register_amd_ci(est_time=10, suite="stage-a-test-1")
 
 

test/registered/quant/test_int8_kernel.py

@@ -11,7 +11,6 @@ from sglang.srt.server_args import ServerArgs, set_global_server_args_for_schedu
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-# INT8 quantization kernel tests
 register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
 
 

test/registered/quant/test_torchao.py

@@ -0,0 +1,98 @@
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang import Engine
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=103, suite="stage-b-test-small-1-gpu")
+from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.srt.utils import kill_process_tree
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_IMAGE_URL,
+    DEFAULT_MODEL_NAME_FOR_TEST,
+    DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestTorchAO(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=["--torchao-config", "int4wo-128"],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_mmlu(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="mmlu",
+            num_examples=64,
+            num_threads=32,
+        )
+
+        metrics = run_eval(args)
+        assert metrics["score"] >= 0.60
+
+    def run_decode(self, max_new_tokens):
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "The capital of France is",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": max_new_tokens,
+                },
+                "ignore_eos": True,
+            },
+        )
+        return response.json()
+
+    def test_throughput(self):
+        import time
+
+        max_tokens = 256
+
+        tic = time.perf_counter()
+        res = self.run_decode(max_tokens)
+        tok = time.perf_counter()
+        print(res["text"])
+        throughput = max_tokens / (tok - tic)
+        print(f"Throughput: {throughput} tokens/s")
+        assert throughput >= 210
+
+
+class TestTorchAOForVLM(CustomTestCase):
+    def test_vlm_generate(self):
+        model_path = DEFAULT_SMALL_VLM_MODEL_NAME_FOR_TEST
+        chat_template = get_chat_template_by_model_path(model_path)
+        text = f"{chat_template.image_token}What is in this picture? Answer: "
+
+        engine = Engine(
+            model_path=model_path,
+            max_total_tokens=512,
+            enable_multimodal=True,
+            torchao_config="fp8wo",
+        )
+        out = engine.generate([text], image_data=[DEFAULT_IMAGE_URL])
+        engine.shutdown()
+        self.assertGreater(len(out), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/quant/test_triton_scaled_mm.py

@@ -8,7 +8,6 @@ from sglang.srt.layers.quantization.fp8_kernel import triton_scaled_mm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-# Triton scaled matrix multiplication tests
 register_cuda_ci(est_time=8, suite="stage-b-test-small-1-gpu")
 register_amd_ci(est_time=12, suite="stage-a-test-1")
 

test/registered/quant/test_w8a8_quantization.py

@@ -14,7 +14,6 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-# W8A8 quantization server integration tests
 register_cuda_ci(est_time=160, suite="stage-b-test-small-1-gpu")