From 2fbc78a08314d6f3c4a054464eb1c8fe383a1595 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 15 Nov 2025 16:34:15 +0800 Subject: [PATCH] Support fast gemm when in batch invariant DeepGEMM fallback (#13259) --- .../sglang/srt/batch_invariant_ops/batch_invariant_ops.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py index c01f8bb33..67bb1c5c7 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -18,6 +18,10 @@ if ENABLE_JIT_DEEPGEMM: _ENABLE_MM_DEEPGEMM = get_bool_env_var( "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( "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) + if _ENABLE_MM_FALLBACK_VARIANT: + return torch.einsum("ik,kj->ij", a, b) + return _matmul_persistent_triton(a=a, b=b, bias=bias)