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

add git gemm warpper for dispatch_bf16_fp32_backend (#25860)

Browse Source
This commit is contained in:
BingjiaWang
2026-05-20 15:24:10 -07:00
committed by GitHub
parent 371b6c9ea0
commit ce7141ef98
1 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/deepseek_v4.py
+7 -8
View File
@@ -21,9 +21,13 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _use_aiter: if _use_aiter:
from aiter.tuned_gemm import tgemm from aiter.tuned_gemm import tgemm
from sglang.srt.layers import deep_gemm_wrapper
if TYPE_CHECKING: if TYPE_CHECKING:
from tvm_ffi.module import Module from tvm_ffi.module import Module
_linear_bf16_fp32_algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get()
def make_name(name: str) -> str: def make_name(name: str) -> str:
return f"dpsk_v4_{name}" return f"dpsk_v4_{name}"
@@ -1013,10 +1017,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: def linear_bf16_fp32(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
from sglang.srt.environ import envs return _dispatch_bf16_fp32_backend(x, y, algo=_linear_bf16_fp32_algo)
algo = envs.SGLANG_OPT_BF16_FP32_GEMM_ALGO.get()
return _dispatch_bf16_fp32_backend(x, y, algo=algo)
def _dispatch_bf16_fp32_backend( def _dispatch_bf16_fp32_backend(
@@ -1026,10 +1027,8 @@ def _dispatch_bf16_fp32_backend(
module = _jit_torch_cublas_bf16_fp32() module = _jit_torch_cublas_bf16_fp32()
return module.linear_bf16_fp32(x, y) return module.linear_bf16_fp32(x, y)
elif algo == "deep_gemm": elif algo == "deep_gemm":
import deep_gemm z = torch.empty(x.size(0), y.size(0), dtype=torch.float32, device=x.device)
deep_gemm_wrapper.gemm_nt_bf16bf16f32(x, y, z)
z = x.new_empty(x.size(0), y.size(0), dtype=torch.float32)
deep_gemm.bf16_gemm_nt(x, y, z)
return z return z
elif _use_aiter: elif _use_aiter:
return tgemm.mm(x, y, otype=torch.float32) return tgemm.mm(x, y, otype=torch.float32)