[Kernel] Deprecate DeepGemm in sgl kernel and apply custom wheel sgl-deep-gemm (#24268)

This commit is contained in:
Baizhou Zhang
2026-05-06 18:59:01 -07:00
committed by GitHub
parent eaf074d50e
commit ecb786c8d7
13 changed files with 87 additions and 19 deletions
+6
View File
@@ -294,6 +294,12 @@ install_sglang_kernel() {
else
echo "CUSTOM_BUILD_SGL_KERNEL=true: keeping freshly built sgl-kernel wheel."
fi
SGL_DEEP_GEMM_VERSION=$(grep -Po -m1 '(?<=sgl-deep-gemm==)[0-9A-Za-z\.\-]+' python/pyproject.toml)
if [ "$CU_MAJOR" = "13" ]; then
$PIP_CMD install "sgl-deep-gemm==${SGL_DEEP_GEMM_VERSION}" --force-reinstall $PIP_INSTALL_SUFFIX
else
$PIP_CMD install "https://github.com/sgl-project/whl/releases/download/v${SGL_DEEP_GEMM_VERSION}/sgl_deep_gemm-${SGL_DEEP_GEMM_VERSION}+cu129-py3-none-manylinux2014_$(uname -m).whl" --force-reinstall $PIP_INSTALL_SUFFIX
fi
mark_step_done "${FUNCNAME[0]}"
}
+8 -4
View File
@@ -235,8 +235,11 @@ def compile_one_shape(kernel_type, n, k, num_groups, m_list):
)
m_list = [m for m in m_list if m <= max_m]
old_mode = deep_gemm.get_compile_mode()
deep_gemm.set_compile_mode(1)
get_compile_mode = getattr(deep_gemm, "get_compile_mode", None)
set_compile_mode = getattr(deep_gemm, "set_compile_mode", None)
old_mode = get_compile_mode() if get_compile_mode is not None else None
if set_compile_mode is not None:
set_compile_mode(1)
try:
if kernel_type == "NORMAL":
lhs_q, lhs_s = _empty_token_fp8((max_m, k))
@@ -255,7 +258,7 @@ def compile_one_shape(kernel_type, n, k, num_groups, m_list):
(lhs_q[:m], lhs_s[:m]),
(rhs_q, rhs_s),
out[:m],
m_indices=m_indices[:m],
m_indices[:m],
)
elif kernel_type == "MASKED":
@@ -274,7 +277,8 @@ def compile_one_shape(kernel_type, n, k, num_groups, m_list):
expected_m=m,
)
finally:
deep_gemm.set_compile_mode(old_mode)
if set_compile_mode is not None and old_mode is not None:
set_compile_mode(old_mode)
torch.cuda.current_stream().synchronize()
torch.cuda.empty_cache()