diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 789878f48a..7353bd3594 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -145,6 +145,9 @@ When batch invariance is enabled, vLLM: 1. Uses deterministic kernel implementations for attention and other operations 2. Ensures consistent numerical behavior across different batch sizes 3. Disables certain optimizations that may introduce non-determinism (such as custom all-reduce operations in tensor parallel mode) +4. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell), + 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 + (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes). !!! note Enabling batch invariance may impact performance compared to the default non-deterministic mode. This trade-off is intentional to guarantee reproducibility. diff --git a/tests/test_config.py b/tests/test_config.py index 02639db26e..0b8a10504c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -12,6 +12,7 @@ from unittest.mock import patch import pydantic import pytest +import torch from huggingface_hub import ResolvedRevision from pydantic import ValidationError @@ -37,7 +38,7 @@ from vllm.config import ( WatermarkConfig, update_config, ) -from vllm.config.compilation import CompilationMode, CUDAGraphMode +from vllm.config.compilation import CompilationMode, CUDAGraphMode, PassConfig from vllm.config.kernel import IrOpPriorityConfig from vllm.config.load import LoadConfig from vllm.config.mamba import MambaBackendEnum @@ -656,6 +657,70 @@ def test_breakable_cudagraph_platform_default( default_breakable_cudagraph_architectures.cache_clear() +@pytest.mark.parametrize( + "case,hidden,heads,intermediate,tp,expected", + [ + ("bf16", 2048, (16, 8), 6144, 1, True), + ("bi-off", 2048, (16, 8), 6144, 1, False), + ("opt-out", 2048, (16, 8), 6144, 1, False), + ("eager", 2048, (16, 8), 6144, 1, False), + ("sp", 2048, (16, 8), 6144, 1, False), + ("fp16", 2048, (16, 8), 6144, 1, False), + ("quantized", 2048, (16, 8), 6144, 1, False), + ("untuned", 4096, (32, 32), 11008, 1, False), + ("tp4-tuned", 4096, (8, 2), 12288, 4, True), + ("list-intermediate", 2048, (16, 8), [2048, 4096], 1, False), + ("no-table", 2048, None, 6144, 1, False), + ], +) +def test_batch_invariant_breakable_cudagraph( + monkeypatch, case, hidden, heads, intermediate, tp, expected +): + from vllm.config.vllm import default_breakable_cudagraph_architectures + from vllm.model_executor.determinism import batch_invariant_configs as bi_configs + + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0" if case == "bi-off" else "1") + monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False) + if case == "opt-out": + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "0") + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "is_rocm", lambda: False) + table = None if case == "no-table" else {(12288, 2048): None, (6144, 4096): None} + monkeypatch.setattr(current_platform, "get_device_capability", lambda: None) + monkeypatch.setattr(bi_configs, "_get_tuned_matmul_arch_family", lambda cap: "test") + monkeypatch.setattr( + bi_configs, + "_BATCH_INVARIANT_MATMUL_TUNED_CONFIGS", + {"test": table} if table else {}, + ) + default_breakable_cudagraph_architectures.cache_clear() + config = object.__new__(VllmConfig) + config.model_config = SimpleNamespace( + architectures=["Qwen3ForCausalLM"], + enforce_eager=case == "eager", + dtype=torch.float16 if case == "fp16" else torch.bfloat16, + quantization="fp8" if case == "quantized" else None, + hf_text_config=SimpleNamespace(intermediate_size=intermediate), + get_hidden_size=lambda: hidden, + get_head_size=lambda: 128, + get_num_attention_heads=lambda pc: heads[0], + get_num_kv_heads=lambda pc: heads[1], + ) + config.parallel_config = SimpleNamespace(tensor_parallel_size=tp) + config.compilation_config = ( + CompilationConfig(pass_config=PassConfig(enable_sp=True)) + if case == "sp" + else CompilationConfig() + ) + try: + assert config._maybe_enable_breakable_cudagraph() is expected + if expected: + assert config.compilation_config.mode == CompilationMode.NONE + finally: + os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None) + default_breakable_cudagraph_architectures.cache_clear() + + @pytest.mark.parametrize( ("model_type", "expected_architecture"), [ diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 35b214f980..54ff2d483b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -776,6 +776,45 @@ class VllmConfig: architectures = set(model_config.architectures) return bool(architectures & default_breakable_cudagraph_architectures()) + def _uses_breakable_cudagraph_for_batch_invariance(self) -> bool: + """Avoid freezing runtime-M tile lookup in compiled forward (#54243). + Breakable graphs look up tuned bf16, unquantized qkv/o/gate_up/down tiles + at capture; lm_head runs outside compiled forward and does not benefit.""" + from vllm.model_executor.determinism import batch_invariant_configs as bi + from vllm.platforms import current_platform + + model = self.model_config + if ( + not envs.VLLM_BATCH_INVARIANT + or model is None + or model.enforce_eager + or model.dtype != torch.bfloat16 + or model.quantization is not None + or not current_platform.is_cuda() + # Sequence parallelism / async TP are torch.compile passes. + or self.compilation_config.pass_config.enable_sp + or self.compilation_config.pass_config.fuse_gemm_comms + ): + return False + family = bi._get_tuned_matmul_arch_family( + current_platform.get_device_capability() + ) + if family is None or family not in bi._BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: + return False + table = bi._BATCH_INVARIANT_MATMUL_TUNED_CONFIGS[family] + parallel = self.parallel_config + tp = parallel.tensor_parallel_size + hidden = model.get_hidden_size() + head = model.get_head_size() + heads = model.get_num_attention_heads(parallel) + kv_heads = model.get_num_kv_heads(parallel) + shapes = [((heads + 2 * kv_heads) * head, hidden), (hidden, heads * head)] + intermediate = getattr(model.hf_text_config, "intermediate_size", None) + # Per-layer sizes (e.g. Gemma3n) are not modeled. + if isinstance(intermediate, int): + shapes += [(2 * intermediate // tp, hidden), (hidden, intermediate // tp)] + return any(shape in table for shape in shapes) + def _maybe_enable_breakable_cudagraph(self) -> bool: if ( "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ @@ -786,6 +825,16 @@ class VllmConfig: "Auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. " "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out." ) + elif ( + "VLLM_USE_BREAKABLE_CUDAGRAPH" not in os.environ + and envs.VLLM_BATCH_INVARIANT + and self._uses_breakable_cudagraph_for_batch_invariance() + ): + os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" + logger.info_once( + "VLLM_BATCH_INVARIANT=1: auto-enabling VLLM_USE_BREAKABLE_CUDAGRAPH=1. " + "Set VLLM_USE_BREAKABLE_CUDAGRAPH=0 to opt out." + ) from vllm.compilation.breakable_cudagraph import ( is_breakable_cudagraph_enabled,