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
[FEATURE] support custom vllm tuned config path for fused moe triton kernels (#22791)
Signed-off-by: Chi Zhang <zhangchi.usc1992@bytedance.com>
Chi Zhang committed
Aug 13, 2025 at 20:27 UTC
98deac3879860b829dd9a30b19bbb2adb9c96e7f
2 files changed
+26
-8
vllm/envs.py
+6
@@ -158,6 +158,7 @@ if TYPE_CHECKING:
158
VLLM_USE_TRTLLM_ATTENTION: Optional[str] = None
159
VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False
160
VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
161
+ VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
162
163
164
def get_default_cache_root():
@@ -1120,6 +1121,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1121
# never removed from memory until the server terminates.
1122
"VLLM_ENABLE_RESPONSES_API_STORE":
1123
lambda: bool(int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0"))),
1124
+
1125
+ # Allows vllm to find tuned config under customized folder
1126
+ "VLLM_TUNED_CONFIG_FOLDER":
1127
+ lambda: os.getenv("VLLM_TUNED_CONFIG_FOLDER", None),
1128
+
1129
}
1130
1131
# --8<-- [end:env-vars-definition]
vllm/model_executor/layers/fused_moe/fused_moe.py
+20
-8
@@ -701,20 +701,32 @@ def get_moe_configs(
701
block_shape = [block_n, block_k] if block_n and block_k else None
702
json_file_name = get_config_file_name(E, N, dtype, block_shape)
703
704
- config_file_path = os.path.join(
704
+ config_file_paths = []
705
+
706
+ # note that we prioritize user defined config
707
+ user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER
708
+ if user_defined_config_folder is not None:
709
+ user_defined_config_file_path = os.path.join(
710
+ user_defined_config_folder, json_file_name)
711
+ config_file_paths.append(user_defined_config_file_path)
712
+
713
+ default_config_file_path = os.path.join(
714
os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
706
- if os.path.exists(config_file_path):
707
- with open(config_file_path) as f:
708
- logger.info("Using configuration from %s for MoE layer.",
709
- config_file_path)
710
- # If a configuration has been found, return it
711
- return {int(key): val for key, val in json.load(f).items()}
715
+ config_file_paths.append(default_config_file_path)
716
+
717
+ for config_file_path in config_file_paths:
718
+ if os.path.exists(config_file_path):
719
+ with open(config_file_path) as f:
720
+ logger.info("Using configuration from %s for MoE layer.",
721
+ config_file_path)
722
+ # If a configuration has been found, return it
723
+ return {int(key): val for key, val in json.load(f).items()}
724
725
# If no optimized configuration is available, we will use the default
726
# configuration
727
logger.warning(
728
("Using default MoE config. Performance might be sub-optimal! "
717
- "Config file not found at %s"), config_file_path)
729
+ "Config file not found at %s"), config_file_paths)
730
return None
731
732