Skip to content

Commit e749da7

Browse files
Fridge003almaslof
authored andcommitted
Fallback to lower triton version for unfound fused moe configs (sgl-project#7013)
1 parent 3cd75fd commit e749da7

1 file changed

Lines changed: 21 additions & 3 deletions

File tree

python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -983,6 +983,8 @@ def get_moe_configs(
983983
kernel on a given batch size bs, the closest batch size in the grid should
984984
be picked and the associated configuration chosen to invoke the kernel.
985985
"""
986+
# Supported Triton versions, should be sorted from the newest to the oldest
987+
supported_triton_versions = ["3.3.1", "3.2.0", "3.1.0"]
986988

987989
# First look up if an optimized configuration is available in the configs
988990
# directory
@@ -1005,12 +1007,28 @@ def get_moe_configs(
10051007
# For example, updating the Triton version might cause all old configs to become suboptimal.
10061008
# To achieve the best performance, consider re-tuning the Triton fused MOE kernel in your environment.
10071009
# For the tuning method, refer to: https://github.com/sgl-project/sglang/tree/main/benchmark/kernels/fused_moe_triton
1008-
log_info_on_rank0(
1009-
logger, f"Using MoE kernel config from {config_file_path}."
1010-
)
1010+
logger.info(f"Using MoE kernel config from {config_file_path}.")
10111011
# If a configuration has been found, return it
10121012
return {int(key): val for key, val in json.load(f).items()}
10131013

1014+
# Searching for other triton versions that supports the same config
1015+
for try_triton_version in supported_triton_versions:
1016+
if try_triton_version == triton_version:
1017+
continue
1018+
try_config_file_path = os.path.join(
1019+
os.path.dirname(os.path.realpath(__file__)),
1020+
"configs",
1021+
f"triton_{try_triton_version.replace('.', '_')}",
1022+
json_file_name,
1023+
)
1024+
if os.path.exists(try_config_file_path):
1025+
with open(try_config_file_path) as f:
1026+
logger.warning(
1027+
f"Config file not found at {config_file_path}. Fallback to triton version {try_triton_version} and use MoE kernel config from {try_config_file_path}. Performance might be sub-optimal!",
1028+
)
1029+
# If a configuration has been found, return it
1030+
return {int(key): val for key, val in json.load(f).items()}
1031+
10141032
# If no optimized configuration is available, we will use the default
10151033
# configuration
10161034
logger.warning(

0 commit comments

Comments
 (0)