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

Refactor DeepGEMM integration (#7150)

Browse Source
This commit is contained in:
fzyzcjy
2025-06-14 11:41:03 +08:00
committed by GitHub
parent 8b8f2e7463
commit b4c41f7276
12 changed files with 207 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

@@ -23,7 +23,8 @@ import torch
import triton
import triton.language as tl
from sglang.srt.layers.quantization.deep_gemm import _ENABLE_JIT_DEEPGEMM
from sglang.math_utils import align
from sglang.srt.layers.quantization import deep_gemm_wrapper
from sglang.srt.utils import (
direct_register_custom_op,
get_device_core_count,
@@ -44,10 +45,6 @@ if _is_cuda:
sgl_per_token_quant_fp8,
)
from sglang.srt.layers.quantization.deep_gemm import (
gemm_nt_f8f8bf16 as deep_gemm_gemm_nt_f8f8bf16,
)
logger = logging.getLogger(__name__)
@@ -67,7 +64,6 @@ else:
fp8_max = torch.finfo(fp8_dtype).max
fp8_min = -fp8_max
if supports_custom_op():
def deep_gemm_fp8_fp8_bf16_nt(
@@ -77,7 +73,7 @@ if supports_custom_op():
Bs: torch.Tensor,
C: torch.Tensor,
) -> None:
deep_gemm_gemm_nt_f8f8bf16((A, As), (B, Bs), C)
deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C)
def deep_gemm_fp8_fp8_bf16_nt_fake(
A: torch.Tensor,
@@ -797,12 +793,12 @@ def w8a8_block_fp8_matmul_deepgemm(
M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size, output_dtype)
# Deepgemm only supports output tensor type as bfloat16
assert C.dtype == torch.bfloat16 and _ENABLE_JIT_DEEPGEMM
assert C.dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
if supports_custom_op():
torch.ops.sglang.deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C)
else:
deep_gemm_gemm_nt_f8f8bf16((A, As), (B, Bs), C)
deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C)
return C
@@ -896,7 +892,7 @@ def w8a8_block_fp8_matmul(
block_size: List[int],
output_dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
if output_dtype == torch.bfloat16 and _ENABLE_JIT_DEEPGEMM:
if output_dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM:
return w8a8_block_fp8_matmul_deepgemm(
A, B, As, Bs, block_size, output_dtype=output_dtype
)