chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Piecewise CUDA Graph] Support ModelOpt FP8 (#13094)

Browse Source
This commit is contained in:
b8zhong
2025-11-17 20:46:24 -08:00
committed by GitHub
parent 4ce8fb3cc2
commit f33860777c
4 changed files with 48 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -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)

8
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -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")