minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Triton fused_moe_kernel support ep moe tuning (#12343)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-10-29 23:16:09 +08:00
committed by GitHub
parent 400bddf24c
commit 52694b60da
4 changed files with 350 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
+101 -28
View File
@@ -50,7 +50,10 @@ def benchmark_config(
per_channel_quant: bool,
block_shape: List[int] = None,
num_iters: int = 100,
ep_size: int = 1,
) -> float:
# In EP mode, each rank only handles a subset of experts
local_experts = num_experts // ep_size
init_dtype = torch.float16 if use_fp8_w8a8 else dtype
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
if use_int8_w8a16 or use_int8_w8a8:
@@ -58,7 +61,7 @@ def benchmark_config(
-127,
127,
(
num_experts,
local_experts,
shard_intermediate_size,
hidden_size,
),
@@ -68,7 +71,7 @@ def benchmark_config(
-127,
127,
(
num_experts,
local_experts,
hidden_size,
shard_intermediate_size // 2,
),
@@ -76,12 +79,14 @@ def benchmark_config(
)
else:
w1 = torch.randn(
num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype
local_experts, shard_intermediate_size, hidden_size, dtype=init_dtype
)
w2 = torch.randn(
num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype
local_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype
)
gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32)
gating_output = torch.randn(
num_iters, num_tokens, local_experts, dtype=torch.float32
)
w1_scale = None
w2_scale = None
@@ -89,18 +94,18 @@ def benchmark_config(
a2_scale = None
if use_int8_w8a16:
w1_scale = torch.randn(
(num_experts, 2 * shard_intermediate_size), dtype=torch.float32
(local_experts, 2 * shard_intermediate_size), dtype=torch.float32
)
w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32)
w2_scale = torch.randn((hidden_size, local_experts), dtype=torch.float32)
if use_fp8_w8a8 or use_int8_w8a8:
if use_int8_w8a8 and block_shape is None:
w1_scale = torch.randn(
num_experts, shard_intermediate_size, dtype=torch.float32
local_experts, shard_intermediate_size, dtype=torch.float32
)
w2_scale = torch.randn(num_experts, hidden_size, dtype=torch.float32)
w2_scale = torch.randn(local_experts, hidden_size, dtype=torch.float32)
elif block_shape is None:
w1_scale = torch.randn(num_experts, dtype=torch.float32)
w2_scale = torch.randn(num_experts, dtype=torch.float32)
w1_scale = torch.randn(local_experts, dtype=torch.float32)
w2_scale = torch.randn(local_experts, dtype=torch.float32)
a1_scale = torch.randn(1, dtype=torch.float32)
a2_scale = torch.randn(1, dtype=torch.float32)
else:
@@ -110,17 +115,17 @@ def benchmark_config(
k_tiles_w1 = (hidden_size + block_k - 1) // block_k
k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k
w1_scale = torch.rand(
(num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32
(local_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32
)
w2_scale = torch.rand(
(num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32
(local_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32
)
if use_fp8_w8a8:
w1 = w1.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn)
w2 = w2.to(torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn)
input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32)
input_gating = torch.randn(num_tokens, local_experts, dtype=torch.float32)
topk_config = TopKConfig(
top_k=topk,
renormalize=True,
@@ -265,6 +270,7 @@ class BenchmarkWorker:
use_int8_w8a16: bool,
per_channel_quant: bool,
block_shape: List[int],
ep_size: int = 1,
) -> Tuple[Dict[str, int], float]:
torch.cuda.manual_seed_all(0)
dtype_str = get_config_dtype_str(
@@ -272,10 +278,12 @@ class BenchmarkWorker:
)
# NOTE(woosuk): The current naming convention uses w2.shape[2], which
# is the intermediate size after silu_and_mul.
# For EP mode, use local expert count for config lookup
local_experts = num_experts // ep_size
block_n = block_shape[0] if block_shape else 0
block_k = block_shape[1] if block_shape else 0
op_config = get_moe_configs(
num_experts,
local_experts,
shard_intermediate_size // 2,
dtype_str,
block_n,
@@ -285,7 +293,7 @@ class BenchmarkWorker:
if op_config is None:
config = get_default_config(
num_tokens,
num_experts,
local_experts,
shard_intermediate_size,
hidden_size,
topk,
@@ -309,6 +317,7 @@ class BenchmarkWorker:
use_int8_w8a16,
per_channel_quant,
block_shape,
ep_size,
)
return config, kernel_time
@@ -326,6 +335,7 @@ class BenchmarkWorker:
per_channel_quant: bool,
block_shape: List[int],
search_space: List[Dict[str, int]],
ep_size: int = 1,
) -> Dict[str, int]:
best_config = None
best_time = float("inf")
@@ -346,6 +356,7 @@ class BenchmarkWorker:
per_channel_quant,
block_shape,
num_iters=10,
ep_size=ep_size,
)
except (triton.runtime.autotuner.OutOfResources, RuntimeError):
# Some configurations may be invalid and fail to compile.
@@ -395,6 +406,7 @@ def get_filename(
use_int8_w8a16: bool,
per_channel_quant: bool,
block_shape: List[int],
ep_size: int = 1,
) -> None:
dtype_str = get_config_dtype_str(
dtype,
@@ -405,9 +417,11 @@ def get_filename(
# NOTE(woosuk): The current naming convention uses w2.shape[2], which
# is the intermediate size after silu_and_mul.
# For EP mode, we use local expert count instead of total expert count
local_experts = num_experts // ep_size
filename = get_config_file_name(
num_experts,
shard_intermediate_size // 2,
local_experts,
shard_intermediate_size if ep_size > 1 else shard_intermediate_size // 2,
dtype_str,
block_shape,
per_channel_quant,
@@ -419,17 +433,35 @@ def get_filename(
def main(args: argparse.Namespace):
print(args)
# Check EP mode constraint: tp_size must be 1 when ep_size > 1
if args.ep_size > 1 and args.tp_size != 1:
raise ValueError(
f"When using Expert Parallelism (ep_size={args.ep_size}), "
f"tp_size must be set to 1, but got tp_size={args.tp_size}. "
f"Please set --tp-size 1 when using --ep-size > 1."
)
config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
if config.architectures[0] == "DbrxForCausalLM":
E = config.ffn_config.moe_num_experts
topk = config.ffn_config.moe_top_k
intermediate_size = config.ffn_config.ffn_hidden_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] == "JambaForCausalLM":
E = config.num_experts
topk = config.num_experts_per_tok
intermediate_size = config.intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] in [
"Qwen2MoeForCausalLM",
"Qwen3MoeForCausalLM",
@@ -438,7 +470,12 @@ def main(args: argparse.Namespace):
E = config.num_experts
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]:
E = (
config.n_routed_experts + (0 if args.disable_shared_experts_fusion else 1)
@@ -447,14 +484,24 @@ def main(args: argparse.Namespace):
)
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] == "Llama4ForConditionalGeneration":
E = config.text_config.num_local_experts + (
0 if args.disable_shared_experts_fusion else 1
)
topk = config.text_config.num_experts_per_tok
intermediate_size = config.text_config.intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] in [
"Grok1ForCausalLM",
"Grok1ImgGen",
@@ -463,7 +510,12 @@ def main(args: argparse.Namespace):
E = config.num_local_experts
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] in [
"BailingMoEForCausalLM",
"BailingMoeForCausalLM",
@@ -472,18 +524,33 @@ def main(args: argparse.Namespace):
E = config.num_experts
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
elif config.architectures[0] in ["Glm4MoeForCausalLM"]:
E = config.n_routed_experts
topk = config.num_experts_per_tok
intermediate_size = config.moe_intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
else:
# Default: Mixtral
E = config.num_local_experts
topk = config.num_experts_per_tok
intermediate_size = config.intermediate_size
shard_intermediate_size = 2 * intermediate_size // args.tp_size
# In EP mode, use original intermediate_size; otherwise apply TP sharding
shard_intermediate_size = (
intermediate_size
if args.ep_size > 1
else 2 * intermediate_size // args.tp_size
)
hidden_size = getattr(config, "hidden_size", None) or config.text_config.hidden_size
dtype = config.torch_dtype
@@ -559,6 +626,7 @@ def main(args: argparse.Namespace):
use_int8_w8a16,
per_channel_quant,
block_shape,
args.ep_size,
)
print(
f"Start tuning over {len(search_space)} configurations to create {filename}..."
@@ -581,6 +649,7 @@ def main(args: argparse.Namespace):
per_channel_quant,
block_shape,
search_space,
args.ep_size,
)
for batch_size in batch_sizes
],
@@ -610,6 +679,7 @@ def main(args: argparse.Namespace):
use_int8_w8a16,
per_channel_quant,
block_shape,
args.ep_size,
)
for batch_size in batch_sizes
],
@@ -625,7 +695,7 @@ if __name__ == "__main__":
parser.add_argument(
"--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1"
)
parser.add_argument("--tp-size", "--tp", type=int, default=2)
parser.add_argument("--tp-size", "--tp", type=int, default=1)
parser.add_argument(
"--dtype",
type=str,
@@ -640,6 +710,9 @@ if __name__ == "__main__":
parser.add_argument("--batch-size", type=int, required=False)
parser.add_argument("--tune", action="store_true")
parser.add_argument("--disable-shared-experts-fusion", action="store_true")
parser.add_argument(
"--ep-size", "--ep", type=int, default=1, help="Expert parallelism size"
)
args = parser.parse_args()
main(args)