Wider coverage of batch invariance torch.compile, cuda graphs, and B200

Author: PaulZhang12

tests/v1/generation/test_batch_invariance.py

@@ -456,7 +456,6 @@ def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):
         model=model,
         max_num_seqs=1,
         tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
-        enforce_eager=True,
         gpu_memory_utilization=0.9,
         max_model_len=2048,
         dtype="bfloat16",
@@ -998,7 +997,6 @@ def LLM_with_max_seqs(
         dtype="bfloat16",
         tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
         enable_prefix_caching=False,
-        enforce_eager=True,
         # Enable for MOE models
         # enable_expert_parallel=True,
     )

vllm/config/compilation.py

@@ -15,6 +15,9 @@
 from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
 from vllm.config.utils import config
 from vllm.logger import init_logger
+from vllm.model_executor.layers.batch_invariant import (
+    vllm_is_batch_invariant,
+)
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.torch_utils import is_torch_equal_or_newer

vllm/model_executor/layers/batch_invariant.py

@@ -733,14 +733,17 @@ def enable_batch_invariant_mode():
     global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
     if _batch_invariant_MODE:
         return
-
+    
     _batch_invariant_MODE = True
     _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
-    _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::bmm", bmm_batch_invariant, "CUDA")
-    _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
+    
+    # Batch invariant matmuls are no longer needed after cublas overrides
+    if not is_torch_equal_or_newer("2.10.0.dev"):
+        _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")
+
     _batch_invariant_LIB.impl(
         "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
     )
@@ -749,6 +752,7 @@ def enable_batch_invariant_mode():
     _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
 
     # Also monkeypatch torch.bmm directly as a fallback
+    _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
     _original_torch_bmm = torch.bmm
     torch.bmm = bmm_batch_invariant
 

vllm/utils/flashinfer.py

@@ -20,6 +20,9 @@
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm.model_executor.layers.batch_invariant import (
+    vllm_is_batch_invariant,
+)
 
 logger = init_logger(__name__)
 
@@ -213,6 +216,9 @@ def _force_use_trtllm_attention(env_value: bool | None) -> bool | None:
     """Cache the env value for VLLM_USE_TRTLLM_ATTENTION"""
     if env_value is not None:
         logger.info_once("VLLM_USE_TRTLLM_ATTENTION is set to %s", env_value)
+    if vllm_is_batch_invariant():
+        logger.info_once("VLLM_USE_TRTLLM_ATTENTION is disabled for batch-invariant")
+        return False
     return env_value