Make bmm batch invariant injection optional (#12118)

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,6 +539,8 @@ def enable_batch_invariant_mode():
         "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
     )
     _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
+
+    if enable_bmm:
         _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
 
         # Also monkeypatch torch.bmm directly as a fallback