[compilation] allow torch.compile with batch invariant inference (#1570)

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/config/model.py

@@ -17,7 +17,6 @@
 from aphrodite.config.scheduler import RunnerType
 from aphrodite.config.utils import assert_hashable, config, getattr_iter
 from aphrodite.logger import init_logger
-from aphrodite.modeling.layers.batch_invariant import aphrodite_is_batch_invariant
 from aphrodite.platforms import current_platform
 from aphrodite.transformers_utils.config import (
     ConfigFormat,
@@ -424,10 +423,6 @@ def __post_init__(
         skip_mm_profiling: bool | None,
         video_pruning_rate: float | None,
     ) -> None:
-        # Enable batch invariance settings if requested
-        if aphrodite_is_batch_invariant():
-            self.enforce_eager = True
-
         # Set the default seed to 0 in V1.
         # NOTE(woosuk): In V1, we use separate processes for workers (unless
         # APHRODITE_ENABLE_V1_MULTIPROCESSING=0), so setting a seed here

aphrodite/envs.py

@@ -254,11 +254,12 @@ def disable_compile_cache() -> bool:
 
 
 def use_aot_compile() -> bool:
+    from aphrodite.modeling.layers.batch_invariant import aphrodite_is_batch_invariant
     from aphrodite.utils.torch_utils import is_torch_equal_or_newer
 
     default_value = "1" if is_torch_equal_or_newer("2.10.0.dev") and not disable_compile_cache() else "0"
 
-    return os.environ.get("APHRODITE_USE_AOT_COMPILE", default_value) == "1"
+    return not aphrodite_is_batch_invariant() and os.environ.get("APHRODITE_USE_AOT_COMPILE", default_value) == "1"
 
 
 def env_with_choices(

aphrodite/modeling/layers/batch_invariant.py

@@ -9,6 +9,7 @@
 import aphrodite.envs as envs
 from aphrodite.logger import init_logger
 from aphrodite.triton_utils import tl, triton
+from aphrodite.utils.torch_utils import is_torch_equal_or_newer
 
 logger = init_logger(__name__)
 
@@ -678,6 +679,10 @@ def linear_batch_invariant(input, weight, bias=None):
 _batch_invariant_MODE = False
 _batch_invariant_LIB = None
 _original_torch_bmm = None
+_original_fp16_reduction_precision = None
+_original_bf16_reduction_precision = None
+_original_cublas_workspace_cfg = None
+_original_cublaslt_workspace_size = None
 
 
 def is_batch_invariant_mode_enabled():
@@ -686,6 +691,8 @@ def is_batch_invariant_mode_enabled():
 
 def enable_batch_invariant_mode():
     global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
+    global _original_fp16_reduction_precision, _original_bf16_reduction_precision
+    global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
     if _batch_invariant_MODE:
         return
 
@@ -705,14 +712,59 @@ def enable_batch_invariant_mode():
     _original_torch_bmm = torch.bmm
     torch.bmm = bmm_batch_invariant
 
+    _original_bf16_reduction_precision = torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction
+    _original_fp16_reduction_precision = torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction
+
+    reduced_precision_val = (False, False) if is_torch_equal_or_newer("2.10.0.dev") else False
+    torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = reduced_precision_val
+    torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = reduced_precision_val
+    torch.backends.cuda.preferred_blas_library(backend="cublaslt")
+
+    if not is_torch_equal_or_newer("2.10.0.dev"):
+        _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"
+
 
 def disable_batch_invariant_mode():
     global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
+    global _original_fp16_reduction_precision, _original_bf16_reduction_precision
+    global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
+    if not _batch_invariant_MODE:
+        return
+
     if _batch_invariant_LIB is not None:
         _batch_invariant_LIB._destroy()
     if _original_torch_bmm is not None:
         torch.bmm = _original_torch_bmm
         _original_torch_bmm = None
+
+    if _original_bf16_reduction_precision is not None:
+        torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = _original_bf16_reduction_precision
+        _original_bf16_reduction_precision = None
+    if _original_fp16_reduction_precision is not None:
+        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = _original_fp16_reduction_precision
+        _original_fp16_reduction_precision = None
+
+    torch.backends.cuda.preferred_blas_library(backend="default")
+
+    if not is_torch_equal_or_newer("2.10.0.dev"):
+        # Set cublas env vars to previous results. If previous results are None,
+        # that means the env vars were not set, so we should remove them.
+        if _original_cublas_workspace_cfg:
+            os.environ["CUBLAS_WORKSPACE_CONFIG"] = _original_cublas_workspace_cfg
+        elif "CUBLAS_WORKSPACE_CONFIG" in os.environ:
+            del os.environ["CUBLAS_WORKSPACE_CONFIG"]
+
+        if _original_cublaslt_workspace_size:
+            os.environ["CUBLASLT_WORKSPACE_SIZE"] = _original_cublaslt_workspace_size
+        elif "CUBLASLT_WORKSPACE_SIZE" in os.environ:
+            del os.environ["CUBLASLT_WORKSPACE_SIZE"]
+
+    _original_cublas_workspace_cfg = None
+    _original_cublaslt_workspace_size = None
+
     _batch_invariant_MODE = False
     _batch_invariant_LIB = None
 
@@ -791,6 +843,9 @@ def override_envs_for_invariance():
     os.environ["NCCL_NTHREADS"] = "1"
     os.environ["NCCL_SOCKET_NTHREADS"] = "1"
 
+    # torch.compile settings
+    os.environ["APHRODITE_USE_AOT_COMPILE"] = "0"
+
 
 def init_batch_invariance():
     # this will hit all the csrc overrides as well

aphrodite/quantization/fp8.py

@@ -319,6 +319,7 @@ def __init__(self, quant_config: Fp8Config):
             self.use_marlin = False
 
         self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
+        self.use_deep_gemm = is_deep_gemm_supported()
 
         self.weight_block_size = self.quant_config.weight_block_size
         self.block_quant = self.weight_block_size is not None
@@ -493,7 +494,11 @@ def apply(
         # if batch invariant mode is enabled, prefer DeepGEMM FP8 path
         # we will use BF16 dequant when DeepGEMM is not supported.
         if aphrodite_is_batch_invariant():
-            if self.block_quant and should_use_deepgemm_for_fp8_linear(torch.bfloat16, layer.weight, None):
+            # Call is_deep_gemm_supported() ahead of time for torch.compile
+            # dynamo has trouble tracing through
+            if self.block_quant and should_use_deepgemm_for_fp8_linear(
+                torch.bfloat16, layer.weight, self.use_deep_gemm
+            ):
                 # use group quant consistent with block size across K
                 assert self.act_q_group_shape is not None
                 q_input, input_scale = QuantFP8(