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 5, 16:00 UTC
[Bugfix] Enable batch-invariant Triton matmul on all Ampere GPUs (SM 8x) (#38427)
Signed-off-by: yusuf <yusufmohammad@live.com> Signed-off-by: yusuf <yusuf@deeplearningmachine.mynet> Signed-off-by: Yusuf Mohammad <79484377+YM2132@users.noreply.github.com> Signed-off-by: <> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: yusuf <yusuf@deeplearningmachine.mynet>
Yusuf Mohammad committed
Apr 2, 2026 at 14:29 UTC
16a65e41736c5e8d27e9e843668f6e3f99d68d9a
1 file changed
+3
-5
vllm/model_executor/layers/batch_invariant.py
+3
-5
@@ -935,11 +935,9 @@ def enable_batch_invariant_mode():
935
_batch_invariant_MODE = True
936
_batch_invariant_LIB = torch.library.Library("aten", "IMPL")
937
938
- if (
939
- current_platform.is_device_capability_family(100)
940
- or current_platform.is_device_capability(80)
941
- or current_platform.is_device_capability(89)
942
- ):
938
+ if current_platform.is_device_capability_family(
939
+ 100
940
+ ) or current_platform.is_device_capability_family(80):
941
# For PyTorch 2.9, B200 uses GEMV for bs=1
942
# Requires https://github.com/pytorch/pytorch/pull/166735
943
_batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")