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 15, 06:00 UTC
[Misc] Turn MOE_DP_CHUNK_SIZE into an env var (#19506)
Varun Sundar Rabindranath committed
Jun 12, 2025 at 14:01 UTC
9d880f594d6bb85a2657c147e2357b42a147c709
2 files changed
+18
-8
vllm/envs.py
+9
@@ -112,6 +112,7 @@ if TYPE_CHECKING:
112
VLLM_DP_SIZE: int = 1
113
VLLM_DP_MASTER_IP: str = ""
114
VLLM_DP_MASTER_PORT: int = 0
115
+ VLLM_MOE_DP_CHUNK_SIZE: int = 256
116
VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
117
VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
118
VLLM_V0_USE_OUTLINES_CACHE: bool = False
@@ -773,6 +774,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
774
"VLLM_DP_MASTER_PORT":
775
lambda: int(os.getenv("VLLM_DP_MASTER_PORT", "0")),
776
777
+ # In the context of executing MoE models with Data-Parallel, Expert-Parallel
778
+ # and Batched All-to-All dispatch/combine kernels, VLLM_MOE_DP_CHUNK_SIZE
779
+ # dictates the quantum of tokens that can be dispatched from a DP
780
+ # rank. All DP ranks process the activations in VLLM_MOE_DP_CHUNK_SIZE
781
+ # units.
782
+ "VLLM_MOE_DP_CHUNK_SIZE":
783
+ lambda: int(os.getenv("VLLM_MOE_DP_CHUNK_SIZE", "256")),
784
+
785
# Randomize inputs during dummy runs when using Data Parallel
786
"VLLM_RANDOMIZE_DP_DUMMY_INPUTS":
787
lambda: os.environ.get("VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "0") == "1",
vllm/model_executor/layers/fused_moe/layer.py
+9
-8
@@ -61,10 +61,6 @@ else:
61
fused_moe_pallas = None # type: ignore
62
logger = init_logger(__name__)
63
64
-# Note: this limit is somewhat arbitrary and might be changed later.
65
-# The size of the activations will be E x MOE_DP_CHUNK_SIZE x hidden_dim.
66
-MOE_DP_CHUNK_SIZE = 256
67
-
64
65
@dataclass
66
class FusedMoEParallelConfig:
@@ -218,7 +214,12 @@ class MoEConfig:
214
# TODO: add more quantization params, blocked, per-token, etc.
215
block_size: int = 128
216
221
- max_num_tokens: int = MOE_DP_CHUNK_SIZE
217
+ max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
218
+
219
+ def __post_init__(self):
220
+ if self.dp_size > 1:
221
+ logger.debug("Using MOEConfig::max_num_tokens=%d",
222
+ self.max_num_tokens)
223
224
@property
225
def tp_size(self):
@@ -913,7 +914,7 @@ class FusedMoE(torch.nn.Module):
914
moe_parallel_config=self.moe_parallel_config,
915
in_dtype=params_dtype,
916
quant_dtype=quant_dtype,
916
- max_num_tokens=MOE_DP_CHUNK_SIZE,
917
+ max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
918
)
919
self.moe_config = moe
920
self.quant_config = quant_config
@@ -952,12 +953,12 @@ class FusedMoE(torch.nn.Module):
953
or self.moe_parallel_config.use_deepep_ll_kernels):
954
act_dtype = vllm_config.model_config.dtype
955
self.batched_hidden_states = torch.zeros(
955
- (MOE_DP_CHUNK_SIZE, self.hidden_size),
956
+ (envs.VLLM_MOE_DP_CHUNK_SIZE, self.hidden_size),
957
dtype=act_dtype,
958
device=torch.cuda.current_device())
959
960
self.batched_router_logits = torch.zeros(
960
- (MOE_DP_CHUNK_SIZE, self.global_num_experts),
961
+ (envs.VLLM_MOE_DP_CHUNK_SIZE, self.global_num_experts),
962
dtype=act_dtype,
963
device=torch.cuda.current_device())
964