From c001deba3796ede0150cbb5a9cf5226b44edfafb Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 26 Oct 2025 10:18:35 +0800 Subject: [PATCH] Make bmm batch invariant injection optional (#12118) --- .../srt/batch_invariant_ops/batch_invariant_ops.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) 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 7b9e3af5a..9db9cc430 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -524,7 +524,9 @@ def is_batch_invariant_mode_enabled(): return _batch_invariant_MODE -def enable_batch_invariant_mode(): +def enable_batch_invariant_mode( + enable_bmm: bool = True, +): global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm if _batch_invariant_MODE: return @@ -537,11 +539,13 @@ def enable_batch_invariant_mode(): "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA" ) _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA") - _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA") - # Also monkeypatch torch.bmm directly as a fallback - _original_torch_bmm = torch.bmm - torch.bmm = bmm_batch_invariant + if enable_bmm: + _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA") + + # Also monkeypatch torch.bmm directly as a fallback + _original_torch_bmm = torch.bmm + torch.bmm = bmm_batch_invariant def disable_batch_invariant_mode():