diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index d6e13b4eb..0a03f5dc4 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1320,10 +1320,3 @@ class Fp8KVCacheMethod(BaseKVCacheMethod): def __init__(self, quant_config: Fp8Config): super().__init__(quant_config) - - -if _is_cuda: - - @torch.library.register_fake("sgl_kernel::fp8_scaled_mm") - def _(mat_a, mat_b, scales_a, scales_b, out_dtype, bias): - return mat_a.new_empty((mat_a.shape[0], mat_b.shape[-1]), dtype=out_dtype) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 5cbde4c3e..f9578e4fe 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -52,6 +52,14 @@ if _use_aiter: if _is_cuda: from sgl_kernel import fp8_blockwise_scaled_mm, fp8_scaled_mm + @torch.library.register_fake("sgl_kernel::fp8_scaled_mm") + def _fp8_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): + # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N]. + M = mat_a.shape[-2] + N = mat_b.shape[-1] + return mat_a.new_empty((M, N), dtype=out_dtype) + + use_vllm_cutlass_w8a8_fp8_kernel = get_bool_env_var("USE_VLLM_CUTLASS_W8A8_FP8_KERNEL") use_triton_w8a8_fp8_kernel = get_bool_env_var("USE_TRITON_W8A8_FP8_KERNEL") diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index a83a0df20..b4ae69f54 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -94,7 +94,7 @@ suites = { TestFile("test_original_logprobs.py", 41), TestFile("test_page_size.py", 60), TestFile("test_penalty.py", 82), - TestFile("test_piecewise_cuda_graph.py", 450), + TestFile("test_piecewise_cuda_graph.py", 600), TestFile("test_priority_scheduling.py", 130), TestFile("test_pytorch_sampling_backend.py", 66), TestFile("test_radix_attention.py", 105), diff --git a/test/srt/test_piecewise_cuda_graph.py b/test/srt/test_piecewise_cuda_graph.py index 19723e5e6..3dc410ebf 100644 --- a/test/srt/test_piecewise_cuda_graph.py +++ b/test/srt/test_piecewise_cuda_graph.py @@ -209,5 +209,44 @@ class TestPiecewiseCudaGraphAWQ(CustomTestCase): self.assertGreaterEqual(metrics["score"], 0.65) +class TestPiecewiseCudaGraphFP8(CustomTestCase): + """Test piecewise CUDA graph with FP8 quantized model""" + + @classmethod + def setUpClass(cls): + cls.model = "nvidia/Llama-3.1-8B-Instruct-FP8" + 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=[ + "--enable-piecewise-cuda-graph", + "--quantization", + "modelopt_fp8", + "--kv-cache-dtype", + "bfloat16", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_mgsm_accuracy(self): + """Test MGSM accuracy with FP8 model""" + num_examples = 1319 + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="mgsm_en", + num_examples=num_examples, + num_threads=min(num_examples, 1024), + ) + metrics = run_eval(args) + self.assertGreaterEqual(metrics["score"], 0.85) + print(f"MGSM Accuracy: {metrics['score']:.3f}") + + if __name__ == "__main__": unittest.main()