A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 18:45 UTC
[Bugfix][Torch 2.12] Fix batch_invariant test with allow_override for torch 2.12 upgrade (#40562)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lucas Kabela committed
Apr 22, 2026 at 13:48 UTC
22fa63cfe849c38bbdb590163043230a1a1a6148
1 file changed
+6
-2
vllm/model_executor/layers/batch_invariant.py
+6
-2
@@ -963,8 +963,12 @@ def enable_batch_invariant_mode():
963
_batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, "CUDA")
964
_batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
965
966
- # Also monkeypatch torch.bmm directly as a fallback
967
- _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
966
+ # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
967
+ # (torch._native.ops.bmm_outer_product), so we need allow_override
968
+ # to replace it at the dispatcher level.
969
+ _batch_invariant_LIB.impl(
970
+ "aten::bmm", bmm_batch_invariant, "CUDA", allow_override=True
971
+ )
972
_original_torch_bmm = torch.bmm
973
torch.bmm = bmm_batch_invariant
974