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 13, 13:15 UTC
[Feature] Add env var `VLLM_MOE_USE_DEEP_GEMM` (#28422)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Nov 10, 2025 at 21:29 UTC
de540c0354b9ecfa979c917a4599f8030d4105be
4 files changed
+19
-2
vllm/envs.py
+6
@@ -147,6 +147,7 @@ if TYPE_CHECKING:
147
VLLM_TPU_MOST_MODEL_LEN: int | None = None
148
VLLM_TPU_USING_PATHWAYS: bool = False
149
VLLM_USE_DEEP_GEMM: bool = True
150
+ VLLM_MOE_USE_DEEP_GEMM: bool = True
151
VLLM_USE_DEEP_GEMM_E8M0: bool = True
152
VLLM_DEEP_GEMM_WARMUP: Literal[
153
"skip",
@@ -1116,6 +1117,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1117
),
1118
# Allow use of DeepGemm kernels for fused moe ops.
1119
"VLLM_USE_DEEP_GEMM": lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "1"))),
1120
+ # Allow use of DeepGemm specifically for MoE fused ops (overrides only MoE).
1121
+ "VLLM_MOE_USE_DEEP_GEMM": lambda: bool(
1122
+ int(os.getenv("VLLM_MOE_USE_DEEP_GEMM", "1"))
1123
+ ),
1124
# Whether to use E8M0 scaling when DeepGEMM is used on Blackwell GPUs.
1125
"VLLM_USE_DEEP_GEMM_E8M0": lambda: bool(
1126
int(os.getenv("VLLM_USE_DEEP_GEMM_E8M0", "1"))
@@ -1569,6 +1574,7 @@ def compute_hash() -> str:
1574
"VLLM_USE_FLASHINFER_SAMPLER",
1575
"VLLM_DISABLED_KERNELS",
1576
"VLLM_USE_DEEP_GEMM",
1577
+ "VLLM_MOE_USE_DEEP_GEMM",
1578
"VLLM_USE_DEEP_GEMM_E8M0",
1579
"VLLM_USE_FUSED_MOE_GROUPED_TOPK",
1580
"VLLM_USE_FLASHINFER_MOE_FP16",
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+9
-1
@@ -966,10 +966,18 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
966
max_num_tokens=max_num_tokens_per_rank,
967
num_dispatchers=prepare_finalize.num_dispatchers(),
968
quant_config=self.moe_quant_config,
969
+ allow_deep_gemm=(
970
+ envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
971
+ ),
972
)
973
else:
974
logger.debug("TritonOrDeepGemmExperts(%s)", self.__class__.__name__)
972
- return TritonOrDeepGemmExperts(self.moe_quant_config, allow_deep_gemm=True)
975
+ return TritonOrDeepGemmExperts(
976
+ self.moe_quant_config,
977
+ allow_deep_gemm=(
978
+ envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
979
+ ),
980
+ )
981
982
def get_fused_moe_quant_config(
983
self, layer: torch.nn.Module
vllm/model_executor/layers/quantization/fp8.py
+1
-1
@@ -158,7 +158,7 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:
158
return Fp8MoeBackend.MARLIN
159
160
# deepGEMM on supported platforms with block-quantized weights
161
- if envs.VLLM_USE_DEEP_GEMM and block_quant:
161
+ if envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM and block_quant:
162
if not has_deep_gemm():
163
logger.warning_once("DeepGEMM backend requested but not available.")
164
elif is_deep_gemm_supported():
vllm/model_executor/warmup/deep_gemm_warmup.py
+3
@@ -148,6 +148,9 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
148
149
150
def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
151
+ if not (envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM):
152
+ return False
153
+
154
if not isinstance(module, FusedMoE):
155
return False
156