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, 13:45 UTC
[BE][Torch 2.12] Remove workaround code for fixed cublas issue (#40845)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com> Signed-off-by: Lucas Kabela <lucasakabela@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Lucas Kabela committed
Apr 28, 2026 at 21:07 UTC
fa1b9840f6d87ef6e3b247a78514ccc1d6e5f1ce
1 file changed
+11
-11
vllm/model_executor/layers/batch_invariant.py
+11
-11
@@ -930,22 +930,17 @@ def enable_batch_invariant_mode():
930
_batch_invariant_MODE = True
931
_batch_invariant_LIB = torch.library.Library("aten", "IMPL")
932
933
- if current_platform.is_device_capability_family(
934
- 100
935
- ) or current_platform.is_device_capability_family(80):
936
- # For PyTorch 2.9, B200 uses GEMV for bs=1
937
- # Requires https://github.com/pytorch/pytorch/pull/166735
933
+ if current_platform.is_device_capability_family(80):
934
+ # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
935
+ # triton persistent matmul overrides for mm/addmm/matmul/linear.
936
_batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
937
_batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
938
_batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
939
_batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
942
-
943
- # Query the shared memory size and set block size
944
- # accordingly to avoid triton OutOfResources
945
- _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
940
else:
947
- # Only source of batch invariance for Hopper is split-k, can disable through
948
- # cuBLAS workspace config
941
+ # Hopper (SM90) and Blackwell (SM100): the only source of batch
942
+ # variance is split-k, which we disable via the cuBLAS workspace
943
+ # config.
944
_original_cublas_workspace_cfg = os.environ.get("CUBLAS_WORKSPACE_CONFIG", None)
945
_original_cublaslt_workspace_size = os.environ.get(
946
"CUBLASLT_WORKSPACE_SIZE", None
@@ -953,6 +948,11 @@ def enable_batch_invariant_mode():
948
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
949
os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
950
951
+ # Triton bmm/persistent-matmul kernels read this for the FP16 N-tile size;
952
+ # set unconditionally because bmm is overridden on all CUDA platforms.
953
+ if current_platform.is_cuda():
954
+ _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
955
+
956
_batch_invariant_LIB.impl(
957
"aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
958
)