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 14, 05:30 UTC
[CLI env var] Add VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH in env variables (#25274)
Signed-off-by: qqma <qqma@amazon.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: qqma <qqma@amazon.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Daisy-Ma-coder committed
Sep 22, 2025 at 10:37 UTC
cfbee3d0e72582b0b0d910ecbd9a7c1028a7205c
5 files changed
+32
-13
tests/compile/piecewise/test_full_cudagraph.py
+9
-2
@@ -46,7 +46,10 @@ backend_configs = {
46
# FA3 on Hopper
47
"FA3":
48
BackendConfig(name="FA3",
49
- env_vars={"VLLM_FLASH_ATTN_VERSION": "3"},
49
+ env_vars={
50
+ "VLLM_FLASH_ATTN_VERSION": "3",
51
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
52
+ },
53
comp_config={
54
"cudagraph_mode": "FULL",
55
},
@@ -66,6 +69,7 @@ backend_configs = {
69
BackendConfig(name="FlashAttentionMLA",
70
env_vars={
71
"VLLM_ATTENTION_BACKEND": "FLASH_ATTN_MLA",
72
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
73
},
74
comp_config={
75
"cudagraph_mode": "FULL_DECODE_ONLY",
@@ -89,7 +93,10 @@ backend_configs = {
93
# FA2
94
"FA2":
95
BackendConfig(name="FA2",
92
- env_vars={"VLLM_FLASH_ATTN_VERSION": "2"},
96
+ env_vars={
97
+ "VLLM_FLASH_ATTN_VERSION": "2",
98
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
99
+ },
100
comp_config={
101
"cudagraph_mode": "FULL",
102
}),
tests/v1/cudagraph/test_cudagraph_mode.py
+9
-2
@@ -47,7 +47,10 @@ backend_configs = {
47
# FA3 on Hopper
48
"FA3":
49
BackendConfig(name="FA3",
50
- env_vars={"VLLM_FLASH_ATTN_VERSION": "3"},
50
+ env_vars={
51
+ "VLLM_FLASH_ATTN_VERSION": "3",
52
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
53
+ },
54
comp_config={
55
"cudagraph_mode": "FULL",
56
},
@@ -67,6 +70,7 @@ backend_configs = {
70
BackendConfig(name="FlashAttentionMLA",
71
env_vars={
72
"VLLM_ATTENTION_BACKEND": "FLASH_ATTN_MLA",
73
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
74
},
75
comp_config={
76
"cudagraph_mode": "FULL_DECODE_ONLY",
@@ -75,7 +79,10 @@ backend_configs = {
79
# FA2
80
"FA2":
81
BackendConfig(name="FA2",
78
- env_vars={"VLLM_FLASH_ATTN_VERSION": "2"},
82
+ env_vars={
83
+ "VLLM_FLASH_ATTN_VERSION": "2",
84
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH": "16",
85
+ },
86
comp_config={
87
"cudagraph_mode": "FULL_AND_PIECEWISE",
88
}),
vllm/envs.py
+8
@@ -119,6 +119,7 @@ if TYPE_CHECKING:
119
VLLM_SERVER_DEV_MODE: bool = False
120
VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128
121
VLLM_MLA_DISABLE: bool = False
122
+ VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH: int = 16
123
VLLM_RAY_PER_WORKER_GPUS: float = 1.0
124
VLLM_RAY_BUNDLE_INDICES: str = ""
125
VLLM_CUDART_SO_PATH: Optional[str] = None
@@ -946,6 +947,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
947
"VLLM_MLA_DISABLE":
948
lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))),
949
950
+ # If set, vLLM will pick up the provided Flash Attention MLA
951
+ # max number splits for cuda graph decode
952
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH":
953
+ lambda: int(os.getenv("VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH",
954
+ "16")),
955
+
956
# Number of GPUs per worker in Ray, if it is set to be a fraction,
957
# it allows ray to schedule multiple actors on a single GPU,
958
# so that users can colocate other actors on the same GPUs as vLLM.
@@ -1379,6 +1386,7 @@ def compute_hash() -> str:
1386
environment_variables_to_hash = [
1387
"VLLM_PP_LAYER_PARTITION",
1388
"VLLM_MLA_DISABLE",
1389
+ "VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH",
1390
"VLLM_USE_TRITON_FLASH_ATTN",
1391
"VLLM_USE_TRITON_AWQ",
1392
"VLLM_DP_RANK",
vllm/v1/attention/backends/flash_attn.py
+3
-4
@@ -8,6 +8,7 @@ import numpy as np
8
import torch
9
10
from vllm import _custom_ops as ops
11
+from vllm import envs
12
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
13
AttentionMetadata, AttentionType,
14
is_quantized_kv_cache)
@@ -33,9 +34,6 @@ from vllm.v1.kv_cache_interface import AttentionSpec
34
35
logger = init_logger(__name__)
36
36
-# NOTE(woosuk): This is an arbitrary number. Tune it if needed.
37
-_DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH = 16
38
-
37
38
class FlashAttentionBackend(AttentionBackend):
39
@@ -215,7 +213,8 @@ class FlashAttentionMetadataBuilder(
213
# When using cuda graph, we need to set the upper bound of the
214
# number of splits so that large enough intermediate buffers are
215
# pre-allocated during capture.
218
- self.max_num_splits = _DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
216
+ self.max_num_splits = (
217
+ envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH)
218
219
# Sliding window size to be used with the AOT scheduler will be
220
# populated on first build() call.
vllm/v1/attention/backends/mla/flashattn_mla.py
+3
-5
@@ -6,6 +6,7 @@ from typing import ClassVar, Optional, Union
6
7
import torch
8
9
+from vllm import envs
10
from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,
11
is_quantized_kv_cache)
12
from vllm.attention.utils.fa_utils import (flash_attn_supports_mla,
@@ -24,10 +25,6 @@ from vllm.vllm_flash_attn import flash_attn_varlen_func, get_scheduler_metadata
25
26
logger = init_logger(__name__)
27
27
-# NOTE(matt): This is an arbitrary number, copied from
28
-# woosuk's implementation in standard FlashAttention backend
29
-_DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH = 16
30
-
28
29
class FlashAttnMLABackend(MLACommonBackend):
30
@@ -97,7 +94,8 @@ class FlashAttnMLAMetadataBuilder(
94
# When using cuda graph, we need to set the upper bound of the
95
# number of splits so that large enough intermediate buffers are
96
# pre-allocated during capture.
100
- self.max_num_splits = _DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
97
+ self.max_num_splits = (
98
+ envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH)
99
100
# TODO(lucas): Until we add support for the DCP custom masking we need
101
# to restrict decodes to q_len == 1 when DCP is enabled.