xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-25 21:51:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update batch invariant to use attention config (#30704)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Matthew Bonanni 2025-12-15 15:24:16 -05:00 committed by GitHub
parent a450c64a30
commit 60dbf7d8f1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
vllm/model_executor/layers/batch_invariant.py
View File

@ -6,7 +6,7 @@ from typing import Any
import torch import torch
import vllm.envs as envs from vllm.attention.backends.registry import AttentionBackendEnum
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.platforms import current_platform from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton from vllm.triton_utils import tl, triton
@ -1004,27 +1004,30 @@ def vllm_is_batch_invariant() -> bool:
return VLLM_BATCH_INVARIANT return VLLM_BATCH_INVARIANT
def override_envs_for_invariance(): def override_envs_for_invariance(
curr_attn_backend = envs.VLLM_ATTENTION_BACKEND attention_backend: AttentionBackendEnum | None,
):
supported_backends = [ supported_backends = [
"FLASH_ATTN", # best supported backend AttentionBackendEnum.FLASH_ATTN, # best supported backend
"FLASHINFER", AttentionBackendEnum.FLASHINFER,
"FLASH_ATTN_MLA", AttentionBackendEnum.FLASH_ATTN_MLA,
"TRITON_MLA", AttentionBackendEnum.TRITON_MLA,
# Not yet supported MLA backends # Not yet supported MLA backends
# "FLASHMLA", # AttentionBackendEnum.FLASHMLA,
# "FLEX_ATTENTION", # IMA issue even if we disable batch invariance # AttentionBackendEnum.FLEX_ATTENTION, # IMA issue
# "FLASHINFER_MLA", https://github.com/vllm-project/vllm/pull/28967 # AttentionBackendEnum.FLASHINFER_MLA, # PR #28967
] ]
if curr_attn_backend not in supported_backends: if attention_backend not in supported_backends:
supported_names = [b.name for b in supported_backends]
backend_name = attention_backend.name if attention_backend else None
error = ( error = (
"VLLM batch_invariant mode requires an attention backend in " "VLLM batch_invariant mode requires an attention backend in "
f"{supported_backends}, but got '{curr_attn_backend}'. " f"{supported_names}, but got '{backend_name}'. "
"Please set the 'VLLM_ATTENTION_BACKEND' environment variable " "Please use --attention-backend or attention_config to set "
"to one of the supported backends before enabling batch_invariant." "one of the supported backends before enabling batch_invariant."
) )
raise RuntimeError(error) raise RuntimeError(error)
if os.environ["VLLM_ATTENTION_BACKEND"] != supported_backends[0]: if attention_backend != supported_backends[0]:
warning = ( warning = (
"You are using a decode-invariant form of batch invariance. " "You are using a decode-invariant form of batch invariance. "
"This will not be invariant between prefill and decode." "This will not be invariant between prefill and decode."
@ -1050,10 +1053,12 @@ def override_envs_for_invariance():
os.environ["VLLM_USE_AOT_COMPILE"] = "0" os.environ["VLLM_USE_AOT_COMPILE"] = "0"
def init_batch_invariance(): def init_batch_invariance(
attention_backend: AttentionBackendEnum | None,
):
# this will hit all the csrc overrides as well # this will hit all the csrc overrides as well
if vllm_is_batch_invariant(): if vllm_is_batch_invariant():
override_envs_for_invariance() override_envs_for_invariance(attention_backend)
enable_batch_invariant_mode() enable_batch_invariant_mode()
# Disable TF32 for batch invariance - it causes non-deterministic rounding # Disable TF32 for batch invariance - it causes non-deterministic rounding

3
vllm/v1/worker/gpu_worker.py
View File

@ -931,10 +931,11 @@ def init_worker_distributed_environment(
backend: str = "nccl", backend: str = "nccl",
) -> None: ) -> None:
"""Initialize the distributed environment.""" """Initialize the distributed environment."""
attention_config = vllm_config.attention_config
parallel_config = vllm_config.parallel_config parallel_config = vllm_config.parallel_config
from vllm.model_executor.layers.batch_invariant import init_batch_invariance from vllm.model_executor.layers.batch_invariant import init_batch_invariance
init_batch_invariance() init_batch_invariance(attention_config.backend)
set_custom_all_reduce(not parallel_config.disable_custom_all_reduce) set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
init_method = distributed_init_method or "env://" init_method = distributed_init_method or "env://"