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

Support fast gemm when in batch invariant DeepGEMM fallback (#13259)

Browse Source
This commit is contained in:
fzyzcjy
2025-11-15 16:34:15 +08:00
committed by GitHub
parent f0b5ccf5f5
commit 2fbc78a083
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py
+7
View File
@@ -18,6 +18,10 @@ if ENABLE_JIT_DEEPGEMM:
_ENABLE_MM_DEEPGEMM = get_bool_env_var( _ENABLE_MM_DEEPGEMM = get_bool_env_var(
"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM", "1" "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM", "1"
) )
# If true, allows to fallback to batch variant gemm when the shape cannot be run in DeepGEMM
_ENABLE_MM_FALLBACK_VARIANT = get_bool_env_var(
"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT", "0"
)
_ENABLE_MM_COMPARISON_TEST = get_bool_env_var( _ENABLE_MM_COMPARISON_TEST = get_bool_env_var(
"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_COMPARISON_TEST" "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_COMPARISON_TEST"
) )
@@ -292,6 +296,9 @@ def matmul_persistent(
return _matmul_persistent_deepgemm(a=a, b=b, bias=bias) return _matmul_persistent_deepgemm(a=a, b=b, bias=bias)
if _ENABLE_MM_FALLBACK_VARIANT:
return torch.einsum("ik,kj->ij", a, b)
return _matmul_persistent_triton(a=a, b=b, bias=bias) return _matmul_persistent_triton(a=a, b=b, bias=bias)