diff --git a/README.md b/README.md
index 864caa5..265eee4 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,11 @@
# xe-fuse
+
+
+[](https://scorecard.dev/viewer/?uri=github.com/IntelLabs/Xe-Fuse)
+
+> ⚠️ **Disclaimer**: This project is currently in active development. The code is **not stable** and **not intended for use in production environments**. Interfaces, features, and behaviors are subject to change without notice.
+
GEMM epilogue fusion framework for Intel Xe GPUs, built on sycl-tla.
Fuses memory-bound Transformer operations (RMSNorm, SwiGLU, RoPE, GeLU, residual-add, etc.) into GEMM epilogues — the ops execute on data still in registers from the accumulator, avoiding separate kernel launches and global memory round-trips.
diff --git a/assets/xe-fuse-logo.png b/assets/xe-fuse-logo.png
new file mode 100644
index 0000000..197b7ec
Binary files /dev/null and b/assets/xe-fuse-logo.png differ
diff --git a/autotune/generate_kernel.py b/autotune/generate_kernel.py
index 18f73c4..cae6e7f 100644
--- a/autotune/generate_kernel.py
+++ b/autotune/generate_kernel.py
@@ -44,8 +44,6 @@
import argparse
import json
-import os
-import sys
from datetime import datetime
from pathlib import Path
@@ -68,8 +66,8 @@
"",
"typename b::MulOp<>::Arguments mul_args{};",
"",
- "typename EVT::Arguments evt_args{accum_args, scale_args, mul_args};"
- ]
+ "typename EVT::Arguments evt_args{accum_args, scale_args, mul_args};",
+ ],
},
"k3": {
"name": "K3_RoPE",
@@ -77,7 +75,12 @@
"tile_shape": "_256, _256, _32",
"evt_typedefs": "using EVT = b::RoPE;",
"aux_data": [
- {"name": "cos_sin", "type": "float", "shape": "M * N * L", "init_seed": 2024}
+ {
+ "name": "cos_sin",
+ "type": "float",
+ "shape": "M * N * L",
+ "init_seed": 2024,
+ }
],
"evt_args": [
"typename b::AuxLoad::Arguments cs_args;",
@@ -88,8 +91,8 @@
"",
"typename xe_fuse::XeRoPECompute::Arguments rope_args{};",
"",
- "typename EVT::Arguments evt_args{cs_args, rope_args};"
- ]
+ "typename EVT::Arguments evt_args{cs_args, rope_args};",
+ ],
},
"k4": {
"name": "K4_RmsNormRoPE",
@@ -98,7 +101,12 @@
"evt_typedefs": "using EVT = b::RoPEComposed, float>;",
"aux_data": [
{"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42},
- {"name": "cos_sin", "type": "float", "shape": "M * N * L", "init_seed": 2024}
+ {
+ "name": "cos_sin",
+ "type": "float",
+ "shape": "M * N * L",
+ "init_seed": 2024,
+ },
],
"evt_args": [
"// child 0: ScaleRows (inner tree)",
@@ -118,8 +126,8 @@
" cute::Stride, int64_t>{}, make_shape(M, N, L));",
"",
"typename xe_fuse::XeRoPEComputeTwoChild::Arguments rope_args{};",
- "typename EVT::Arguments evt_args{rms_args, cs_args, rope_args};"
- ]
+ "typename EVT::Arguments evt_args{rms_args, cs_args, rope_args};",
+ ],
},
"k4v2": {
"name": "K4v2_RoPEScaled",
@@ -128,7 +136,12 @@
"evt_typedefs": "using EVT = b::RoPEScaled;",
"aux_data": [
{"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42},
- {"name": "cos_sin", "type": "float", "shape": "M * N * L", "init_seed": 2024}
+ {
+ "name": "cos_sin",
+ "type": "float",
+ "shape": "M * N * L",
+ "init_seed": 2024,
+ },
],
"evt_args": [
"// child 0: ColBroadcast",
@@ -147,8 +160,8 @@
"// root: XeRoPEScaledCompute (merged)",
"typename xe_fuse::XeRoPEScaledCompute::Arguments rope_args{};",
"",
- "typename EVT::Arguments evt_args{scale_args, cs_args, rope_args};"
- ]
+ "typename EVT::Arguments evt_args{scale_args, cs_args, rope_args};",
+ ],
},
"k0a": {
"name": "K0a_ResidualGamma",
@@ -156,8 +169,13 @@
"tile_shape": "_256, _256, _32",
"evt_typedefs": "using EVT = b::ScaleCols, TileShape, float>;",
"aux_data": [
- {"name": "residual", "type": "bf16", "shape": "M * N * L", "init_seed": 2021},
- {"name": "gamma", "type": "float", "shape": "N * L", "init_seed": 99}
+ {
+ "name": "residual",
+ "type": "bf16",
+ "shape": "M * N * L",
+ "init_seed": 2021,
+ },
+ {"name": "gamma", "type": "float", "shape": "N * L", "init_seed": 99},
],
"evt_args": [
"// RowBroadcast (child 0 of outer Mul)",
@@ -178,8 +196,8 @@
"",
"// Outer Mul",
"typename b::MulOp<>::Arguments mul_args{};",
- "typename EVT::Arguments evt_args{gamma_args, inner_args, mul_args};"
- ]
+ "typename EVT::Arguments evt_args{gamma_args, inner_args, mul_args};",
+ ],
},
"k2": {
"name": "K2_RmsNormSwiGLU",
@@ -201,8 +219,8 @@
"",
"// Outer: SwiGLU",
"typename xe_fuse::XePairwiseCompute::Arguments swiglu_args{};",
- "typename EVT::Arguments evt_args{rms_args, swiglu_args};"
- ]
+ "typename EVT::Arguments evt_args{rms_args, swiglu_args};",
+ ],
},
"k2_geglu": {
"name": "K2_RmsNormGeGLU",
@@ -224,8 +242,8 @@
"",
"// Outer: GeGLU",
"typename xe_fuse::XePairwiseCompute::Arguments geglu_args{};",
- "typename EVT::Arguments evt_args{rms_args, geglu_args};"
- ]
+ "typename EVT::Arguments evt_args{rms_args, geglu_args};",
+ ],
},
"k1v2": {
"name": "K1v2_ScaleRowsMerged",
@@ -245,8 +263,8 @@
"// root: XeScaleRowsCompute (merged)",
"typename xe_fuse::XeScaleRowsCompute::Arguments visitor_args{};",
"",
- "typename EVT::Arguments evt_args{scale_args, visitor_args};"
- ]
+ "typename EVT::Arguments evt_args{scale_args, visitor_args};",
+ ],
},
"k2v2": {
"name": "K2v2_SwiGLUScaled",
@@ -266,8 +284,8 @@
"// root: XeScaleRowsSwiGLUCompute (merged)",
"typename xe_fuse::XeScaleRowsSwiGLUCompute::Arguments visitor_args{};",
"",
- "typename EVT::Arguments evt_args{scale_args, visitor_args};"
- ]
+ "typename EVT::Arguments evt_args{scale_args, visitor_args};",
+ ],
},
"k2v2_geglu": {
"name": "K2v2_GeGLUScaled",
@@ -287,8 +305,8 @@
"// root: XeScaleRowsGeGLUCompute (merged)",
"typename xe_fuse::XeScaleRowsGeGLUCompute::Arguments visitor_args{};",
"",
- "typename EVT::Arguments evt_args{scale_args, visitor_args};"
- ]
+ "typename EVT::Arguments evt_args{scale_args, visitor_args};",
+ ],
},
"w8a8_dequant": {
"name": "W8A8_Dequant",
@@ -304,7 +322,12 @@
"evt_typedefs": "using EVT = b::DequantW8A8;",
"aux_data": [
{"name": "scale_token", "type": "float", "shape": "M * L", "init_seed": 42},
- {"name": "scale_channel", "type": "float", "shape": "N * L", "init_seed": 99}
+ {
+ "name": "scale_channel",
+ "type": "float",
+ "shape": "N * L",
+ "init_seed": 99,
+ },
],
"evt_args": [
"// Inner Mul: Acc * scale_token[m] (ColBroadcast, Idx=0)",
@@ -325,8 +348,8 @@
"// Compose: EVT = Mul(Mul(Acc, ColBcast), RowBcast)",
"using InnerMul = b::Mul>;",
"typename InnerMul::Arguments inner_args{accum_args, token_args, inner_mul_args};",
- "typename EVT::Arguments evt_args{inner_args, channel_args, outer_mul_args};"
- ]
+ "typename EVT::Arguments evt_args{inner_args, channel_args, outer_mul_args};",
+ ],
},
"w8a8_dequant_biased": {
"name": "W8A8_DequantBiased",
@@ -342,8 +365,13 @@
"evt_typedefs": "using EVT = b::DequantW8A8Biased;",
"aux_data": [
{"name": "scale_token", "type": "float", "shape": "M * L", "init_seed": 42},
- {"name": "scale_channel", "type": "float", "shape": "N * L", "init_seed": 99},
- {"name": "bias", "type": "float", "shape": "N * L", "init_seed": 77}
+ {
+ "name": "scale_channel",
+ "type": "float",
+ "shape": "N * L",
+ "init_seed": 99,
+ },
+ {"name": "bias", "type": "float", "shape": "N * L", "init_seed": 77},
],
"evt_args": [
"// Inner Mul: Acc * scale_token[m] (ColBroadcast, Idx=0)",
@@ -373,9 +401,9 @@
"using DequantMul = b::DequantW8A8;",
"typename InnerMul::Arguments inner_args{accum_args, token_args, inner_mul_args};",
"typename DequantMul::Arguments dequant_args{inner_args, channel_args, mid_mul_args};",
- "typename EVT::Arguments evt_args{dequant_args, bias_args, add_args};"
- ]
- }
+ "typename EVT::Arguments evt_args{dequant_args, bias_args, add_args};",
+ ],
+ },
}
@@ -387,7 +415,9 @@ def generate_aux_allocations(aux_data: list[dict]) -> str:
name = aux["name"]
shape = aux["shape"]
seed = aux.get("init_seed", 2020)
- lines.append(f" cutlass::DeviceAllocation<{ctype}> block_{name}(static_cast({shape}));")
+ lines.append(
+ f" cutlass::DeviceAllocation<{ctype}> block_{name}(static_cast({shape}));"
+ )
lines.append(f" initialize_block(block_{name}, {seed});")
return "\n".join(lines)
@@ -399,6 +429,7 @@ def generate_cpp(spec: dict, defaults: dict | None = None) -> str:
if template_path.exists():
try:
from jinja2 import Template
+
with open(template_path) as f:
tmpl = Template(f.read())
@@ -411,8 +442,10 @@ def generate_cpp(spec: dict, defaults: dict | None = None) -> str:
align_cd = spec.get("alignment_cd", 8)
make_gemm_extra = ""
if align_ab != 8 or align_cd != 8:
- make_gemm_extra = (f",\n cutlass::layout::RowMajor, cutlass::layout::RowMajor, "
- f"{align_ab}, {align_cd}")
+ make_gemm_extra = (
+ f",\n cutlass::layout::RowMajor, cutlass::layout::RowMajor, "
+ f"{align_ab}, {align_cd}"
+ )
return tmpl.render(
kernel_name=spec["name"],
@@ -465,8 +498,10 @@ def generate_cpp_inline(spec: dict, defaults: dict | None = None) -> str:
# MakeGemm template args beyond the 7 positional defaults
make_gemm_extra = ""
if align_ab != 8 or align_cd != 8:
- make_gemm_extra = (f",\n cutlass::layout::RowMajor, cutlass::layout::RowMajor, "
- f"{align_ab}, {align_cd}")
+ make_gemm_extra = (
+ f",\n cutlass::layout::RowMajor, cutlass::layout::RowMajor, "
+ f"{align_ab}, {align_cd}"
+ )
return f"""\
// Auto-generated xe-fuse kernel benchmark
@@ -675,29 +710,32 @@ def generate_standalone_cpp(spec: dict, defaults: dict | None = None) -> str:
],
"op_code": [
"xe_fuse::standalone::scale_rows(q, block_D.get(), block_scale.get(), M, N, L);"
- ]
+ ],
},
"sa_residual_gamma": {
"name": "SA_ResidualGamma",
"evt_description": "D[m,n] = gamma[n] * (D + residual) (standalone)",
"standalone": True,
"aux_data": [
- {"name": "residual", "type": "bf16", "shape": "M * N * L", "init_seed": 2021},
- {"name": "gamma", "type": "float", "shape": "N * L", "init_seed": 99}
+ {
+ "name": "residual",
+ "type": "bf16",
+ "shape": "M * N * L",
+ "init_seed": 2021,
+ },
+ {"name": "gamma", "type": "float", "shape": "N * L", "init_seed": 99},
],
"op_code": [
"xe_fuse::standalone::residual_gamma(q, block_D.get(), block_residual.get(),",
- " block_gamma.get(), M, N, L);"
- ]
+ " block_gamma.get(), M, N, L);",
+ ],
},
"sa_swiglu": {
"name": "SA_SwiGLU",
"evt_description": "D = SwiGLU(D) pairwise (standalone)",
"standalone": True,
"aux_data": [],
- "op_code": [
- "xe_fuse::standalone::swiglu(q, block_D.get(), M, N, L);"
- ]
+ "op_code": ["xe_fuse::standalone::swiglu(q, block_D.get(), M, N, L);"],
},
"sa_rope_scaled": {
"name": "SA_RoPEScaled",
@@ -705,16 +743,21 @@ def generate_standalone_cpp(spec: dict, defaults: dict | None = None) -> str:
"standalone": True,
"aux_data": [
{"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42},
- {"name": "cos_sin", "type": "float", "shape": "M * N * L", "init_seed": 2024},
- {"name": "tmp", "type": "bf16", "shape": "M * N * L", "init_seed": 0}
+ {
+ "name": "cos_sin",
+ "type": "float",
+ "shape": "M * N * L",
+ "init_seed": 2024,
+ },
+ {"name": "tmp", "type": "bf16", "shape": "M * N * L", "init_seed": 0},
],
"op_code": [
"q.memcpy(block_tmp.get(), block_D.get(), static_cast(M) * N * L * sizeof(bf16));",
"q.wait();",
"xe_fuse::standalone::rope_scaled(q, block_D.get(), block_tmp.get(),",
- " block_scale.get(), block_cos_sin.get(), M, N, L);"
- ]
- }
+ " block_scale.get(), block_cos_sin.get(), M, N, L);",
+ ],
+ },
}
@@ -723,17 +766,25 @@ def main():
all_presets = {**PRESETS, **STANDALONE_PRESETS}
parser.add_argument("--spec", help="Path to kernel spec JSON file")
- parser.add_argument("--preset", choices=list(all_presets.keys()),
- help="Use a built-in kernel preset")
+ parser.add_argument(
+ "--preset",
+ choices=list(all_presets.keys()),
+ help="Use a built-in kernel preset",
+ )
parser.add_argument("--output", "-o", help="Output .cpp path")
parser.add_argument("--m", type=int, default=4096)
parser.add_argument("--n", type=int, default=4096)
parser.add_argument("--k", type=int, default=4096)
parser.add_argument("--iterations", type=int, default=200)
- parser.add_argument("--tile", type=str, default=None,
- help="Override tile shape, e.g. '128x256x32' or 'auto'")
- parser.add_argument("--list-presets", action="store_true",
- help="List available presets and exit")
+ parser.add_argument(
+ "--tile",
+ type=str,
+ default=None,
+ help="Override tile shape, e.g. '128x256x32' or 'auto'",
+ )
+ parser.add_argument(
+ "--list-presets", action="store_true", help="List available presets and exit"
+ )
args = parser.parse_args()
if args.list_presets:
@@ -762,6 +813,7 @@ def main():
if args.tile:
if args.tile == "auto":
from tile_selector import select_tile
+
kernel_tag = args.preset.split("_")[0] if args.preset else "bare"
spec["tile_shape"] = select_tile(args.m, args.n, args.k, kernel_tag)
else:
diff --git a/autotune/generate_pipeline.py b/autotune/generate_pipeline.py
index 27cdf5e..b770223 100644
--- a/autotune/generate_pipeline.py
+++ b/autotune/generate_pipeline.py
@@ -23,8 +23,12 @@
from model_presets import MODEL_PRESETS, list_presets
-def generate_pipeline_cpp(config: dict, preset_name: str = "custom",
- seq_len: int = 2048, autotune: bool = False) -> str:
+def generate_pipeline_cpp(
+ config: dict,
+ preset_name: str = "custom",
+ seq_len: int = 2048,
+ autotune: bool = False,
+) -> str:
template_path = Path(__file__).parent / "pipeline_template.cpp.j2"
if not template_path.exists():
@@ -34,7 +38,10 @@ def generate_pipeline_cpp(config: dict, preset_name: str = "custom",
try:
from jinja2 import Template
except ImportError:
- print("ERROR: jinja2 required. Install with: uv pip install jinja2", file=sys.stderr)
+ print(
+ "ERROR: jinja2 required. Install with: uv pip install jinja2",
+ file=sys.stderr,
+ )
sys.exit(1)
with open(template_path) as f:
@@ -55,6 +62,7 @@ def generate_pipeline_cpp(config: dict, preset_name: str = "custom",
if autotune:
from tile_selector import select_tile
+
k = "k4" if config["use_rope"] else "k1"
tile_vars["tile_shape"] = select_tile(seq_len, H, H, k)
# else: template defaults to _256, _256, _32
@@ -76,16 +84,29 @@ def generate_pipeline_cpp(config: dict, preset_name: str = "custom",
def main():
parser = argparse.ArgumentParser(description="xe-fuse pipeline code generator")
- parser.add_argument("--preset", choices=list(MODEL_PRESETS.keys()),
- help="Use a built-in model preset")
+ parser.add_argument(
+ "--preset",
+ choices=list(MODEL_PRESETS.keys()),
+ help="Use a built-in model preset",
+ )
parser.add_argument("--config", help="Path to custom model config JSON")
parser.add_argument("--output", "-o", help="Output .cpp path")
- parser.add_argument("--autotune", action="store_true",
- help="Auto-select tile shapes per GEMM stage based on sweep data")
- parser.add_argument("--seq-len", type=int, default=2048,
- help="Target sequence length for tile selection (default: 2048)")
- parser.add_argument("--list-presets", action="store_true",
- help="List available model presets and exit")
+ parser.add_argument(
+ "--autotune",
+ action="store_true",
+ help="Auto-select tile shapes per GEMM stage based on sweep data",
+ )
+ parser.add_argument(
+ "--seq-len",
+ type=int,
+ default=2048,
+ help="Target sequence length for tile selection (default: 2048)",
+ )
+ parser.add_argument(
+ "--list-presets",
+ action="store_true",
+ help="List available model presets and exit",
+ )
args = parser.parse_args()
if args.list_presets:
@@ -106,8 +127,9 @@ def main():
parser.error("Provide --preset or --config")
return
- code = generate_pipeline_cpp(config, preset_name,
- seq_len=args.seq_len, autotune=args.autotune)
+ code = generate_pipeline_cpp(
+ config, preset_name, seq_len=args.seq_len, autotune=args.autotune
+ )
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
with open(args.output, "w") as f:
@@ -116,11 +138,14 @@ def main():
n_ffn = 2 * config["I"] if config["gated_ffn"] else config["I"]
print(f"Generated: {args.output}")
print(f" Model: {config['name']}")
- print(f" Dims: H={config['H']}, H_kv={config['H_kv']}, I={config['I']}, N_ffn={n_ffn}")
+ print(
+ f" Dims: H={config['H']}, H_kv={config['H_kv']}, I={config['I']}, N_ffn={n_ffn}"
+ )
print(f" Q/K: {'K4 (RMSNorm+RoPE)' if config['use_rope'] else 'K1 (RMSNorm)'}")
print(f" FFN: {config['ffn_activation']}")
if args.autotune:
from tile_selector import select_tile
+
k = "k4" if config["use_rope"] else "k1"
tile = select_tile(args.seq_len, config["H"], config["H"], k)
print(f" Autotune: M={args.seq_len} -> tile={tile}")
diff --git a/autotune/model_presets.py b/autotune/model_presets.py
index 858b683..0f789e3 100644
--- a/autotune/model_presets.py
+++ b/autotune/model_presets.py
@@ -15,8 +15,8 @@
"llama3_8b": {
"name": "LLaMA 3 8B",
"H": 4096,
- "H_kv": 1024, # 8 KV heads * 128 head_dim (GQA)
- "I": 14336, # intermediate_size
+ "H_kv": 1024, # 8 KV heads * 128 head_dim (GQA)
+ "I": 14336, # intermediate_size
"num_layers": 32,
"use_rope": True,
"ffn_activation": "swiglu",
@@ -25,7 +25,7 @@
"llama2_7b": {
"name": "LLaMA 2 7B",
"H": 4096,
- "H_kv": 4096, # MHA (no GQA)
+ "H_kv": 4096, # MHA (no GQA)
"I": 11008,
"num_layers": 32,
"use_rope": True,
@@ -35,7 +35,7 @@
"llama3_70b": {
"name": "LLaMA 3 70B",
"H": 8192,
- "H_kv": 1024, # 8 KV heads * 128 head_dim
+ "H_kv": 1024, # 8 KV heads * 128 head_dim
"I": 28672,
"num_layers": 80,
"use_rope": True,
@@ -45,7 +45,7 @@
"gemma2_9b": {
"name": "Gemma 2 9B",
"H": 3584,
- "H_kv": 2048, # 8 KV heads * 256 head_dim
+ "H_kv": 2048, # 8 KV heads * 256 head_dim
"I": 14336,
"num_layers": 42,
"use_rope": True,
@@ -55,7 +55,7 @@
"gemma2_27b": {
"name": "Gemma 2 27B",
"H": 4608,
- "H_kv": 2048, # 16 KV heads * 128 head_dim
+ "H_kv": 2048, # 16 KV heads * 128 head_dim
"I": 36864,
"num_layers": 46,
"use_rope": True,
@@ -65,7 +65,7 @@
"mistral_7b": {
"name": "Mistral 7B",
"H": 4096,
- "H_kv": 1024, # 8 KV heads * 128 head_dim
+ "H_kv": 1024, # 8 KV heads * 128 head_dim
"I": 14336,
"num_layers": 32,
"use_rope": True,
@@ -75,7 +75,7 @@
"qwen25_7b": {
"name": "Qwen 2.5 7B",
"H": 3584,
- "H_kv": 512, # 4 KV heads * 128 head_dim
+ "H_kv": 512, # 4 KV heads * 128 head_dim
"I": 18944,
"num_layers": 28,
"use_rope": True,
@@ -85,7 +85,7 @@
"qwen25_72b": {
"name": "Qwen 2.5 72B",
"H": 8192,
- "H_kv": 1024, # 8 KV heads * 128 head_dim
+ "H_kv": 1024, # 8 KV heads * 128 head_dim
"I": 29568,
"num_layers": 80,
"use_rope": True,
@@ -95,7 +95,7 @@
"phi3_mini": {
"name": "Phi-3 Mini 3.8B",
"H": 3072,
- "H_kv": 3072, # MHA (no GQA)
+ "H_kv": 3072, # MHA (no GQA)
"I": 8192,
"num_layers": 32,
"use_rope": True,
@@ -105,7 +105,7 @@
"phi3_medium": {
"name": "Phi-3 Medium 14B",
"H": 5120,
- "H_kv": 5120, # MHA (no GQA)
+ "H_kv": 5120, # MHA (no GQA)
"I": 17920,
"num_layers": 40,
"use_rope": True,
@@ -124,8 +124,12 @@ def get_preset(name: str) -> dict:
def list_presets() -> None:
print("Available model presets:")
- print(f" {'Name':<16s} {'Model':<20s} {'H':>6s} {'H_kv':>6s} {'I':>6s} {'FFN':>8s} {'RoPE'}")
+ print(
+ f" {'Name':<16s} {'Model':<20s} {'H':>6s} {'H_kv':>6s} {'I':>6s} {'FFN':>8s} {'RoPE'}"
+ )
print(" " + "-" * 78)
for key, cfg in MODEL_PRESETS.items():
- print(f" {key:<16s} {cfg['name']:<20s} {cfg['H']:>6d} {cfg['H_kv']:>6d} "
- f"{cfg['I']:>6d} {cfg['ffn_activation']:>8s} {'yes' if cfg['use_rope'] else 'no':>4s}")
+ print(
+ f" {key:<16s} {cfg['name']:<20s} {cfg['H']:>6d} {cfg['H_kv']:>6d} "
+ f"{cfg['I']:>6d} {cfg['ffn_activation']:>8s} {'yes' if cfg['use_rope'] else 'no':>4s}"
+ )
diff --git a/autotune/tile_selector.py b/autotune/tile_selector.py
index dcd6a39..533ca23 100644
--- a/autotune/tile_selector.py
+++ b/autotune/tile_selector.py
@@ -1,87 +1,47 @@
"""
xe-fuse autotune: tile shape selector.
-Selects the best tile shape for a GEMM based on (M, N, K, groups) dimensions.
-Rules derived from empirical sweeps on Intel Arc Pro B70 (BMG G31) and
-CRI SKU4 (Xe3P, 32 XeCores, 2500 MHz):
- - Round 1: 7 tiles × 32 shapes × 6 M values (444 points)
- - Round 3: 25 tiles × 44 shapes including skinny-M DPAS (420 points)
- - Comprehensive sweep: 25 tiles × 79 shapes incl. N=384/3584 (822 points)
- - CRI grouped GEMM: ww27_bd_benchmarks (fp8/mxfp8, G=4/8/16, MoE shapes)
- MFU formula: TFLOPS*1000/(32*8192*2.5*0.4375[*2 for fp8])*100
- CRI peak: ~287 TFLOPS bf16, ~573 TFLOPS fp8
- - G>16 extrapolation: G=24/48/96/192 shapes from bytedance_ops_WW35 (shapes only,
- no perf data). Extrapolated from G>=8 tall-M rule — G>=16 branch covers all G>16.
- One mxfp4 data point (G=32, M=171, N=3584, K=1280) shows 192x512x128 winning
- over baseline at 23.6%; fp8 equivalent unknown. G>=96 small-M shapes flagged.
+Selects the best tile shape for a GEMM based on (M, N, K) dimensions.
+Rules derived from sweeps on Intel Arc Pro B70 (BMG G31).
Usage:
from tile_selector import select_tile
tile = select_tile(M=128, N=4096, K=4096)
# Returns "_128, _128, _32"
-
- tile = select_tile(M=683, N=3584, K=1280, groups=8)
- # Returns "_352, _256, _64" (71% MFU on CRI fp8)
"""
-import csv
-import os
-from pathlib import Path
-
TILES = {
- # --- BMG B70 tiles (empirical sweep on Arc Pro B70 / BMG G31) ---
"512x128x32": "_512, _128, _32",
- "512x64x32": "_512, _64, _32",
+ "512x64x32": "_512, _64, _32",
"384x128x32": "_384, _128, _32",
- "384x64x32": "_384, _64, _32",
+ "384x64x32": "_384, _64, _32",
"256x512x32": "_256, _512, _32",
"256x256x32": "_256, _256, _32",
"256x128x32": "_256, _128, _32",
- "256x64x32": "_256, _64, _32",
- "256x64x64": "_256, _64, _64",
+ "256x64x32": "_256, _64, _32",
+ "256x64x64": "_256, _64, _64",
"192x128x32": "_192, _128, _32",
"128x256x32": "_128, _256, _32",
"128x128x32": "_128, _128, _32",
"128x128x64": "_128, _128, _64",
- "128x64x32": "_128, _64, _32",
- "128x64x64": "_128, _64, _64",
- "64x256x32": "_64, _256, _32",
- "64x128x32": "_64, _128, _32",
- "64x64x32": "_64, _64, _32",
- "32x256x32": "_32, _256, _32",
- "32x256x64": "_32, _256, _64",
- "32x128x32": "_32, _128, _32",
- "32x128x64": "_32, _128, _64",
- "16x256x32": "_16, _256, _32",
- "16x128x32": "_16, _128, _32",
- "8x256x32": "_8, _256, _32",
- "8x128x32": "_8, _128, _32",
- "4x256x32": "_4, _256, _32",
- "2x256x32": "_2, _256, _32",
- "1x256x32": "_1, _256, _32",
- "1x128x32": "_1, _128, _32",
- # --- CRI SKU4 tiles (best_vs_baseline_ordered.xlsx + ww27_bd_benchmarks) ---
- # Winners: +35–42 pp MFU over 256x256 baseline for MoE shapes
- "352x256x64": "_352, _256, _64",
- "448x256x64": "_448, _256, _64",
- "192x640x64": "_192, _640, _64",
- "448x320x64": "_448, _320, _64",
- "320x512x128": "_320, _512, _128",
- "192x640x128": "_192, _640, _128",
- "96x896x128": "_96, _896, _128",
- "352x256x128": "_352, _256, _128",
- "448x256x128": "_448, _256, _128",
- "448x320x128": "_448, _320, _128",
- "128x896x128": "_128, _896, _128",
- "320x512x64": "_320, _512, _64",
- # ww27_bd_benchmarks new tiles
- "128x896x64": "_128, _896, _64",
- "64x896x64": "_64, _896, _64",
- "96x896x64": "_96, _896, _64",
- "608x128x64": "_608, _128, _64",
- # G>16 extrapolation (tile_recalc model + G=32 mxfp4 data point)
- "192x512x64": "_192, _512, _64",
+ "128x64x32": "_128, _64, _32",
+ "128x64x64": "_128, _64, _64",
+ "64x256x32": "_64, _256, _32",
+ "64x128x32": "_64, _128, _32",
+ "64x64x32": "_64, _64, _32",
+ "32x256x32": "_32, _256, _32",
+ "32x256x64": "_32, _256, _64",
+ "32x128x32": "_32, _128, _32",
+ "32x128x64": "_32, _128, _64",
+ "16x256x32": "_16, _256, _32",
+ "16x128x32": "_16, _128, _32",
+ "8x256x32": "_8, _256, _32",
+ "8x128x32": "_8, _128, _32",
+ "4x256x32": "_4, _256, _32",
+ "2x256x32": "_2, _256, _32",
+ "1x256x32": "_1, _256, _32",
+ "1x128x32": "_1, _128, _32",
}
@@ -89,128 +49,62 @@ def select_tile(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1) -
"""Select optimal tile shape string for CUTLASS cute::Shape<>.
Args:
- M: number of rows (total across all groups for grouped GEMM)
+ M: number of rows
N: output columns (hidden dim, FFN dim, etc.)
K: reduction dim (input hidden dim)
kernel: "bare", "k1", "k2", or "k4" (affects register pressure)
- groups: number of GEMM groups (>1 enables CRI grouped GEMM rules)
+ groups: unused, kept for API compatibility
Returns:
CUTLASS tile shape string, e.g. "_128, _128, _32"
"""
tile_k = 32 # default; overridden below for shapes where K=64 validated better
- # ── CRI Grouped GEMM (G > 1) ────────────────────────────────────────────
- # Rules from ww27_bd_benchmarks (fp8/mxfp8, CRI SKU4, 32 XeCores).
- # Key insight: G>=8 → tall-M tiles for wave fill; G=4 → wider-N tiles.
- # Wide-N tiles (64_1792, 64_1280) are WRONG for G>=8 regardless of dtype.
- # G>16 (G=24/48/96/192 from ByteDance WW35): no fp8 perf data — extrapolated
- # from G=16 rule (groups>=16 branch covers all G>16). Exception: G>=32 with
- # M<=200 (e.g. G=96/192 decode) has one mxfp4 data point (G=32, M=171) showing
- # 192x512x128 winning at 23.6% MFU; fp8 tile unknown — needs fulsim validation.
- if groups > 1:
- if N >= 3000 and K <= 1500:
- # N=3584, K=1280 family (short K — GroupScheduler overhead dominant)
- if groups >= 16:
- # tall-M: 69% MFU at M=427 G=16 with 448x256x64
- # 64_1792_64 gives 17-20% MFU (wrong tile class)
- # G=24/48/96/192: extrapolated from G=16, no fp8 data.
- # Small M (≤200): tile_recalc + G=32 mxfp4 data (192x512x128 at 23.6%)
- # prefer 192x512 — tile_m matches group boundary; 448 spans ~2.6 groups
- if M <= 200:
- tile_m, tile_n, tile_k = 192, 512, 64 # model-extrapolated, needs fulsim
- else:
- tile_m, tile_n, tile_k = 448, 256, 64
- elif groups >= 8:
- if M <= 800:
- tile_m, tile_n, tile_k = 352, 256, 64 # 71% at M=683 G=8
- elif M <= 1500:
- tile_m, tile_n, tile_k = 128, 896, 64 # 64% at M=1366 G=8
- elif M <= 3000:
- tile_m, tile_n, tile_k = 64, 896, 64 # 37% at M=2731 G=8
- else:
- tile_m, tile_n, tile_k = 96, 896, 64 # 54% at M=5462 G=8
- else:
- # G=4: per-group M is large enough for wider N tiles
- tile_m, tile_n, tile_k = 320, 512, 64 # 38% at M=2560 G=4
- elif N >= 2000 and K >= 3000:
- # N=2560, K=3584 family (longer K — better arithmetic intensity)
- if groups >= 16:
- # G=16 wrong tile: 64_1280_64 gives 35%; use tall-M
- # G=24/48/96/192: extrapolated from G=16, no fp8 data.
- # Small M (≤200): tile_recalc + ww27 G=8 winner (192x640x64 at 53%)
- if M <= 200:
- tile_m, tile_n, tile_k = 192, 640, 64 # model-extrapolated, needs fulsim
- else:
- tile_m, tile_n, tile_k = 448, 256, 64
- elif groups >= 8:
- if M <= 200:
- tile_m, tile_n, tile_k = 192, 640, 64 # 53% at M=171 G=8
- elif M <= 800:
- tile_m, tile_n, tile_k = 352, 256, 64 # 83% at M=683 G=8
- else:
- tile_m, tile_n, tile_k = 608, 128, 64 # 88% at M=5462 G=8
- else:
- # G=4
- if M <= 500:
- tile_m, tile_n, tile_k = 448, 320, 64 # 77% at M=427 G=4
- elif M <= 1500:
- tile_m, tile_n, tile_k = 352, 256, 64 # 84% at M=1366 G=4
- else:
- tile_m, tile_n, tile_k = 320, 512, 64 # 61% at M=2560 G=4
- else:
- # Unknown grouped shape — fall through to bare GEMM heuristic
- tile_m, tile_n = 256, 256
-
- key = f"{tile_m}x{tile_n}x{tile_k}"
- return TILES.get(key, "_256, _256, _32")
-
# ── Skinny-M DPAS tiles for GEMV (M ≤ 64) ──
- # Rescue3 sweep: skinny tiles give 2-7x over 64x256 at small M.
- # Best tile depends on (M, N, K) combination.
+ # Skinny tiles give 2-7x over 64x256 at small M.
if M <= 1:
if N >= 8192:
- tile_m, tile_n = 4, 256 # M=1 N=8192 K=1024: 3.8x speedup
+ tile_m, tile_n = 4, 256
elif N >= 4096:
- tile_m, tile_n = 16, 128 # M=1 N=4096: 1.8x
+ tile_m, tile_n = 16, 128
else:
- tile_m, tile_n = 8, 128 # M=1 N=1024: 6x
+ tile_m, tile_n = 8, 128
elif M <= 2:
if K <= 1024 and N >= 8192:
tile_m, tile_n = 8, 256
else:
- tile_m, tile_n = 4, 256 # M=2 N=1024 K=8192: 5.7x
+ tile_m, tile_n = 4, 256
elif M <= 4:
if N >= 8192 and K <= 1024:
- tile_m, tile_n = 8, 128 # M=4 N=8192 K=1024: 3.6x
+ tile_m, tile_n = 8, 128
else:
tile_m, tile_n = 32, 256
elif M <= 8:
if N <= 1024:
- tile_m, tile_n = 2, 256 # M=8 N=1024 K=8192: 6.6x
+ tile_m, tile_n = 2, 256
elif N >= 8192 and K <= 1024:
- tile_m, tile_n = 8, 256 # M=8 N=8192 K=1024: 2.1x
+ tile_m, tile_n = 8, 256
else:
tile_m, tile_n = 16, 128
elif M <= 16:
if N <= 1024:
- tile_m, tile_n = 8, 128 # M=16 N=1024: 4x
+ tile_m, tile_n = 8, 128
elif N >= 8192 and K <= 1024:
- tile_m, tile_n = 16, 256 # M=16 N=8192 K=1024: 2.6x
+ tile_m, tile_n = 16, 256
else:
tile_m, tile_n = 16, 128
elif M <= 32:
if N <= 1024:
- tile_m, tile_n = 8, 128 # M=32 N=1024: 3.3x
+ tile_m, tile_n = 8, 128
elif N >= 8192 and K <= 1024:
- tile_m, tile_n = 32, 256 # M=32 N=8192 K=1024: 2.2x
+ tile_m, tile_n = 32, 256
else:
tile_m, tile_n = 16, 256
elif M <= 64:
if N <= 1024 and K >= 4096:
- tile_m, tile_n, tile_k = 32, 128, 64 # val 307282: 7.3% MFU (K=64 2x over K=32)
+ tile_m, tile_n, tile_k = 32, 128, 64
elif N <= 1024:
- tile_m, tile_n = 8, 128 # M=64 N=1024: 2.9x
+ tile_m, tile_n = 8, 128
elif N >= 4096:
tile_m, tile_n = 64, 256
else:
@@ -218,74 +112,70 @@ def select_tile(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1) -
# ── Standard tiles for M ≥ 128 ──
elif M <= 128:
if N <= 384 and K > 1024:
- tile_m, tile_n = 128, 128 # N=384 K=3584: 128x128 best
- elif N <= 384:
- tile_m, tile_n = 128, 64
+ tile_m, tile_n = 128, 128
elif N <= 1024:
- tile_m, tile_n = 128, 64 # sweep 307257: 9.3% MFU (was 5.9% with 128x128)
+ tile_m, tile_n = 128, 64
elif N <= 4096:
- tile_m, tile_n = 64, 64 # sweep 307257: 28.7% MFU (was 24.3% with 64x256)
+ tile_m, tile_n = 64, 64
elif K <= 1024:
- tile_m, tile_n = 32, 128 # deep sweep 307267: 24.9% MFU (was 22.5% with 64x128)
+ tile_m, tile_n = 32, 128
else:
tile_m, tile_n = 128, 256
elif M <= 256:
if N <= 384:
tile_m, tile_n = 256, 64
elif N <= 1024:
- tile_m, tile_n, tile_k = 256, 64, 64 # val 307282: 25.1% MFU (K=64 vs 18.1% K=32)
+ tile_m, tile_n, tile_k = 256, 64, 64
elif N <= 4096:
- tile_m, tile_n = 256, 128 # M=256 N=4096: 1.8x over 256x256
+ tile_m, tile_n = 256, 128
else:
tile_m, tile_n = 256, 128
elif M <= 384:
if N <= 384 and K > 1024:
- tile_m, tile_n = 128, 128 # N=384 K=3584: 128x128 3x over 256x64
+ tile_m, tile_n = 128, 128
elif N <= 384:
tile_m, tile_n = 256, 64
elif N <= 1024:
- tile_m, tile_n = 384, 64 # val 307282: 24.0% MFU (was 19.8% with 64x128)
+ tile_m, tile_n = 384, 64
elif N <= 4096 and K <= 384:
- tile_m, tile_n = 192, 128 # N=3584 K=384: 192x128 2x over 384x128
+ tile_m, tile_n = 192, 128
else:
- tile_m, tile_n = 384, 128 # M=384: 1.5-1.6x over 256x256
+ tile_m, tile_n = 384, 128
elif M <= 512:
if N <= 384 and K > 1024:
- tile_m, tile_n = 128, 128 # N=384 K=3584: 128x128 1.4x over 256x64
+ tile_m, tile_n = 128, 128
elif N <= 384:
tile_m, tile_n = 256, 64
elif N <= 1024:
- tile_m, tile_n, tile_k = 32, 128, 64 # val 307282: 50.1% MFU (K=64 vs 31.6% K=32)
+ tile_m, tile_n, tile_k = 32, 128, 64
elif N <= 4096 and K <= 384:
- tile_m, tile_n = 128, 128 # N=3584 K=384: 128x128 1.5x
+ tile_m, tile_n = 128, 128
elif N <= 4096:
- tile_m, tile_n = 256, 128 # M=512 N=4096: 256x128 ≈ best
+ tile_m, tile_n = 256, 128
else:
- tile_m, tile_n = 256, 256 # M=512 N=8192: 256x256 best
+ tile_m, tile_n = 256, 256
elif M <= 640:
if N <= 384 and K > 1024:
tile_m, tile_n = 128, 128
- elif N <= 384:
- tile_m, tile_n = 384, 64
elif N <= 1024:
- tile_m, tile_n = 384, 64 # val 307282: 27.9% MFU (43.9% was run-to-run variance)
+ tile_m, tile_n = 384, 64
elif N <= 4096 and K <= 384:
tile_m, tile_n = 128, 128
elif N >= 8192 and K <= 1024:
tile_m, tile_n = 128, 256
elif N >= 4096 and K >= 4096:
- tile_m, tile_n = 256, 512 # CRI SKU4 archstudy: +19% over 256x256 for N=K=4096
+ tile_m, tile_n = 256, 512
else:
tile_m, tile_n = 256, 256 if N >= 4096 else 128
elif M <= 768:
if N <= 384 and K > 1024:
- tile_m, tile_n = 128, 128 # N=384 K=3584: 128x128 dominant
+ tile_m, tile_n = 128, 128
elif N <= 384:
tile_m, tile_n = 384, 64
elif N <= 1024:
- tile_m, tile_n = 32, 256 # val 307282: 51.9% MFU! (768/32=24 M-tiles, 96% wave fill)
+ tile_m, tile_n = 32, 256
elif N <= 4096 and K <= 384:
- tile_m, tile_n = 128, 128 # N=3584 K=384: 128x128
+ tile_m, tile_n = 128, 128
elif N >= 8192 and K <= 1024:
tile_m, tile_n = 128, 256
else:
@@ -296,7 +186,7 @@ def select_tile(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1) -
elif N <= 384:
tile_m, tile_n = 384, 64
elif N <= 1024:
- tile_m, tile_n, tile_k = 32, 256, 64 # val 307282: 35.7% MFU (K=64 vs 30.2% K=32)
+ tile_m, tile_n, tile_k = 32, 256, 64
elif N <= 4096 and K <= 384:
tile_m, tile_n = 128, 128
elif N >= 8192 and K <= 1024:
@@ -307,44 +197,44 @@ def select_tile(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1) -
# M >= 1024
if N <= 384 and K > 1024:
if M <= 1024:
- tile_m, tile_n = 256, 64 # M=1024 N=384 K=3584: 256x64
+ tile_m, tile_n = 256, 64
elif M <= 2048:
- tile_m, tile_n = 256, 128 # M=2048: 256x128 best
+ tile_m, tile_n = 256, 128
elif M <= 3328:
- tile_m, tile_n = 384, 128 # M=3328: 384x128 best
+ tile_m, tile_n = 384, 128
elif M <= 4096:
- tile_m, tile_n = 256, 64 # M=4096: 256x64
+ tile_m, tile_n = 256, 64
elif M <= 7680:
- tile_m, tile_n = 128, 64 # val 307282: 47.7% MFU (was 42.3% with 384x64)
+ tile_m, tile_n = 128, 64
else:
- tile_m, tile_n, tile_k = 128, 128, 64 # val 307282: 53.4% MFU (K=64 vs 40.7% K=32)
+ tile_m, tile_n, tile_k = 128, 128, 64
elif N <= 384:
tile_m, tile_n = 512, 64
elif N <= 1024:
if M <= 1024:
- tile_m, tile_n, tile_k = 32, 128, 64 # val 307282: 56.6% MFU (K=64+S3, was 56.5% K=32)
+ tile_m, tile_n, tile_k = 32, 128, 64
else:
- tile_m, tile_n = 256, 256 # large M: 256x256 still good (65.4% at M=8192)
+ tile_m, tile_n = 256, 256
elif N <= 4096 and K <= 384:
if M <= 1024:
- tile_m, tile_n = 256, 64 # N=3584 K=384 M=1024: 256x64
+ tile_m, tile_n = 256, 64
elif M <= 2048:
- tile_m, tile_n = 128, 256 # M=2048: 128x256
+ tile_m, tile_n = 128, 256
elif M <= 3328:
- tile_m, tile_n = 256, 64 # M=3328: 256x64
+ tile_m, tile_n = 256, 64
elif M <= 4096:
- tile_m, tile_n = 128, 256 # M=4096: 128x256
+ tile_m, tile_n = 128, 256
else:
- tile_m, tile_n = 256, 128 # M=7680: 256x128
+ tile_m, tile_n = 256, 128
elif N >= 11264:
- tile_m, tile_n = 512, 128 # N=11264: 1.5x over 256x128
+ tile_m, tile_n = 512, 128
elif N >= 8192 and K <= 1024:
if M <= 1024:
- tile_m, tile_n = 256, 128 # M=1024 N=8192 K=1024: 256x128
+ tile_m, tile_n = 256, 128
else:
- tile_m, tile_n = 192, 128 # large M: 192x128
+ tile_m, tile_n = 192, 128
elif N >= 4096 and K >= 4096:
- tile_m, tile_n = 256, 512 # CRI SKU4 archstudy: compute_bound at 92.6% EU active for M=1024
+ tile_m, tile_n = 256, 512
else:
tile_m, tile_n = 256, 256 if M >= 512 else 128
@@ -365,40 +255,9 @@ def select_tile(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1) -
return TILES.get(key, "_256, _256, _32")
-def select_tile_from_csv(M: int, N: int, K: int, kernel: str = "Bare_GEMM(bf16)", groups: int = 1) -> str:
- """Look up the best tile from the sweep CSV if available, else fall back to heuristic."""
- csv_path = Path(__file__).parent.parent / "tests" / "best_tiles.csv"
- if not csv_path.exists():
- return select_tile(M, N, K, kernel, groups)
-
- best_tf = 0.0
- best_tile = None
-
- with open(csv_path) as f:
- reader = csv.DictReader(f)
- for row in reader:
- if (row["kernel"] == kernel and
- int(row["M"]) == M and
- int(row["N"]) == N and
- int(row["K"]) == K):
- tf = float(row["tflops"])
- if tf > best_tf:
- best_tf = tf
- best_tile = row["best_tile"]
-
- if best_tile and best_tile in TILES:
- return TILES[best_tile]
-
- return select_tile(M, N, K, kernel, groups)
-
-
-def tile_shape_str(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1) -> str:
- """Convenience: returns the tile shape for generate_kernel/pipeline use."""
- return select_tile(M, N, K, kernel, groups)
-
-
if __name__ == "__main__":
import argparse
+
parser = argparse.ArgumentParser(description="xe-fuse tile selector")
parser.add_argument("--m", type=int, required=True)
parser.add_argument("--n", type=int, required=True)
@@ -408,4 +267,6 @@ def tile_shape_str(M: int, N: int, K: int, kernel: str = "bare", groups: int = 1
args = parser.parse_args()
tile = select_tile(args.m, args.n, args.k, args.kernel, args.groups)
- print(f"M={args.m} N={args.n} K={args.k} groups={args.groups} kernel={args.kernel} -> tile={tile}")
+ print(
+ f"M={args.m} N={args.n} K={args.k} groups={args.groups} kernel={args.kernel} -> tile={tile}"
+ )
diff --git a/tests/bench_e2e_comparison.py b/tests/bench_e2e_comparison.py
index 3d4290e..e43cc04 100644
--- a/tests/bench_e2e_comparison.py
+++ b/tests/bench_e2e_comparison.py
@@ -17,32 +17,58 @@
sbatch run_bench_e2e.sh
"""
+import argparse
import time
+
import torch
-import argparse
def load_vllm_ops():
import vllm_xpu_kernels._C # noqa: F401
+
return torch.ops._C
MODEL_PRESETS = {
"llama3_8b": {
- "name": "LLaMA 3 8B", "H": 4096, "H_kv": 1024, "I": 14336,
- "head_dim": 128, "num_heads": 32, "num_kv_heads": 8, "act": "swiglu",
+ "name": "LLaMA 3 8B",
+ "H": 4096,
+ "H_kv": 1024,
+ "I": 14336,
+ "head_dim": 128,
+ "num_heads": 32,
+ "num_kv_heads": 8,
+ "act": "swiglu",
},
"gemma2_9b": {
- "name": "Gemma 2 9B", "H": 3584, "H_kv": 2048, "I": 14336,
- "head_dim": 256, "num_heads": 16, "num_kv_heads": 8, "act": "geglu",
+ "name": "Gemma 2 9B",
+ "H": 3584,
+ "H_kv": 2048,
+ "I": 14336,
+ "head_dim": 256,
+ "num_heads": 16,
+ "num_kv_heads": 8,
+ "act": "geglu",
},
"qwen25_7b": {
- "name": "Qwen 2.5 7B", "H": 3584, "H_kv": 512, "I": 18944,
- "head_dim": 128, "num_heads": 28, "num_kv_heads": 4, "act": "swiglu",
+ "name": "Qwen 2.5 7B",
+ "H": 3584,
+ "H_kv": 512,
+ "I": 18944,
+ "head_dim": 128,
+ "num_heads": 28,
+ "num_kv_heads": 4,
+ "act": "swiglu",
},
"phi3_mini": {
- "name": "Phi-3 Mini 3.8B", "H": 3072, "H_kv": 3072, "I": 8192,
- "head_dim": 128, "num_heads": 24, "num_kv_heads": 24, "act": "swiglu",
+ "name": "Phi-3 Mini 3.8B",
+ "H": 3072,
+ "H_kv": 3072,
+ "I": 8192,
+ "head_dim": 128,
+ "num_heads": 24,
+ "num_kv_heads": 24,
+ "act": "swiglu",
},
}
@@ -164,12 +190,14 @@ def run_gemms():
def pipeline_flops(M, H, H_kv, N_ffn):
"""Total FLOPs for 4 GEMMs in the pipeline."""
- return 2.0 * M * (H*H + H_kv*H + H*H + N_ffn*H)
+ return 2.0 * M * (H * H + H_kv * H + H * H + N_ffn * H)
def main():
parser = argparse.ArgumentParser(description="xe-fuse vs real vllm: E2E comparison")
- parser.add_argument("--preset", default="llama3_8b", choices=list(MODEL_PRESETS.keys()))
+ parser.add_argument(
+ "--preset", default="llama3_8b", choices=list(MODEL_PRESETS.keys())
+ )
parser.add_argument("--m", type=int, default=2048)
parser.add_argument("--iters", type=int, default=200)
parser.add_argument("--all", action="store_true", help="Run all presets")
@@ -191,10 +219,12 @@ def main():
N_flops = pipeline_flops(args.m, H, H_kv, 2 * I)
M = args.m
- print(f"\n{'='*60}")
+ print(f"\n{'=' * 60}")
print(f" {cfg['name']} (H={H}, H_kv={H_kv}, I={I})")
- print(f" Pipeline: Q(RMSNorm+RoPE) + V(RMSNorm) + O(ResAdd+RMSNorm) + FFN({cfg['act']})")
- print(f"{'='*60}")
+ print(
+ f" Pipeline: Q(RMSNorm+RoPE) + V(RMSNorm) + O(ResAdd+RMSNorm) + FFN({cfg['act']})"
+ )
+ print(f"{'=' * 60}")
# Benchmark oneDNN GEMMs only (no ops)
t_gemm_only = bench_onednn_gemm_only(cfg, M, iters=args.iters)
@@ -208,16 +238,16 @@ def main():
t_ops_overhead = t_vllm_e2e - t_gemm_only
ops_pct = t_ops_overhead / t_vllm_e2e * 100
- print(f"\n oneDNN GEMMs only (4x torch.mm):")
+ print("\n oneDNN GEMMs only (4x torch.mm):")
print(f" Time: {t_gemm_only:.4f} ms")
print(f" Throughput: {gemm_tflops:.1f} TFlop/s")
- print(f"\n vllm full pipeline (oneDNN GEMM + real vllm ops):")
+ print("\n vllm full pipeline (oneDNN GEMM + real vllm ops):")
print(f" Time: {t_vllm_e2e:.4f} ms")
print(f" Throughput: {vllm_tflops:.1f} TFlop/s")
print(f" Ops overhead: {t_ops_overhead:.4f} ms ({ops_pct:.1f}% of pipeline)")
# Structured output for parsing
- print(f"\n=== STRUCTURED OUTPUT ===")
+ print("\n=== STRUCTURED OUTPUT ===")
print(f"E2E: {preset_name}")
print(f"MODEL: {cfg['name']}")
print(f"DIMS: M={M} H={H} H_kv={H_kv} I={I}")
diff --git a/tests/bench_vllm_real.py b/tests/bench_vllm_real.py
index f1e5039..9ca4368 100644
--- a/tests/bench_vllm_real.py
+++ b/tests/bench_vllm_real.py
@@ -14,13 +14,16 @@
sbatch run_bench_vllm_real.sh
"""
+import argparse
import time
+
import torch
-import argparse
+
def load_vllm_ops():
"""Load vllm-xpu-kernels C extension."""
import vllm_xpu_kernels._C # noqa: F401
+
return torch.ops._C
@@ -80,7 +83,9 @@ def bench_gelu_and_mul(ops, M, N, warmup=20, iters=200):
return (time.perf_counter() - t0) / iters * 1e3
-def bench_rotary_embedding(ops, M, num_heads, num_kv_heads, head_dim, warmup=20, iters=200):
+def bench_rotary_embedding(
+ ops, M, num_heads, num_kv_heads, head_dim, warmup=20, iters=200
+):
"""Benchmark rotary_embedding: in-place NeoX RoPE on Q+K."""
rot_dim = head_dim
@@ -99,8 +104,20 @@ def bench_rotary_embedding(ops, M, num_heads, num_kv_heads, head_dim, warmup=20,
return (time.perf_counter() - t0) / iters * 1e3
-def bench_pipeline_standalone(ops, M, H, H_kv, I, head_dim, num_heads,
- num_kv_heads, eps, act_fn, warmup=20, iters=200):
+def bench_pipeline_standalone(
+ ops,
+ M,
+ H,
+ H_kv,
+ I,
+ head_dim,
+ num_heads,
+ num_kv_heads,
+ eps,
+ act_fn,
+ warmup=20,
+ iters=200,
+):
"""Benchmark a full pipeline of vllm standalone ops (no GEMM).
Measures the total cost of post-GEMM operations that xe-fuse eliminates."""
rot_dim = head_dim
@@ -124,7 +141,9 @@ def bench_pipeline_standalone(ops, M, H, H_kv, I, head_dim, num_heads,
v_normed = torch.empty_like(v_out)
for _ in range(warmup):
- ops.rms_norm(q_normed, q_out, torch.randn(Q_dim, dtype=torch.bfloat16, device="xpu"), eps)
+ ops.rms_norm(
+ q_normed, q_out, torch.randn(Q_dim, dtype=torch.bfloat16, device="xpu"), eps
+ )
ops.rotary_embedding(positions, q_normed, None, head_dim, cos_sin_cache, True)
ops.rms_norm(v_normed, v_out, weight_kv, eps)
ops.fused_add_rms_norm(o_proj, residual, weight_h, eps)
@@ -151,27 +170,53 @@ def bench_pipeline_standalone(ops, M, H, H_kv, I, head_dim, num_heads,
MODEL_PRESETS = {
"llama3_8b": {
- "name": "LLaMA 3 8B", "H": 4096, "H_kv": 1024, "I": 14336,
- "head_dim": 128, "num_heads": 32, "num_kv_heads": 8, "act": "swiglu",
+ "name": "LLaMA 3 8B",
+ "H": 4096,
+ "H_kv": 1024,
+ "I": 14336,
+ "head_dim": 128,
+ "num_heads": 32,
+ "num_kv_heads": 8,
+ "act": "swiglu",
},
"gemma2_9b": {
- "name": "Gemma 2 9B", "H": 3584, "H_kv": 2048, "I": 14336,
- "head_dim": 256, "num_heads": 16, "num_kv_heads": 8, "act": "geglu",
+ "name": "Gemma 2 9B",
+ "H": 3584,
+ "H_kv": 2048,
+ "I": 14336,
+ "head_dim": 256,
+ "num_heads": 16,
+ "num_kv_heads": 8,
+ "act": "geglu",
},
"qwen25_7b": {
- "name": "Qwen 2.5 7B", "H": 3584, "H_kv": 512, "I": 18944,
- "head_dim": 128, "num_heads": 28, "num_kv_heads": 4, "act": "swiglu",
+ "name": "Qwen 2.5 7B",
+ "H": 3584,
+ "H_kv": 512,
+ "I": 18944,
+ "head_dim": 128,
+ "num_heads": 28,
+ "num_kv_heads": 4,
+ "act": "swiglu",
},
"phi3_mini": {
- "name": "Phi-3 Mini 3.8B", "H": 3072, "H_kv": 3072, "I": 8192,
- "head_dim": 128, "num_heads": 24, "num_kv_heads": 24, "act": "swiglu",
+ "name": "Phi-3 Mini 3.8B",
+ "H": 3072,
+ "H_kv": 3072,
+ "I": 8192,
+ "head_dim": 128,
+ "num_heads": 24,
+ "num_kv_heads": 24,
+ "act": "swiglu",
},
}
def main():
parser = argparse.ArgumentParser(description="Real vllm-xpu-kernels benchmark")
- parser.add_argument("--preset", default="llama3_8b", choices=list(MODEL_PRESETS.keys()))
+ parser.add_argument(
+ "--preset", default="llama3_8b", choices=list(MODEL_PRESETS.keys())
+ )
parser.add_argument("--m", type=int, default=2048)
parser.add_argument("--iters", type=int, default=200)
parser.add_argument("--all", action="store_true", help="Run all presets")
@@ -196,7 +241,9 @@ def main():
num_kv_heads = cfg["num_kv_heads"]
M = args.m
- print(f"\n--- {cfg['name']} (H={H}, H_kv={H_kv}, I={I}, head_dim={head_dim}) ---")
+ print(
+ f"\n--- {cfg['name']} (H={H}, H_kv={H_kv}, I={I}, head_dim={head_dim}) ---"
+ )
weight = torch.randn(max(H, H_kv), dtype=torch.bfloat16, device="xpu")
@@ -207,12 +254,24 @@ def main():
t_act = bench_silu_and_mul(ops, M, I, iters=args.iters)
else:
t_act = bench_gelu_and_mul(ops, M, I, iters=args.iters)
- t_rope = bench_rotary_embedding(ops, M, num_heads, num_kv_heads, head_dim, iters=args.iters)
+ t_rope = bench_rotary_embedding(
+ ops, M, num_heads, num_kv_heads, head_dim, iters=args.iters
+ )
# Pipeline of all ops
t_pipeline = bench_pipeline_standalone(
- ops, M, H, H_kv, I, head_dim, num_heads, num_kv_heads,
- eps, cfg["act"], iters=args.iters)
+ ops,
+ M,
+ H,
+ H_kv,
+ I,
+ head_dim,
+ num_heads,
+ num_kv_heads,
+ eps,
+ cfg["act"],
+ iters=args.iters,
+ )
# Memory traffic estimates (bytes)
bf = 2
@@ -220,19 +279,38 @@ def main():
fused_rms_bytes = M * H * bf * 4 + H * bf # read inp+res, write both + weight
act_bytes = M * 2 * I * bf + M * I * bf # read [M,2I], write [M,I]
Q_dim = num_heads * head_dim
- rope_bytes = M * Q_dim * bf * 2 + M * head_dim * bf # read+write Q + read cos_sin
- pipeline_bytes_total = rms_bytes + rope_bytes + M * H_kv * bf * 2 + H_kv * bf + fused_rms_bytes + act_bytes
-
- print(f" rms_norm: {t_rms:.4f} ms ({rms_bytes / (t_rms * 1e-3) / 1e9:.1f} GB/s)")
- print(f" fused_add_rms_norm: {t_fused_rms:.4f} ms ({fused_rms_bytes / (t_fused_rms * 1e-3) / 1e9:.1f} GB/s)")
- print(f" {cfg['act']}: {t_act:.4f} ms ({act_bytes / (t_act * 1e-3) / 1e9:.1f} GB/s)")
- print(f" rotary_embedding: {t_rope:.4f} ms ({rope_bytes / (t_rope * 1e-3) / 1e9:.1f} GB/s)")
- print(f" ---")
- print(f" Pipeline (all ops): {t_pipeline:.4f} ms ({pipeline_bytes_total / (t_pipeline * 1e-3) / 1e9:.1f} GB/s)")
+ rope_bytes = (
+ M * Q_dim * bf * 2 + M * head_dim * bf
+ ) # read+write Q + read cos_sin
+ pipeline_bytes_total = (
+ rms_bytes
+ + rope_bytes
+ + M * H_kv * bf * 2
+ + H_kv * bf
+ + fused_rms_bytes
+ + act_bytes
+ )
+
+ print(
+ f" rms_norm: {t_rms:.4f} ms ({rms_bytes / (t_rms * 1e-3) / 1e9:.1f} GB/s)"
+ )
+ print(
+ f" fused_add_rms_norm: {t_fused_rms:.4f} ms ({fused_rms_bytes / (t_fused_rms * 1e-3) / 1e9:.1f} GB/s)"
+ )
+ print(
+ f" {cfg['act']}: {t_act:.4f} ms ({act_bytes / (t_act * 1e-3) / 1e9:.1f} GB/s)"
+ )
+ print(
+ f" rotary_embedding: {t_rope:.4f} ms ({rope_bytes / (t_rope * 1e-3) / 1e9:.1f} GB/s)"
+ )
+ print(" ---")
+ print(
+ f" Pipeline (all ops): {t_pipeline:.4f} ms ({pipeline_bytes_total / (t_pipeline * 1e-3) / 1e9:.1f} GB/s)"
+ )
print(f" Sum of individual: {t_rms + t_fused_rms + t_act + t_rope:.4f} ms")
# Structured output
- print(f"\n=== STRUCTURED OUTPUT ===")
+ print("\n=== STRUCTURED OUTPUT ===")
print(f"VLLM_REAL: {preset_name}")
print(f"MODEL: {cfg['name']}")
print(f"DIMS: M={M} H={H} H_kv={H_kv} I={I} head_dim={head_dim}")
@@ -241,7 +319,9 @@ def main():
print(f"ACTIVATION: {t_act:.4f} ms {cfg['act']}")
print(f"ROPE: {t_rope:.4f} ms")
print(f"OPS_PIPELINE: {t_pipeline:.4f} ms")
- print(f"OPS_PIPELINE_BW: {pipeline_bytes_total / (t_pipeline * 1e-3) / 1e9:.1f} GB/s")
+ print(
+ f"OPS_PIPELINE_BW: {pipeline_bytes_total / (t_pipeline * 1e-3) / 1e9:.1f} GB/s"
+ )
if __name__ == "__main__":