iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[XPU][MoE] Tune Triton fused MoE for Intel XPU (#53065)

Browse Source
This commit is contained in:
pmanczak
2026-10-08 00:02:13 +08:00
committed by GitHub
parent 9367d8b969
commit f7f62d1916
4 changed files with 427 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_moe.py
+43 -4
View File
@@ -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 {}),
}
vllm/model_executor/layers/fused_moe/configs/E=128,N=384,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json
+164
View File
@@ -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"
}
}
vllm/model_executor/layers/fused_moe/configs/E=64,N=1024,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json
+164
View File
@@ -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"
}
}
vllm/model_executor/layers/fused_moe/configs/E=64,N=512,device_name=Intel(R)_Arc(TM)_Pro_B70_Graphics.json
+56
View File
@@ -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"
}
}