From 0997c78d2c99110ef313008c57dfd9096c9bcc4b Mon Sep 17 00:00:00 2001 From: Stefan He Date: Fri, 14 Nov 2025 11:40:12 -0800 Subject: [PATCH] Support FP8 Per Token Quant Piecewise (#13272) --- python/sglang/srt/layers/quantization/fp8.py | 7 +++++++ python/sglang/srt/layers/quantization/fp8_kernel.py | 4 ++++ 2 files changed, 11 insertions(+) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 14d8671e9..7b4870c97 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -1322,3 +1322,10 @@ 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_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 09174d0e4..2be8a8cf5 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/python/sglang/srt/layers/quantization/fp8_kernel.py @@ -1849,3 +1849,7 @@ if _is_cuda: input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0 ): return + + @torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8") + def _(input, output_q, output_s): + return