obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-13 09:18:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix][Torch 2.12] Fix batch_invariant test with allow_override for torch 2.12 upgrade (#40562)

Browse Source 
Signed-off-by: Lucas Kabela <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Lucas Kabela
2026-04-22 13:48:55 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 8f87eb4622
commit 22fa63cfe8
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/batch_invariant.py
+6 -2
View File
@@ -963,8 +963,12 @@ def enable_batch_invariant_mode():
_batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, "CUDA")
_batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
# Also monkeypatch torch.bmm directly as a fallback
_batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
# torch 2.12+ registers a built-in Triton bmm kernel for CUDA
# (torch._native.ops.bmm_outer_product), so we need allow_override
# to replace it at the dispatcher level.
_batch_invariant_LIB.impl(
"aten::bmm", bmm_batch_invariant, "CUDA", allow_override=True
)
_original_torch_bmm = torch.bmm
torch.bmm = bmm_batch_invariant