diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py index 0167bf7a0..a74972c33 100644 --- a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py +++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py @@ -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) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 7548086dd..4da5449f9 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -138,6 +138,7 @@ class TritonRunnerCore(MoeRunnerCore): running_state["config"], running_state.get("down_config"), running_state.get("down_moe_use_tma", False), + running_state.get("up_moe_use_tma", False), b1=quant_info.b13, b2=quant_info.b2, use_fp8_w8a8=quant_info.use_fp8_w8a8, @@ -280,6 +281,7 @@ def pre_permute_standard_to_triton( config, down_config, down_moe_use_tma, + up_moe_use_tma, sorted_token_ids, expert_ids, num_tokens_post_padded, @@ -299,6 +301,7 @@ def pre_permute_standard_to_triton( running_state["config"] = config running_state["down_config"] = down_config running_state["down_moe_use_tma"] = down_moe_use_tma + running_state["up_moe_use_tma"] = up_moe_use_tma return TritonRunnerInput( hidden_states=hidden_states, diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json new file mode 100644 index 000000000..5f70c4971 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": false + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": false + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": false + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": false + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": false + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": true + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": false + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": true + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5, + "USE_TMA": true + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4, + "USE_TMA": true + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 4, + "USE_TMA": false + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": false + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4, + "USE_TMA": false + } +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128]_down.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128]_down.json new file mode 100644 index 000000000..7d5ca2e74 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128]_down.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3, + "USE_TMA": true + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 4, + "num_stages": 2, + "USE_TMA": true + } +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index 1db3f9470..cd9338916 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -7,6 +7,7 @@ from __future__ import annotations import functools +import logging from typing import TYPE_CHECKING, Any, Dict, List, Optional import torch @@ -93,6 +94,8 @@ if not _is_cuda and not _is_hip and not _is_xpu: padding_size = get_moe_padding_size(_use_aiter) +logger = logging.getLogger(__name__) + def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool: return num_tokens <= 32 and not is_batch_invariant_mode_enabled() @@ -363,7 +366,7 @@ def swiglu_no_interleaved_with_alpha_and_limit(x, gemm1_alpha, gemm1_limit): @functools.lru_cache() -def _down_moe_use_tma(): +def _moe_support_tma(): return support_tensor_descriptor() @@ -409,11 +412,26 @@ def _prepare_fused_moe_run( per_channel_quant=per_channel_quant, return_down_config=True, ) - down_moe_use_tma = ( - _down_moe_use_tma() - and down_config is not None - and down_config.pop("USE_TMA", False) - ) + # Copy config to avoid mutating the lru_cached dict returned by + # get_moe_configs; we pop USE_TMA below. + config = dict(config) + # Up-projection TMA is opt-in: only enabled when the up config file + # explicitly carries "USE_TMA": true (produced by tuning). By default the + # existing up config files do not contain this key, so existing users are + # unaffected unless they re-tune with the updated script. + up_tma_requested = config.pop("USE_TMA", False) + up_moe_use_tma = _moe_support_tma() and up_tma_requested + if up_moe_use_tma: + logger.warning_once( + "Up MoE TMA is enabled (USE_TMA=true in the up-projection config). " + "This requires a config produced by the updated tuning script. " + ) + down_tma_requested = down_config is not None and down_config.pop("USE_TMA", False) + down_moe_use_tma = _moe_support_tma() and down_tma_requested + if down_moe_use_tma: + logger.warning_once( + "Down MoE TMA is enabled (USE_TMA=true in the down-projection config)." + ) sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( topk_ids, config["BLOCK_SIZE_M"], E @@ -423,6 +441,7 @@ def _prepare_fused_moe_run( config, down_config, down_moe_use_tma, + up_moe_use_tma, sorted_token_ids, expert_ids, num_tokens_post_padded, @@ -441,6 +460,7 @@ def _fused_moe_kernel_sequence( config: Dict[str, Any], down_config: Optional[Dict[str, Any]], down_moe_use_tma: bool, + up_moe_use_tma: bool, *, b1: Optional[torch.Tensor], b2: Optional[torch.Tensor], @@ -568,6 +588,7 @@ def _fused_moe_kernel_sequence( per_channel_quant=per_channel_quant, block_shape=block_shape, c_sorted=down_moe_use_tma, + b_use_tma=up_moe_use_tma, filter_expert=filter_expert, ) @@ -927,6 +948,7 @@ def fused_experts_impl( config, down_config, down_moe_use_tma, + up_moe_use_tma, sorted_token_ids, expert_ids, num_tokens_post_padded, @@ -955,6 +977,7 @@ def fused_experts_impl( config, down_config, down_moe_use_tma, + up_moe_use_tma, b1=b1, b2=b2, use_fp8_w8a8=use_fp8_w8a8,