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

[kernel] add triton moe TMA up support (#33559)

Browse Source 
Co-authored-by: undefined <zhouchen.arrebol@jd.com>
Co-authored-by: xq25478 <xq25478@qq.com>
Co-authored-by: xieminghe.simon <xieminghe.simon@jd.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
xieminghe1
2026-08-13 13:29:32 +08:00
committed by GitHub
co-authored by undefined xq25478 xieminghe.simon Xiaoyu Zhang
parent 69bf601e3c
commit ef7208d41d
5 changed files with 669 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py
+309 -80
View File
@@ -143,6 +143,9 @@ def benchmark_config(
block_shape: List[int] = None,
ep_size: int = 1,
num_iters: int = 100,
enable_up_tma: bool = False,
tune_round: str = "both",
down_use_tma_map: dict = None,
) -> float:
ncu_enable = os.getenv("NCU_ENABLE", "0") == "1"
if ncu_enable:
@@ -312,14 +315,29 @@ def benchmark_config(
moe_inputs[k].expert_ids.copy_(expert_ids_)
moe_inputs[k].num_tokens_post_padded.copy_(num_tokens_post_padded_)
def get_kernel_wrapper(moe_use_tma, inner_iter, use_cuda_graph):
compute_type = (
tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
)
moe_runner_config = MoeRunnerConfig(
inplace=True,
)
apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input
compute_type = tl.bfloat16 if hidden_states.dtype == torch.bfloat16 else tl.float16
moe_runner_config = MoeRunnerConfig(
inplace=True,
)
apply_router_weight_on_input = moe_runner_config.apply_router_weight_on_input
use_cuda_graph = True if not ncu_enable else False
# Determine which kernels to build based on tune_round:
# "both" — up (c_sorted=False, no TMA) + down (no-tma + tma) [default, enable_up_tma=False]
# "down" — down (no-tma + tma) only [round 1 of two-round tune]
# "up" — up (no-tma + tma) only, c_sorted from down_use_tma_map [round 2 of two-round tune]
build_up = tune_round in ("both", "up")
build_down = tune_round in ("both", "down")
# For "up" round, c_sorted must match what down TMA decided at runtime
if tune_round == "up":
c_sorted = down_use_tma_map.get(config["BLOCK_SIZE_M"], False)
else:
c_sorted = False
up_kernels = []
if build_up:
kernel0 = KernelWrapper(
A=hidden_states,
B=w1,
@@ -340,12 +358,44 @@ def benchmark_config(
use_int4_w4a16=use_int4_w4a16,
per_channel_quant=False,
block_shape=block_shape,
b_use_tma=moe_use_tma,
c_sorted=moe_use_tma,
b_use_tma=False,
c_sorted=c_sorted,
filter_expert=False,
use_cuda_graph=use_cuda_graph,
inner_iter=inner_iter,
)
up_kernels.append(kernel0)
if enable_up_tma or tune_round == "up":
kernel0_tma = KernelWrapper(
A=hidden_states,
B=w1,
bias=None,
C=intermediate_cache1,
A_scale=a1_scale,
B_scale=w1_scale,
B_zp=None,
topk_weights=topk_output_.topk_weights,
moe_inputs=moe_inputs,
mul_routed_weight=apply_router_weight_on_input,
top_k=topk,
config=config,
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=use_int8_w8a8,
use_int8_w8a16=use_int8_w8a16,
use_int4_w4a16=use_int4_w4a16,
per_channel_quant=False,
block_shape=block_shape,
b_use_tma=True,
c_sorted=c_sorted,
filter_expert=False,
use_cuda_graph=use_cuda_graph,
inner_iter=inner_iter,
)
up_kernels.append(kernel0_tma)
down_kernels = []
if build_down:
kernel1 = KernelWrapper(
A=intermediate_cache2,
B=w2,
@@ -366,25 +416,45 @@ def benchmark_config(
use_int4_w4a16=use_int4_w4a16,
per_channel_quant=False,
block_shape=block_shape,
a_use_tma=moe_use_tma,
b_use_tma=moe_use_tma,
a_use_tma=False,
b_use_tma=False,
filter_expert=False,
use_cuda_graph=use_cuda_graph,
inner_iter=inner_iter,
)
return kernel0, kernel1
use_cuda_graph = True if not ncu_enable else False
kernel0, kernel1 = get_kernel_wrapper(False, inner_iter, use_cuda_graph)
kernel_tma0, kernel_tma1 = get_kernel_wrapper(True, inner_iter, use_cuda_graph)
down_kernels.append(kernel1)
kernel1_tma = KernelWrapper(
A=intermediate_cache2,
B=w2,
bias=None,
C=intermediate_cache3,
A_scale=a2_scale,
B_scale=w2_scale,
B_zp=None,
topk_weights=topk_output_.topk_weights,
moe_inputs=moe_inputs,
mul_routed_weight=not apply_router_weight_on_input,
top_k=1,
config=config,
compute_type=compute_type,
use_fp8_w8a8=use_fp8_w8a8,
use_int8_w8a8=use_int8_w8a8,
use_int8_w8a16=use_int8_w8a16,
use_int4_w4a16=use_int4_w4a16,
per_channel_quant=False,
block_shape=block_shape,
a_use_tma=True,
b_use_tma=True,
filter_expert=False,
use_cuda_graph=use_cuda_graph,
inner_iter=inner_iter,
)
down_kernels.append(kernel1_tma)
# JIT compilation & warmup
if not ncu_enable:
kernel0.forward_cost()
kernel1.forward_cost()
kernel_tma0.forward_cost()
kernel_tma1.forward_cost()
for k in up_kernels + down_kernels:
k.forward_cost()
ts0 = []
ts1 = []
@@ -393,31 +463,42 @@ def benchmark_config(
for i in range(num_iters // inner_iter):
prepare(i, inner_iter)
ts0.append(kernel0.forward_cost())
ts1.append(kernel1.forward_cost())
ts_tma0.append(kernel_tma0.forward_cost())
ts_tma1.append(kernel_tma1.forward_cost())
if build_up:
ts0.append(kernel0.forward_cost()) # up no-tma
if len(up_kernels) > 1:
ts_tma0.append(kernel0_tma.forward_cost()) # up tma
if build_down:
ts1.append(kernel1.forward_cost()) # down no-tma
ts_tma1.append(kernel1_tma.forward_cost()) # down tma
torch.cuda.synchronize()
avg = sum(ts0) / (num_iters) * 1000 # us
avg1 = sum(ts1) / (num_iters) * 1000 # us
avg_tma = sum(ts_tma0) / (num_iters) * 1000 # us
avg1_tma = sum(ts_tma1) / (num_iters) * 1000 # us
avg = sum(ts0) / (num_iters) * 1000 if ts0 else float("inf")
avg1 = sum(ts1) / (num_iters) * 1000 if ts1 else float("inf")
avg_tma = sum(ts_tma0) / (num_iters) * 1000 if ts_tma0 else float("inf")
avg1_tma = sum(ts_tma1) / (num_iters) * 1000 if ts_tma1 else float("inf")
return avg, avg_tma, avg1, avg1_tma
class BestConfigTrace:
def __init__(self, name, down_moe=False):
def __init__(self, name, down_moe=False, enable_up_tma=False):
self.name = name
self.down_moe = down_moe
self.enable_up_tma = enable_up_tma
self.best_costs_m = {} # block_m: best_cost
def update(self, config, time_cost_all):
block_m = config["BLOCK_SIZE_M"]
if not self.down_moe:
time_cost = time_cost_all[0]
# time_cost_all = (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma)
if self.enable_up_tma:
# For up_proj, pick the faster of TMA vs no-TMA.
time_cost = min(time_cost_all[0], time_cost_all[1])
else:
# Up TMA not enabled — always use no-TMA cost.
time_cost = time_cost_all[0]
else:
# For down_proj, pick the faster of TMA vs no-TMA.
time_cost = min(time_cost_all[2], time_cost_all[3])
if (
block_m not in self.best_costs_m
@@ -436,7 +517,16 @@ class BestConfigTrace:
return {}
config, _, time_cost_all = self.best_costs_m[block_m]
if not self.down_moe:
return config
if self.enable_up_tma:
# up_proj: use TMA when the TMA variant is faster.
return {
**config,
"USE_TMA": time_cost_all[0] > time_cost_all[1],
}
else:
# Up TMA not enabled — do not add USE_TMA key so the runtime
# defaults to no-TMA for up-projection.
return config
else:
return {
**config,
@@ -471,26 +561,72 @@ class BenchmarkWorker:
cfg: Dict[str, int],
topk_ids_dir: str,
ep_size: int = 1,
enable_up_tma: bool = False,
) -> Tuple[Dict[str, int], float]:
torch.cuda.manual_seed_all(0)
topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)]
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
kernel_time = benchmark_config(
cfg,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
)
if enable_up_tma:
# Two-step: first measure down to determine c_sorted,
# then measure up with the correct c_sorted.
_, _, kt1_no_tma, kt1_tma = benchmark_config(
cfg,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
enable_up_tma=True,
tune_round="down",
)
down_use_tma = kt1_no_tma > kt1_tma
kt0_no_tma, kt0_tma, _, _ = benchmark_config(
cfg,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
enable_up_tma=True,
tune_round="up",
down_use_tma_map={cfg["BLOCK_SIZE_M"]: down_use_tma},
)
kernel_time = (kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma)
else:
kernel_time = benchmark_config(
cfg,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
enable_up_tma=False,
tune_round="both",
)
return cfg, kernel_time
def tune(
@@ -509,42 +645,124 @@ class BenchmarkWorker:
search_space: List[Dict[str, int]],
topk_ids_dir: str,
ep_size: int = 1,
enable_up_tma: bool = False,
) -> Dict[str, int]:
trace0 = BestConfigTrace("kernel0", down_moe=False)
trace1 = BestConfigTrace("kernel1", down_moe=True)
topk_ids_list = [load_topk_ids(topk_ids_dir, i) for i in range(100)]
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for config in tqdm(search_space):
try:
kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma = benchmark_config(
if not enable_up_tma:
# Default path: single round, up c_sorted=False, no up TMA.
# Down TMA is still tuned.
trace0 = BestConfigTrace("kernel0", down_moe=False, enable_up_tma=False)
trace1 = BestConfigTrace("kernel1", down_moe=True)
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for config in tqdm(search_space):
try:
kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma = benchmark_config(
config,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
num_iters=100,
enable_up_tma=False,
tune_round="both",
)
except triton.runtime.autotuner.OutOfResources:
continue
trace0.update(
config,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
num_iters=100,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
trace1.update(
config,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
else:
# Two-round coupled tuning: first tune down to determine c_sorted,
# then tune up with the correct c_sorted that matches runtime.
trace0 = BestConfigTrace("kernel0", down_moe=False, enable_up_tma=True)
trace1 = BestConfigTrace("kernel1", down_moe=True)
# === Round 1: Down-only ===
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for config in tqdm(search_space, desc="Round 1 (down)"):
try:
_, _, kt1_no_tma, kt1_tma = benchmark_config(
config,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
num_iters=100,
enable_up_tma=True,
tune_round="down",
)
except triton.runtime.autotuner.OutOfResources:
continue
trace1.update(
config,
(float("inf"), float("inf"), kt1_no_tma, kt1_tma),
)
# Extract down TMA decision per BLOCK_SIZE_M from round 1 results
down_use_tma_map = {}
for block_m, (_, _, time_cost_all) in trace1.best_costs_m.items():
down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3]
print(
f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: "
f"{down_use_tma_map}"
)
# === Round 2: Up with c_sorted from round 1 ===
with torch.cuda.device(self.device_id) if is_hip() else nullcontext():
for config in tqdm(search_space, desc="Round 2 (up)"):
try:
kt0_no_tma, kt0_tma, _, _ = benchmark_config(
config,
num_tokens,
num_experts,
shard_intermediate_size,
hidden_size,
topk,
dtype,
use_fp8_w8a8,
use_int8_w8a8,
use_int8_w8a16,
use_int4_w4a16,
topk_ids_list,
block_shape,
ep_size=ep_size,
num_iters=100,
enable_up_tma=True,
tune_round="up",
down_use_tma_map=down_use_tma_map,
)
except triton.runtime.autotuner.OutOfResources:
continue
trace0.update(
config,
(kt0_no_tma, kt0_tma, float("inf"), float("inf")),
)
except triton.runtime.autotuner.OutOfResources:
# Some configurations may be invalid and fail to compile.
continue
trace0.update(
config,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
trace1.update(
config,
(kt0_no_tma, kt0_tma, kt1_no_tma, kt1_tma),
)
now = datetime.now()
print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}")
@@ -731,6 +949,7 @@ def main(args: argparse.Namespace):
search_space,
topk_ids_dir,
args.ep_size,
enable_up_tma=args.enable_tune_up_tma,
)
else:
cfg = {
@@ -757,6 +976,7 @@ def main(args: argparse.Namespace):
cfg,
topk_ids_dir,
args.ep_size,
enable_up_tma=args.enable_tune_up_tma,
)
print(f"{t0=}, {t0_tma=}, {t1=}, {t1_tma=}")
return
@@ -823,6 +1043,7 @@ def main(args: argparse.Namespace):
search_space,
topk_ids_dir,
args.ep_size,
args.enable_tune_up_tma,
)
for batch_size in batch_sizes
],
@@ -890,6 +1111,14 @@ if __name__ == "__main__":
parser.add_argument("--configs", type=int, nargs="+", required=False)
parser.add_argument("--topk-ids-dir", type=str, required=True)
parser.add_argument("--cmp-configs", type=str, nargs="+", required=False)
parser.add_argument(
"--enable-tune-up-tma",
action="store_true",
help="Enable up-projection TMA tuning in addition to down-projection TMA. "
"When set, the up config file will contain a USE_TMA flag. "
"When not set (default), only down-projection TMA is tuned and the up "
"config will not contain a USE_TMA key.",
)
args = parser.parse_args()
main(args)