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, 04:45 UTC
[BugFix] Add an env to disable moe chunking to work around compile incompatibility (#19642)
Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com>
Ye (Charlotte) Qi committed
Jun 22, 2025 at 15:17 UTC
33d51f599e414c3639bdea7e5e8fb0c7bb61cc3d
2 files changed
+12
-1
vllm/envs.py
+7
@@ -49,6 +49,7 @@ if TYPE_CHECKING:
49
VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
50
VLLM_XLA_CHECK_RECOMPILATION: bool = False
51
VLLM_FUSED_MOE_CHUNK_SIZE: int = 64 * 1024
52
+ VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING: bool = True
53
VLLM_USE_RAY_SPMD_WORKER: bool = False
54
VLLM_USE_RAY_COMPILED_DAG: bool = False
55
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "auto"
@@ -535,6 +536,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
536
lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
537
"VLLM_FUSED_MOE_CHUNK_SIZE":
538
lambda: int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768")),
539
+ # Control whether to use fused MoE activation chunking. Current chunking
540
+ # logic is incompatible with torch.compile and causes IMA. See issue
541
+ # https://github.com/vllm-project/vllm/issues/19631.
542
+ "VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING":
543
+ lambda: bool(
544
+ int(os.getenv("VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING", "1"))),
545
546
# If set, vllm will skip the deprecation warnings.
547
"VLLM_NO_DEPRECATION_WARNING":
vllm/model_executor/layers/fused_moe/modular_kernel.py
+5
-1
@@ -225,6 +225,10 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
225
else:
226
raise ValueError(f"Unsupported FusedMoe activation: {activation}")
227
228
+ def enable_chunking(self):
229
+ return envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and \
230
+ self.supports_chunking()
231
+
232
@abstractmethod
233
def apply(
234
self,
@@ -400,7 +404,7 @@ class FusedMoEModularKernel(torch.nn.Module):
404
else:
405
_, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
406
403
- if self.fused_experts.supports_chunking():
407
+ if self.fused_experts.enable_chunking():
408
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
409
num_chunks = cdiv(M, CHUNK_SIZE)
410
else: