From 22fa63cfe849c38bbdb590163043230a1a1a6148 Mon Sep 17 00:00:00 2001 From: Lucas Kabela Date: Wed, 22 Apr 2026 13:48:55 -0700 Subject: [PATCH] [Bugfix][Torch 2.12] Fix batch_invariant test with allow_override for torch 2.12 upgrade (#40562) Signed-off-by: Lucas Kabela Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/layers/batch_invariant.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 152333beec..08756ee04d 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -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