Support FP8 Per Token Quant Piecewise (#13272)

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)

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