[Perf] Use breakable CUDA graphs (no torch.compile) by default under VLLM_BATCH_INVARIANT so the tuned matmul configs see the runtime M (#57586)
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com>
Lio Einaudi committed
Sep 24, 2026 at 02:39 UTC
711fc55c10a2bbbefcb12456c43ba5eca01be3f6
3 files changed
+118
-1
docs/features/batch_invariance.md
+3
@@ -145,6 +145,9 @@ When batch invariance is enabled, vLLM:
145
1. Uses deterministic kernel implementations for attention and other operations
146
2. Ensures consistent numerical behavior across different batch sizes
147
3. Disables certain optimizations that may introduce non-determinism (such as custom all-reduce operations in tensor parallel mode)
148
+4. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),
149
+ runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out
150
+ (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).
151
152
!!! note
153
Enabling batch invariance may impact performance compared to the default non-deterministic mode. This trade-off is intentional to guarantee reproducibility.
tests/test_config.py
+66
-1
@@ -12,6 +12,7 @@ from unittest.mock import patch
12
13
import pydantic
14
import pytest
15
+import torch
16
from huggingface_hub import ResolvedRevision
17
from pydantic import ValidationError
18
@@ -37,7 +38,7 @@ from vllm.config import (
38
WatermarkConfig,
39
update_config,
40
)
40
-from vllm.config.compilation import CompilationMode, CUDAGraphMode
41
+from vllm.config.compilation import CompilationMode, CUDAGraphMode, PassConfig
42
from vllm.config.kernel import IrOpPriorityConfig
43
from vllm.config.load import LoadConfig
44
from vllm.config.mamba import MambaBackendEnum
@@ -656,6 +657,70 @@ def test_breakable_cudagraph_platform_default(
657
default_breakable_cudagraph_architectures.cache_clear()
658
659
660
+@pytest.mark.parametrize(
661
+ "case,hidden,heads,intermediate,tp,expected",
662
+ [
663
+ ("bf16", 2048, (16, 8), 6144, 1, True),
664
+ ("bi-off", 2048, (16, 8), 6144, 1, False),
665
+ ("opt-out", 2048, (16, 8), 6144, 1, False),
666
+ ("eager", 2048, (16, 8), 6144, 1, False),
667
+ ("sp", 2048, (16, 8), 6144, 1, False),
668
+ ("fp16", 2048, (16, 8), 6144, 1, False),
669
+ ("quantized", 2048, (16, 8), 6144, 1, False),
670
+ ("untuned", 4096, (32, 32), 11008, 1, False),
671
+ ("tp4-tuned", 4096, (8, 2), 12288, 4, True),
672
+ ("list-intermediate", 2048, (16, 8), [2048, 4096], 1, False),
673
+ ("no-table", 2048, None, 6144, 1, False),
674
+ ],
675
+)
676
+def test_batch_invariant_breakable_cudagraph(
677
+ monkeypatch, case, hidden, heads, intermediate, tp, expected
678
+):
679
+ from vllm.config.vllm import default_breakable_cudagraph_architectures
680
+ from vllm.model_executor.determinism import batch_invariant_configs as bi_configs
681
+
682
+ monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0" if case == "bi-off" else "1")
683
+ monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False)
684
+ if case == "opt-out":
685
+ monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0")
686
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
687
+ monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
688
+ table = None if case == "no-table" else {(12288, 2048): None, (6144, 4096): None}
689
+ monkeypatch.setattr(current_platform, "get_device_capability", lambda: None)
690
+ monkeypatch.setattr(bi_configs, "_get_tuned_matmul_arch_family", lambda cap: "test")
691
+ monkeypatch.setattr(
692
+ bi_configs,
693
+ "_BATCH_INVARIANT_MATMUL_TUNED_CONFIGS",
694
+ {"test": table} if table else {},
695
+ )
696
+ default_breakable_cudagraph_architectures.cache_clear()
697
+ config = object.__new__(VllmConfig)
698
+ config.model_config = SimpleNamespace(
699
+ architectures=["Qwen3ForCausalLM"],
700
+ enforce_eager=case == "eager",
701
+ dtype=torch.float16 if case == "fp16" else torch.bfloat16,
702
+ quantization="fp8" if case == "quantized" else None,
703
+ hf_text_config=SimpleNamespace(intermediate_size=intermediate),
704
+ get_hidden_size=lambda: hidden,
705
+ get_head_size=lambda: 128,
706
+ get_num_attention_heads=lambda pc: heads[0],
707
+ get_num_kv_heads=lambda pc: heads[1],
708
+ )
709
+ config.parallel_config = SimpleNamespace(tensor_parallel_size=tp)
710
+ config.compilation_config = (
711
+ CompilationConfig(pass_config=PassConfig(enable_sp=True))
712
+ if case == "sp"
713
+ else CompilationConfig()
714
+ )
715
+ try:
716
+ assert config._maybe_enable_breakable_cudagraph() is expected
717
+ if expected:
718
+ assert config.compilation_config.mode == CompilationMode.NONE
719
+ finally:
720
+ os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None)
721
+ default_breakable_cudagraph_architectures.cache_clear()
722
+
723
+
724
@pytest.mark.parametrize(
725
("model_type", "expected_architecture"),
726
[
vllm/config/vllm.py
+49
@@ -776,6 +776,45 @@ class VllmConfig:
776
architectures = set(model_config.architectures)
777
return bool(architectures & default_breakable_cudagraph_architectures())
778
779
+ def _uses_breakable_cudagraph_for_batch_invariance(self) -> bool:
780
+ """Avoid freezing runtime-M tile lookup in compiled forward (#54243).
781
+ Breakable graphs look up tuned bf16, unquantized qkv/o/gate_up/down tiles
782
+ at capture; lm_head runs outside compiled forward and does not benefit."""
783
+ from vllm.model_executor.determinism import batch_invariant_configs as bi
784
+ from vllm.platforms import current_platform
785
+
786
+ model = self.model_config
787
+ if (
788
+ not envs.VLLM_BATCH_INVARIANT
789
+ or model is None
790
+ or model.enforce_eager
791
+ or model.dtype != torch.bfloat16
792
+ or model.quantization is not None
793
+ or not current_platform.is_cuda()
794
+ # Sequence parallelism / async TP are torch.compile passes.
795
+ or self.compilation_config.pass_config.enable_sp
796
+ or self.compilation_config.pass_config.fuse_gemm_comms
797
+ ):
798
+ return False
799
+ family = bi._get_tuned_matmul_arch_family(
800
+ current_platform.get_device_capability()
801
+ )
802
+ if family is None or family not in bi._BATCH_INVARIANT_MATMUL_TUNED_CONFIGS:
803
+ return False
804
+ table = bi._BATCH_INVARIANT_MATMUL_TUNED_CONFIGS[family]
805
+ parallel = self.parallel_config
806
+ tp = parallel.tensor_parallel_size
807
+ hidden = model.get_hidden_size()
808
+ head = model.get_head_size()
809
+ heads = model.get_num_attention_heads(parallel)
810
+ kv_heads = model.get_num_kv_heads(parallel)
811
+ shapes = [((heads + 2 * kv_heads) * head, hidden), (hidden, heads * head)]
812
+ intermediate = getattr(model.hf_text_config, "intermediate_size", None)
813
+ # Per-layer sizes (e.g. Gemma3n) are not modeled.
814
+ if isinstance(intermediate, int):
815
+ shapes += [(2 * intermediate // tp, hidden), (hidden, intermediate // tp)]
816
+ return any(shape in table for shape in shapes)
817
+
818
def _maybe_enable_breakable_cudagraph(self) -> bool:
819
if (
820
"VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
@@ -786,6 +825,16 @@ class VllmConfig:
825
"Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
826
"Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
827
)
828
+ elif (
829
+ "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ
830
+ and envs.VLLM_BATCH_INVARIANT
831
+ and self._uses_breakable_cudagraph_for_batch_invariance()
832
+ ):
833
+ os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
834
+ logger.info_once(
835
+ "VLLM_BATCH_INVARIANT=1: auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. "
836
+ "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out."
837
+ )
838
839
from vllm.compilation.breakable_cudagraph import (
840
is_breakable_cudagraph_enabled,