Batch invariant torch.compile

Author: PaulZhang12

vllm/model_executor/layers/batch_invariant.py

@@ -11,6 +11,7 @@
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import is_torch_equal_or_newer
 
 logger = init_logger(__name__)
 
@@ -741,8 +742,6 @@ def disable_batch_invariant_mode():
     _batch_invariant_MODE = False
     _batch_invariant_LIB = None
 
-    torch.backends.cuda.preferred_blas_library(backend="default")
-
 
 @contextlib.contextmanager
 def set_batch_invariant_mode(enabled: bool = True):