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
[UX] Enforce valid choices for envs like VLLM_ATTENTION_BACKEND, etc (#24761)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
Sep 16, 2025 at 23:42 UTC
dd83a157f12c0ba7f1357f7954cf85aff2c3b882
1 file changed
+77
-24
vllm/envs.py
+77
-24
@@ -6,7 +6,7 @@ import json
6
import os
7
import sys
8
import tempfile
9
-from typing import TYPE_CHECKING, Any, Callable, Optional
9
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
10
11
if TYPE_CHECKING:
12
VLLM_HOST_IP: str = ""
@@ -56,11 +56,12 @@ if TYPE_CHECKING:
56
VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING: bool = True
57
VLLM_USE_RAY_SPMD_WORKER: bool = False
58
VLLM_USE_RAY_COMPILED_DAG: bool = False
59
- VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "auto"
59
+ VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl",
60
+ "shm"] = "auto"
61
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
62
VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
63
VLLM_XLA_USE_SPMD: bool = False
63
- VLLM_WORKER_MULTIPROC_METHOD: str = "fork"
64
+ VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork"
65
VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets")
66
VLLM_IMAGE_FETCH_TIMEOUT: int = 5
67
VLLM_VIDEO_FETCH_TIMEOUT: int = 30
@@ -77,7 +78,8 @@ if TYPE_CHECKING:
78
VLLM_DOCKER_BUILD_CONTEXT: bool = False
79
VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL: bool = False
80
VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False
80
- CMAKE_BUILD_TYPE: Optional[str] = None
81
+ CMAKE_BUILD_TYPE: Optional[Literal["Debug", "Release",
82
+ "RelWithDebInfo"]] = None
83
VERBOSE: bool = False
84
VLLM_ALLOW_LONG_MAX_MODEL_LEN: bool = False
85
VLLM_RPC_TIMEOUT: int = 10000 # ms
@@ -140,22 +142,25 @@ if TYPE_CHECKING:
142
VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
143
VLLM_USE_FLASHINFER_MOE_FP8: bool = False
144
VLLM_USE_FLASHINFER_MOE_FP4: bool = False
143
- VLLM_FLASHINFER_MOE_BACKEND: str = "throughput"
145
+ VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput",
146
+ "latency"] = "throughput"
147
VLLM_XGRAMMAR_CACHE_MB: int = 0
148
VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
149
VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
150
VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
151
VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5557
149
- VLLM_ALL2ALL_BACKEND: str = "naive"
152
+ VLLM_ALL2ALL_BACKEND: Literal["naive", "pplx", "deepep_high_throughput",
153
+ "deepep_low_latency"] = "naive"
154
VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
155
VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
156
VLLM_SLEEP_WHEN_IDLE: bool = False
157
VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
158
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
155
- VLLM_KV_CACHE_LAYOUT: Optional[str] = None
159
+ VLLM_KV_CACHE_LAYOUT: Optional[Literal["NHD", "HND"]] = None
160
VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
161
VLLM_USE_NVFP4_CT_EMULATIONS: bool = False
158
- VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: str = "NONE"
162
+ VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal["FP", "INT8", "INT6", "INT4",
163
+ "NONE"] = "NONE"
164
VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
165
VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: Optional[int] = None
166
VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 120
@@ -207,6 +212,48 @@ def maybe_convert_bool(value: Optional[str]) -> Optional[bool]:
212
return bool(int(value))
213
214
215
+def env_with_choices(
216
+ env_name: str,
217
+ default: Optional[str],
218
+ choices: Union[list[str], Callable[[], list[str]]],
219
+ case_sensitive: bool = True) -> Callable[[], Optional[str]]:
220
+ """
221
+ Create a lambda that validates environment variable against allowed choices
222
+
223
+ Args:
224
+ env_name: Name of the environment variable
225
+ default: Default value if not set (can be None)
226
+ choices: List of valid string options or callable that returns list
227
+ case_sensitive: Whether validation should be case sensitive
228
+
229
+ Returns:
230
+ Lambda function for environment_variables dict
231
+ """
232
+
233
+ def _get_validated_env() -> Optional[str]:
234
+ value = os.getenv(env_name)
235
+ if value is None:
236
+ return default
237
+
238
+ # Resolve choices if it's a callable (for lazy loading)
239
+ actual_choices = choices() if callable(choices) else choices
240
+
241
+ if not case_sensitive:
242
+ check_value = value.lower()
243
+ check_choices = [choice.lower() for choice in actual_choices]
244
+ else:
245
+ check_value = value
246
+ check_choices = actual_choices
247
+
248
+ if check_value not in check_choices:
249
+ raise ValueError(f"Invalid value '{value}' for {env_name}. "
250
+ f"Valid options: {actual_choices}.")
251
+
252
+ return value
253
+
254
+ return _get_validated_env
255
+
256
+
257
def get_vllm_port() -> Optional[int]:
258
"""Get the port from VLLM_PORT environment variable.
259
@@ -287,7 +334,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
334
# If not set, defaults to "Debug" or "RelWithDebInfo"
335
# Available options: "Debug", "Release", "RelWithDebInfo"
336
"CMAKE_BUILD_TYPE":
290
- lambda: os.getenv("CMAKE_BUILD_TYPE"),
337
+ env_with_choices("CMAKE_BUILD_TYPE", None,
338
+ ["Debug", "Release", "RelWithDebInfo"]),
339
340
# If set, vllm will print verbose logs during installation
341
"VERBOSE":
@@ -476,7 +524,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
524
lambda: int(os.getenv("VLLM_TRACE_FUNCTION", "0")),
525
526
# Backend for attention computation
479
- # Available options:
527
+ # Example options:
528
# - "TORCH_SDPA": use torch.nn.MultiheadAttention
529
# - "FLASH_ATTN": use FlashAttention
530
# - "XFORMERS": use XFormers
@@ -486,8 +534,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
534
# - "FLASH_ATTN_MLA": use FlashAttention for MLA
535
# - "FLASHINFER_MLA": use FlashInfer for MLA
536
# - "CUTLASS_MLA": use CUTLASS for MLA
537
+ # All possible options loaded dynamically from _Backend enum
538
"VLLM_ATTENTION_BACKEND":
490
- lambda: os.getenv("VLLM_ATTENTION_BACKEND", None),
539
+ env_with_choices("VLLM_ATTENTION_BACKEND", None,
540
+ lambda: list(__import__('vllm.platforms.interface', \
541
+ fromlist=['_Backend'])._Backend.__members__.keys())),
542
543
# If set, vllm will use flashinfer sampler
544
"VLLM_USE_FLASHINFER_SAMPLER":
@@ -550,7 +601,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
601
# - "shm": use shared memory and gRPC for communication
602
# This flag is ignored if VLLM_USE_RAY_COMPILED_DAG is not set.
603
"VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE":
553
- lambda: os.getenv("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto"),
604
+ env_with_choices("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "auto",
605
+ ["auto", "nccl", "shm"]),
606
607
# If the env var is set, it enables GPU communication overlap
608
# (experimental feature) in Ray's Compiled Graph. This flag is ignored if
@@ -569,7 +621,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
621
# Use dedicated multiprocess context for workers.
622
# Both spawn and fork work
623
"VLLM_WORKER_MULTIPROC_METHOD":
572
- lambda: os.getenv("VLLM_WORKER_MULTIPROC_METHOD", "fork"),
624
+ env_with_choices("VLLM_WORKER_MULTIPROC_METHOD", "fork",
625
+ ["spawn", "fork"]),
626
627
# Path to the cache for storing downloaded assets
628
"VLLM_ASSETS_CACHE":
@@ -833,7 +886,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
886
# Choice of quantization level: FP, INT8, INT6, INT4 or NONE
887
# Recommended for large models to get allreduce
888
"VLLM_ROCM_QUICK_REDUCE_QUANTIZATION":
836
- lambda: os.getenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE").upper(),
889
+ env_with_choices("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE",
890
+ ["FP", "INT8", "INT6", "INT4", "NONE"]),
891
892
# Custom quick allreduce kernel for MI3* cards
893
# Due to the lack of the bfloat16 asm instruction, bfloat16
@@ -1075,21 +1129,20 @@ environment_variables: dict[str, Callable[[], Any]] = {
1129
# - "deepep_high_throughput", use deepep high-throughput kernels
1130
# - "deepep_low_latency", use deepep low-latency kernels
1131
"VLLM_ALL2ALL_BACKEND":
1078
- lambda: os.getenv("VLLM_ALL2ALL_BACKEND", "naive"),
1132
+ env_with_choices("VLLM_ALL2ALL_BACKEND", "naive",
1133
+ ["naive", "pplx",
1134
+ "deepep_high_throughput", "deepep_low_latency"]),
1135
1080
- # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support. Both
1081
- # require compute capability 10.0 or above.
1136
+ # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support.
1137
+ # Both require compute capability 10.0 or above.
1138
# Available options:
1139
# - "throughput": [default]
1140
# Uses CUTLASS kernels optimized for high-throughput batch inference.
1141
# - "latency":
1142
# Uses TensorRT-LLM kernels optimized for low-latency inference.
1087
- # To set this backend, define the environment variable:
1088
- # export VLLM_FLASHINFER_MOE_BACKEND=latency.
1089
- # If not set, defaults to "throughput".
1090
- "VLLM_FLASHINFER_MOE_BACKEND": lambda: os.getenv(
1091
- "VLLM_FLASHINFER_MOE_BACKEND", "throughput"
1092
- ),
1143
+ "VLLM_FLASHINFER_MOE_BACKEND":
1144
+ env_with_choices("VLLM_FLASHINFER_MOE_BACKEND", "throughput",
1145
+ ["throughput", "latency"]),
1146
1147
# Control the maximum number of tokens per expert supported by the
1148
# NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for
@@ -1145,7 +1198,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
1198
# leave the layout choice to the backend. Mind that backends may only
1199
# implement and support a subset of all possible layouts.
1200
"VLLM_KV_CACHE_LAYOUT":
1148
- lambda: os.getenv("VLLM_KV_CACHE_LAYOUT", None),
1201
+ env_with_choices("VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]),
1202
1203
# Enable checking whether the generated logits contain NaNs,
1204
# indicating corrupted output. Useful for debugging low level bugs