Skip to content

[BE][Torch 2.12] Remove workaround code for fixed cublas issue #40845

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vllm-bot merged 5 commits into from
Apr 29, 2026
Merged

[BE][Torch 2.12] Remove workaround code for fixed cublas issue #40845

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
092337b
Remove workaround code for fixed cublas issue
Lucaskabela Apr 24, 2026
980b12a
Merge branch 'main' into lucaskabela/fix_b200_batch_invariance_hack
Lucaskabela Apr 27, 2026
d92565b
Update vllm/model_executor/layers/batch_invariant.py
Lucaskabela Apr 27, 2026
844f4a7
Merge branch 'main' into lucaskabela/fix_b200_batch_invariance_hack
yewentao256 Apr 28, 2026
232b87b
Merge branch 'main' into lucaskabela/fix_b200_batch_invariance_hack
Lucaskabela Apr 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions vllm/model_executor/layers/batch_invariant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -930,29 +930,29 @@ def enable_batch_invariant_mode():
_batch_invariant_MODE = True
_batch_invariant_LIB = torch.library.Library("aten", "IMPL")

if current_platform.is_device_capability_family(
100
) or current_platform.is_device_capability_family(80):
# For PyTorch 2.9, B200 uses GEMV for bs=1
# Requires https://github.com/pytorch/pytorch/pull/166735
if current_platform.is_device_capability_family(80):
# SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
# triton persistent matmul overrides for mm/addmm/matmul/linear.
_batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
_batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
_batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
_batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")

# Query the shared memory size and set block size
# accordingly to avoid triton OutOfResources
_fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
else:
# Only source of batch invariance for Hopper is split-k, can disable through
# cuBLAS workspace config
# Hopper (SM90) and Blackwell (SM100): the only source of batch
# variance is split-k, which we disable via the cuBLAS workspace
# config.
_original_cublas_workspace_cfg = os.environ.get("CUBLAS_WORKSPACE_CONFIG", None)
_original_cublaslt_workspace_size = os.environ.get(
"CUBLASLT_WORKSPACE_SIZE", None
)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"

# Triton bmm/persistent-matmul kernels read this for the FP16 N-tile size;
# set unconditionally because bmm is overridden on all CUDA platforms.
if current_platform.is_cuda():
_fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128

_batch_invariant_LIB.impl(
"aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
)
Expand Down
Loading