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

Make bmm batch invariant injection optional (#12118)

Browse Source
This commit is contained in:
fzyzcjy
2025-10-26 10:18:35 +08:00
committed by GitHub
parent b4d2da106e
commit c001deba37
1 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py
View File

@@ -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():