From f7f62d19160010bdc84b31af407bb8b88e5392d6 Mon Sep 17 00:00:00 2001 From: pmanczak Date: Wed, 7 Oct 2026 18:02:13 +0200 Subject: [PATCH] [XPU][MoE] Tune Triton fused MoE for Intel XPU (#53065) --- benchmarks/kernels/benchmark_moe.py | 47 ++++- ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 164 ++++++++++++++++++ ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 164 ++++++++++++++++++ ...ame=Intel(R)_Arc(TM)_Pro_B70_Graphics.json | 56 ++++++ 4 files changed, 427 insertions(+), 4 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json create mode 100644 vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 1fcc7ddde7..f746e9aab3 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -303,9 +303,13 @@ def benchmark_config( run() torch.accelerator.synchronize() - # Capture 10 invocations with CUDA graph - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): + # Capture 10 invocations via graph(); CUDA needs a side stream + graph = ( + torch.cuda.CUDAGraph() + if current_platform.is_cuda_alike() + else torch.xpu.XPUGraph() + ) + with current_platform.graph(graph): for _ in range(10): run() torch.accelerator.synchronize() @@ -361,11 +365,33 @@ def get_rocm_tuning_space(use_fp16): return param_ranges +def get_xpu_tuning_space(use_fp16): + # BLOCK_SIZE_N=32 and num_warps=32 never won; grf_mode=256 beat 128 on B70. + block_m_range = [16, 32, 64, 128, 256] + block_n_range = [64, 128] + block_k_range = [32, 64] + num_warps_range = [4, 8, 16] + group_m_range = [1, 8, 16] + num_stage_range = [3, 4] + + return { + "BLOCK_SIZE_M": block_m_range, + "BLOCK_SIZE_N": block_n_range, + "BLOCK_SIZE_K": block_k_range, + "GROUP_SIZE_M": group_m_range, + "num_warps": num_warps_range, + "num_stages": num_stage_range, + "grf_mode": ["256"], + } + + def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int]]: configs: list[BenchmarkConfig] = [] if current_platform.is_rocm(): param_ranges = get_rocm_tuning_space(use_fp16) + elif current_platform.is_xpu(): + param_ranges = get_xpu_tuning_space(use_fp16) else: # Reduced search space for faster tuning. # TODO(woosuk): Increase the search space and use a performance model to @@ -408,6 +434,14 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int return configs +def prune_xpu_search_space(num_tokens, search_space, topk, num_experts): + # Cap the tile by rows routed to one expert, not total rows: wider tiles + # just pad and are the slowest to compile. + rows_per_expert = (num_tokens * topk + num_experts - 1) // num_experts + max_block_m = max(16, triton.next_power_of_2(rows_per_expert)) + return [c for c in search_space if c["BLOCK_SIZE_M"] <= max_block_m] + + def prune_rocm_search_space( num_tokens, shard_intermediate_size, hidden_size, search_space, is_fp16, topk ): @@ -522,7 +556,7 @@ def merge_unique_dicts(list1, list2): @ray.remote(num_gpus=1) class BenchmarkWorker: def __init__(self, seed: int) -> None: - torch.set_default_device("cuda") + torch.set_default_device(current_platform.device_type) set_random_seed(seed) self.seed = seed # Get the device ID to allocate tensors and kernels @@ -619,6 +653,10 @@ class BenchmarkWorker: is_fp16, topk, ) + elif current_platform.is_xpu(): + search_space = prune_xpu_search_space( + num_tokens, search_space, topk, num_experts + ) need_device_guard = False if current_platform.is_rocm(): @@ -690,6 +728,7 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: else {} ), **({"kpack": config["kpack"]} if "kpack" in config else {}), + **({"grf_mode": config["grf_mode"]} if "grf_mode" in config else {}), **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}), } diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 0000000000..6c1788d496 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 0000000000..0cdfc22664 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,164 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 4, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 4, + "grf_mode": "256" + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + } +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json b/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json new file mode 100644 index 0000000000..46bee68b89 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json @@ -0,0 +1,56 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "128": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "512": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3, + "grf_mode": "256" + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 16, + "num_stages": 3, + "grf_mode": "256" + } +}