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 11, 10:00 UTC
Update batch invariant to use attention config (#30704)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Matthew Bonanni committed
Dec 15, 2025 at 15:24 UTC
60dbf7d8f13689b17c88840f3ae4e7a222305f2b
2 files changed
+24
-18
vllm/model_executor/layers/batch_invariant.py
+22
-17
@@ -6,7 +6,7 @@ from typing import Any
6
7
import torch
8
9
-import vllm.envs as envs
9
+from vllm.attention.backends.registry import AttentionBackendEnum
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
from vllm.triton_utils import tl, triton
@@ -1004,27 +1004,30 @@ def vllm_is_batch_invariant() -> bool:
1004
return VLLM_BATCH_INVARIANT
1005
1006
1007
-def override_envs_for_invariance():
1008
- curr_attn_backend = envs.VLLM_ATTENTION_BACKEND
1007
+def override_envs_for_invariance(
1008
+ attention_backend: AttentionBackendEnum | None,
1009
+):
1010
supported_backends = [
1010
- "FLASH_ATTN", # best supported backend
1011
- "FLASHINFER",
1012
- "FLASH_ATTN_MLA",
1013
- "TRITON_MLA",
1011
+ AttentionBackendEnum.FLASH_ATTN, # best supported backend
1012
+ AttentionBackendEnum.FLASHINFER,
1013
+ AttentionBackendEnum.FLASH_ATTN_MLA,
1014
+ AttentionBackendEnum.TRITON_MLA,
1015
# Not yet supported MLA backends
1015
- # "FLASHMLA",
1016
- # "FLEX_ATTENTION", # IMA issue even if we disable batch invariance
1017
- # "FLASHINFER_MLA", https://github.com/vllm-project/vllm/pull/28967
1016
+ # AttentionBackendEnum.FLASHMLA,
1017
+ # AttentionBackendEnum.FLEX_ATTENTION, # IMA issue
1018
+ # AttentionBackendEnum.FLASHINFER_MLA, # PR #28967
1019
]
1019
- if curr_attn_backend not in supported_backends:
1020
+ if attention_backend not in supported_backends:
1021
+ supported_names = [b.name for b in supported_backends]
1022
+ backend_name = attention_backend.name if attention_backend else None
1023
error = (
1024
"VLLM batch_invariant mode requires an attention backend in "
1022
- f"{supported_backends}, but got '{curr_attn_backend}'. "
1023
- "Please set the 'VLLM_ATTENTION_BACKEND' environment variable "
1024
- "to one of the supported backends before enabling batch_invariant."
1025
+ f"{supported_names}, but got '{backend_name}'. "
1026
+ "Please use --attention-backend or attention_config to set "
1027
+ "one of the supported backends before enabling batch_invariant."
1028
)
1029
raise RuntimeError(error)
1027
- if os.environ["VLLM_ATTENTION_BACKEND"] != supported_backends[0]:
1030
+ if attention_backend != supported_backends[0]:
1031
warning = (
1032
"You are using a decode-invariant form of batch invariance. "
1033
"This will not be invariant between prefill and decode."
@@ -1050,10 +1053,12 @@ def override_envs_for_invariance():
1053
os.environ["VLLM_USE_AOT_COMPILE"] = "0"
1054
1055
1053
-def init_batch_invariance():
1056
+def init_batch_invariance(
1057
+ attention_backend: AttentionBackendEnum | None,
1058
+):
1059
# this will hit all the csrc overrides as well
1060
if vllm_is_batch_invariant():
1056
- override_envs_for_invariance()
1061
+ override_envs_for_invariance(attention_backend)
1062
enable_batch_invariant_mode()
1063
1064
# Disable TF32 for batch invariance - it causes non-deterministic rounding
vllm/v1/worker/gpu_worker.py
+2
-1
@@ -931,10 +931,11 @@ def init_worker_distributed_environment(
931
backend: str = "nccl",
932
) -> None:
933
"""Initialize the distributed environment."""
934
+ attention_config = vllm_config.attention_config
935
parallel_config = vllm_config.parallel_config
936
from vllm.model_executor.layers.batch_invariant import init_batch_invariance
937
937
- init_batch_invariance()
938
+ init_batch_invariance(attention_config.backend)
939
set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
940
941
init_method = distributed_init_method or "env://"