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 13, 02:00 UTC
[Refactor] Remove Unused Func in Batch Invariant (#28881)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Nov 17, 2025 at 23:29 UTC
3ddcf4601171797b6e63eda6b5956136441b3408
1 file changed
-73
vllm/model_executor/layers/batch_invariant.py
-73
@@ -1,8 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import contextlib
3
import os
5
-from collections import namedtuple
4
from collections.abc import Callable
5
from functools import cache
6
from typing import Any
@@ -725,10 +723,6 @@ _original_cublas_workspace_cfg = None
723
_original_cublaslt_workspace_size = None
724
725
728
-def is_batch_invariant_mode_enabled():
729
- return _batch_invariant_MODE
730
-
731
-
726
def enable_batch_invariant_mode():
727
global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
728
global _original_fp16_reduction_precision, _original_bf16_reduction_precision
@@ -791,73 +785,6 @@ def enable_batch_invariant_mode():
785
torch.backends.cuda.preferred_blas_library(backend="cublaslt")
786
787
794
-def disable_batch_invariant_mode():
795
- global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
796
- global _original_fp16_reduction_precision, _original_bf16_reduction_precision
797
- global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
798
- if not _batch_invariant_MODE:
799
- return
800
-
801
- if _batch_invariant_LIB is not None:
802
- _batch_invariant_LIB._destroy()
803
- if _original_torch_bmm is not None:
804
- torch.bmm = _original_torch_bmm
805
- _original_torch_bmm = None
806
-
807
- if _original_bf16_reduction_precision is not None:
808
- torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (
809
- _original_bf16_reduction_precision
810
- )
811
- _original_bf16_reduction_precision = None
812
- if _original_fp16_reduction_precision is not None:
813
- torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = (
814
- _original_fp16_reduction_precision
815
- )
816
- _original_fp16_reduction_precision = None
817
-
818
- torch.backends.cuda.preferred_blas_library(backend="default")
819
-
820
- if not is_torch_equal_or_newer("2.10.0.dev"):
821
- # Set cublas env vars to previous results. If previous results are None,
822
- # that means the env vars were not set, so we should remove them.
823
- if _original_cublas_workspace_cfg:
824
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = _original_cublas_workspace_cfg
825
- elif "CUBLAS_WORKSPACE_CONFIG" in os.environ:
826
- del os.environ["CUBLAS_WORKSPACE_CONFIG"]
827
-
828
- if _original_cublaslt_workspace_size:
829
- os.environ["CUBLASLT_WORKSPACE_SIZE"] = _original_cublaslt_workspace_size
830
- elif "CUBLASLT_WORKSPACE_SIZE" in os.environ:
831
- del os.environ["CUBLASLT_WORKSPACE_SIZE"]
832
-
833
- _original_cublas_workspace_cfg = None
834
- _original_cublaslt_workspace_size = None
835
-
836
- _batch_invariant_MODE = False
837
- _batch_invariant_LIB = None
838
-
839
-
840
-@contextlib.contextmanager
841
-def set_batch_invariant_mode(enabled: bool = True):
842
- global _batch_invariant_MODE, _batch_invariant_LIB
843
- old_data = (_batch_invariant_MODE, _batch_invariant_LIB)
844
- if enabled:
845
- enable_batch_invariant_mode()
846
- else:
847
- disable_batch_invariant_mode()
848
- yield
849
- if _batch_invariant_LIB is not None:
850
- _batch_invariant_LIB._destroy()
851
- _batch_invariant_MODE, _batch_invariant_LIB = old_data
852
-
853
-
854
-AttentionBlockSize = namedtuple("AttentionBlockSize", ["block_m", "block_n"])
855
-
856
-
857
-def get_batch_invariant_attention_block_size() -> AttentionBlockSize:
858
- return AttentionBlockSize(block_m=16, block_n=16)
859
-
860
-
788
@cache
789
def vllm_is_batch_invariant():
790
env_key = "VLLM_BATCH_INVARIANT"