[FEATURE] support custom vllm tuned config path for fused moe triton kernels (#22791)
Signed-off-by: Chi Zhang <zhangchi.usc1992@bytedance.com>
This commit is contained in:
@@ -701,20 +701,32 @@ def get_moe_configs(
|
||||
block_shape = [block_n, block_k] if block_n and block_k else None
|
||||
json_file_name = get_config_file_name(E, N, dtype, block_shape)
|
||||
|
||||
config_file_path = os.path.join(
|
||||
config_file_paths = []
|
||||
|
||||
# note that we prioritize user defined config
|
||||
user_defined_config_folder = envs.VLLM_TUNED_CONFIG_FOLDER
|
||||
if user_defined_config_folder is not None:
|
||||
user_defined_config_file_path = os.path.join(
|
||||
user_defined_config_folder, json_file_name)
|
||||
config_file_paths.append(user_defined_config_file_path)
|
||||
|
||||
default_config_file_path = os.path.join(
|
||||
os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
|
||||
if os.path.exists(config_file_path):
|
||||
with open(config_file_path) as f:
|
||||
logger.info("Using configuration from %s for MoE layer.",
|
||||
config_file_path)
|
||||
# If a configuration has been found, return it
|
||||
return {int(key): val for key, val in json.load(f).items()}
|
||||
config_file_paths.append(default_config_file_path)
|
||||
|
||||
for config_file_path in config_file_paths:
|
||||
if os.path.exists(config_file_path):
|
||||
with open(config_file_path) as f:
|
||||
logger.info("Using configuration from %s for MoE layer.",
|
||||
config_file_path)
|
||||
# If a configuration has been found, return it
|
||||
return {int(key): val for key, val in json.load(f).items()}
|
||||
|
||||
# If no optimized configuration is available, we will use the default
|
||||
# configuration
|
||||
logger.warning(
|
||||
("Using default MoE config. Performance might be sub-optimal! "
|
||||
"Config file not found at %s"), config_file_path)
|
||||
"Config file not found at %s"), config_file_paths)
|
||||
return None
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user