EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

fix: resolve tuning fused moe issue (#9587)

Browse Source
This commit is contained in:
Yineng Zhang
2025-08-25 01:41:15 -07:00
committed by GitHub
parent ebd9dbe71b
commit f8b757bcac
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
View File

@@ -22,7 +22,7 @@ from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
) )
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
from sglang.srt.layers.moe.topk import TopKConfig, select_experts from sglang.srt.layers.moe.topk import TopKConfig, select_experts
from sglang.srt.utils import is_hip, is_rocm from sglang.srt.utils import is_hip
_is_hip = is_hip() _is_hip = is_hip()
@@ -287,7 +287,7 @@ class BenchmarkWorker:
) )
else: else:
config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))] config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))]
with torch.cuda.device(self.device_id) if is_rocm() else nullcontext(): with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
kernel_time = benchmark_config( kernel_time = benchmark_config(
config, config,
num_tokens, num_tokens,
@@ -319,7 +319,7 @@ class BenchmarkWorker:
) -> Dict[str, int]: ) -> Dict[str, int]:
best_config = None best_config = None
best_time = float("inf") best_time = float("inf")
with torch.cuda.device(self.device_id) if is_rocm() else nullcontext(): with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for config in tqdm(search_space): for config in tqdm(search_space):
try: try:
kernel_time = benchmark_config( kernel_time = benchmark_config(