A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 17:30 UTC
[Bugfix] Make Attention Backend Auto-Selection Batch-Invariance-Aware (#40193)
Signed-off-by: Srreyansh Sethi <srreyansh.sethi@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Srreyansh Sethi committed
Apr 23, 2026 at 07:57 UTC
b7a26050200e20917871a4a6b09df0bc9ea3fdc7
9 files changed
+38
-50
examples/rl/rlhf_async_new_apis.py
+1
-8
@@ -131,16 +131,9 @@ class TrainModel:
131
from vllm.model_executor.layers.batch_invariant import (
132
init_batch_invariance,
133
)
134
- from vllm.platforms import current_platform
135
- from vllm.v1.attention.backends.registry import AttentionBackendEnum
134
135
# need to init all env vars for batch invariance which affect nccl ops
138
- attn_backend = (
139
- AttentionBackendEnum.TRITON_ATTN
140
- if current_platform.is_rocm()
141
- else AttentionBackendEnum.FLASH_ATTN
142
- )
143
- init_batch_invariance(attn_backend)
136
+ init_batch_invariance()
137
138
self.model = AutoModelForCausalLM.from_pretrained(
139
model_name, dtype=torch.bfloat16
vllm/model_executor/layers/batch_invariant.py
+3
-39
@@ -14,7 +14,6 @@ from vllm.triton_utils import tl, triton
14
from vllm.utils.mem_utils import get_max_shared_memory_bytes
15
from vllm.utils.platform_utils import num_compute_units
16
from vllm.utils.torch_utils import is_torch_equal_or_newer
17
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
17
18
logger = init_logger(__name__)
19
@@ -991,40 +990,7 @@ def enable_batch_invariant_mode():
990
torch.backends.cuda.preferred_blas_library(backend="cublaslt")
991
992
994
-def override_envs_for_invariance(
995
- attention_backend: AttentionBackendEnum | None,
996
-):
997
- decode_invariant_backends = [
998
- AttentionBackendEnum.FLASH_ATTN, # best supported backend
999
- AttentionBackendEnum.TRITON_ATTN,
1000
- ]
1001
- supported_backends = decode_invariant_backends + [
1002
- # FlashInfer temporarily disabled due to invariant CTA sizes.
1003
- # See FlashInfer issue #2424
1004
- # AttentionBackendEnum.FLASHINFER,
1005
- AttentionBackendEnum.FLASH_ATTN_MLA,
1006
- AttentionBackendEnum.TRITON_MLA,
1007
- # Not yet supported MLA backends
1008
- # AttentionBackendEnum.FLASHMLA,
1009
- # AttentionBackendEnum.FLEX_ATTENTION, # IMA issue
1010
- # AttentionBackendEnum.FLASHINFER_MLA, # PR #28967
1011
- ]
1012
- if attention_backend not in supported_backends:
1013
- supported_names = [b.name for b in supported_backends]
1014
- backend_name = attention_backend.name if attention_backend else None
1015
- error = (
1016
- "VLLM batch_invariant mode requires an attention backend in "
1017
- f"{supported_names}, but got '{backend_name}'. "
1018
- "Please use --attention-backend or attention_config to set "
1019
- "one of the supported backends before enabling batch_invariant."
1020
- )
1021
- raise RuntimeError(error)
1022
- if attention_backend not in decode_invariant_backends:
1023
- warning = (
1024
- "You are using a non-decode-invariant form of batch invariance. "
1025
- "This will not be invariant between prefill and decode."
1026
- )
1027
- logger.warning_once(warning)
993
+def override_envs_for_invariance():
994
os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
995
996
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
@@ -1045,12 +1011,10 @@ def override_envs_for_invariance(
1011
os.environ["VLLM_USE_AOT_COMPILE"] = "0"
1012
1013
1048
-def init_batch_invariance(
1049
- attention_backend: AttentionBackendEnum | None,
1050
-):
1014
+def init_batch_invariance():
1015
# this will hit all the csrc overrides as well
1016
if envs.VLLM_BATCH_INVARIANT:
1053
- override_envs_for_invariance(attention_backend)
1017
+ override_envs_for_invariance()
1018
enable_batch_invariant_mode()
1019
1020
# Disable TF32 for batch invariance - it causes non-deterministic rounding
vllm/v1/attention/backend.py
+7
@@ -236,6 +236,10 @@ class AttentionBackend(ABC):
236
"""
237
return False
238
239
+ @classmethod
240
+ def supports_batch_invariance(cls) -> bool:
241
+ return False
242
+
243
@classmethod
244
def supports_attn_type(cls, attn_type: str) -> bool:
245
"""Check if backend supports a given attention type.
@@ -278,6 +282,7 @@ class AttentionBackend(ABC):
282
device_capability: "DeviceCapability",
283
attn_type: str,
284
use_non_causal: bool = False,
285
+ use_batch_invariant: bool = False,
286
) -> list[str]:
287
invalid_reasons = []
288
if not cls.supports_head_size(head_size):
@@ -312,6 +317,8 @@ class AttentionBackend(ABC):
317
invalid_reasons.append(f"attention type {attn_type} not supported")
318
if use_non_causal and not cls.supports_non_causal():
319
invalid_reasons.append("non-causal attention not supported")
320
+ if use_batch_invariant and not cls.supports_batch_invariance():
321
+ invalid_reasons.append("batch invariance not supported")
322
combination_reason = cls.supports_combination(
323
head_size,
324
dtype,
vllm/v1/attention/backends/flash_attn.py
+4
@@ -103,6 +103,10 @@ class FlashAttentionBackend(AttentionBackend):
103
def get_name() -> str:
104
return "FLASH_ATTN"
105
106
+ @classmethod
107
+ def supports_batch_invariance(cls) -> bool:
108
+ return True
109
+
110
@classmethod
111
def supports_non_causal(cls) -> bool:
112
return True
vllm/v1/attention/backends/mla/flashattn_mla.py
+4
@@ -56,6 +56,10 @@ class FlashAttnMLABackend(MLACommonBackend):
56
def get_name() -> str:
57
return "FLASH_ATTN_MLA"
58
59
+ @classmethod
60
+ def supports_batch_invariance(cls) -> bool:
61
+ return True
62
+
63
@staticmethod
64
def get_builder_cls() -> type["FlashAttnMLAMetadataBuilder"]:
65
return FlashAttnMLAMetadataBuilder
vllm/v1/attention/backends/mla/triton_mla.py
+4
@@ -55,6 +55,10 @@ class TritonMLABackend(MLACommonBackend):
55
def get_name() -> str:
56
return "TRITON_MLA"
57
58
+ @classmethod
59
+ def supports_batch_invariance(cls) -> bool:
60
+ return True
61
+
62
@staticmethod
63
def get_impl_cls() -> type["TritonMLAImpl"]:
64
return TritonMLAImpl
vllm/v1/attention/backends/triton_attn.py
+4
@@ -296,6 +296,10 @@ class TritonAttentionBackend(AttentionBackend):
296
def get_name() -> str:
297
return "TRITON_ATTN"
298
299
+ @classmethod
300
+ def supports_batch_invariance(cls) -> bool:
301
+ return True
302
+
303
@staticmethod
304
def get_impl_cls() -> type["TritonAttentionImpl"]:
305
return TritonAttentionImpl
vllm/v1/attention/selector.py
+10
-1
@@ -6,6 +6,7 @@ from typing import NamedTuple, cast, get_args
6
7
import torch
8
9
+import vllm.envs as envs
10
from vllm.config.cache import CacheDType
11
from vllm.logger import init_logger
12
from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -30,6 +31,7 @@ class AttentionSelectorConfig(NamedTuple):
31
use_per_head_quant_scales: bool = False
32
attn_type: str = AttentionType.DECODER
33
use_non_causal: bool = False
34
+ use_batch_invariant: bool = False
35
36
def __repr__(self):
37
return (
@@ -43,7 +45,8 @@ class AttentionSelectorConfig(NamedTuple):
45
f"use_mm_prefix={self.use_mm_prefix}, "
46
f"use_per_head_quant_scales={self.use_per_head_quant_scales}, "
47
f"attn_type={self.attn_type}, "
46
- f"use_non_causal={self.use_non_causal})"
48
+ f"use_non_causal={self.use_non_causal}, "
49
+ f"use_batch_invariant={self.use_batch_invariant})"
50
)
51
52
@@ -95,6 +98,7 @@ def get_attn_backend(
98
use_per_head_quant_scales=use_per_head_quant_scales,
99
attn_type=attn_type or AttentionType.DECODER,
100
use_non_causal=use_non_causal,
101
+ use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
102
)
103
104
return _cached_get_attn_backend(
@@ -162,4 +166,9 @@ def _cached_get_mamba_attn_backend(
166
) from e
167
168
mamba_attn_backend = selected_backend.get_class()
169
+ if envs.VLLM_BATCH_INVARIANT and not mamba_attn_backend.supports_batch_invariance():
170
+ raise RuntimeError(
171
+ "VLLM batch_invariant mode is not supported for "
172
+ f"{mamba_attn_backend.get_name()}."
173
+ )
174
return mamba_attn_backend
vllm/v1/worker/gpu_worker.py
+1
-2
@@ -1027,11 +1027,10 @@ def init_worker_distributed_environment(
1027
backend: str = "nccl",
1028
) -> None:
1029
"""Initialize the distributed environment."""
1030
- attention_config = vllm_config.attention_config
1030
parallel_config = vllm_config.parallel_config
1031
from vllm.model_executor.layers.batch_invariant import init_batch_invariance
1032
1034
- init_batch_invariance(attention_config.backend)
1033
+ init_batch_invariance()
1034
override_envs_for_eplb(parallel_config)
1035
set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
1036