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 11, 05:00 UTC
[Bug] Fix batch invariant in torch 2.10 (#30907)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Wentao Ye committed
Dec 18, 2025 at 10:27 UTC
66287582339d1f27c951f0108aac9aa7377fe643
1 file changed
+20
-24
vllm/model_executor/layers/batch_invariant.py
+20
-24
@@ -933,30 +933,26 @@ def enable_batch_invariant_mode():
933
_batch_invariant_MODE = True
934
_batch_invariant_LIB = torch.library.Library("aten", "IMPL")
935
936
- # Batch invariant matmuls are no longer needed after cublas overrides
937
- if not is_torch_equal_or_newer("2.10.0.dev"):
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
- ):
943
- # For PyTorch 2.9, B200 uses GEMV for bs=1
944
- # Requires https://github.com/pytorch/pytorch/pull/166735
945
- _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
946
- _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
947
- _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
948
- _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
949
- else:
950
- # Only source of batch invariance for Hopper is split-k, can disable through
951
- # cuBLAS workspace config
952
- _original_cublas_workspace_cfg = os.environ.get(
953
- "CUBLAS_WORKSPACE_CONFIG", None
954
- )
955
- _original_cublaslt_workspace_size = os.environ.get(
956
- "CUBLASLT_WORKSPACE_SIZE", None
957
- )
958
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
959
- os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
936
+ if (
937
+ current_platform.is_device_capability_family(100)
938
+ or current_platform.is_device_capability(80)
939
+ or current_platform.is_device_capability(89)
940
+ ):
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")
944
+ _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
945
+ _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
946
+ _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
947
+ else:
948
+ # Only source of batch invariance for Hopper is split-k, can disable through
949
+ # cuBLAS workspace config
950
+ _original_cublas_workspace_cfg = os.environ.get("CUBLAS_WORKSPACE_CONFIG", None)
951
+ _original_cublaslt_workspace_size = os.environ.get(
952
+ "CUBLASLT_WORKSPACE_SIZE", None
953
+ )
954
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
955
+ os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
956
957
_batch_invariant_LIB.impl(
958
"aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"