diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5a1dbc7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,16 @@ +build/ +install/ +__pycache__/ +*.pyc +.venv/ +fulsim_logs_*/ +autotune/*.out +tests/*.out +*.out +# generated sweep data / plots / reference papers +autotune/*.csv +autotune/*.xlsx +tests/*.csv +tests/*.pdf +tests/*.png +*.pdf diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..2fb6583 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,135 @@ +cmake_minimum_required(VERSION 3.22 FATAL_ERROR) + +# Re-entry guard for add_subdirectory / FetchContent nesting +if(xe_fuse_LOADED) + return() +endif() +set(xe_fuse_LOADED ON) + +project(XeFuse + VERSION 0.1.0 + DESCRIPTION "Fused GEMM+epilogue kernels for Intel Xe (Battlemage/BMG)" + LANGUAGES CXX) + +# ── Options ──────────────────────────────────────────────────────────────────── +option(XE_FUSE_BUILD_TESTS "Build xe-fuse tests" ${PROJECT_IS_TOP_LEVEL}) +option(XE_FUSE_BUILD_EXAMPLES "Build xe-fuse examples" ${PROJECT_IS_TOP_LEVEL}) + +set(SYCL_TLA_DIR "" CACHE PATH + "Path to an existing sycl-tla checkout. Leave empty to auto-fetch via FetchContent.") +set(SYCL_TLA_GIT_TAG "main" CACHE STRING + "sycl-tla git tag/branch/SHA to fetch when SYCL_TLA_DIR is empty.") +set(DPCPP_SYCL_TARGET "intel_gpu_bmg_g31" CACHE STRING + "Comma-separated SYCL device target(s): intel_gpu_bmg_g31, intel_gpu_bmg_g21, intel_gpu_pvc, spir64, ...") + +# ── C++ standard ────────────────────────────────────────────────────────────── +set(CMAKE_CXX_STANDARD 17 CACHE STRING "" FORCE) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +if(NOT CMAKE_CXX_COMPILER_ID MATCHES "IntelLLVM") + message(WARNING + "xe-fuse targets icpx (Intel oneAPI DPC++). " + "Pass -DCMAKE_TOOLCHAIN_FILE=cmake/toolchain-intel-xe.cmake or set " + "CMAKE_CXX_COMPILER=icpx. Other compilers are unsupported.") +endif() + +# ── sycl-tla cmake options (must be set BEFORE sycl-tla's CMakeLists runs) ──── +set(CUTLASS_ENABLE_SYCL ON CACHE BOOL "" FORCE) +set(CUTLASS_ENABLE_TOOLS ON CACHE BOOL "" FORCE) # needed for cutlass_tools_util_includes +set(CUTLASS_ENABLE_HEADERS_ONLY OFF CACHE BOOL "" FORCE) +set(CUTLASS_ENABLE_EXAMPLES OFF CACHE BOOL "" FORCE) +set(CUTLASS_ENABLE_TESTS OFF CACHE BOOL "" FORCE) +set(CUTLASS_ENABLE_LIBRARY OFF CACHE BOOL "" FORCE) +set(CUTLASS_ENABLE_BENCHMARKS OFF CACHE BOOL "" FORCE) +set(CUTLASS_ENABLE_PROFILER OFF CACHE BOOL "" FORCE) + +# ── Locate or fetch sycl-tla ────────────────────────────────────────────────── +if(SYCL_TLA_DIR) + if(NOT EXISTS "${SYCL_TLA_DIR}/CMakeLists.txt") + message(FATAL_ERROR + "SYCL_TLA_DIR='${SYCL_TLA_DIR}' does not contain CMakeLists.txt. " + "Check the path or leave SYCL_TLA_DIR empty to auto-fetch.") + endif() + message(STATUS "xe-fuse: using sycl-tla at ${SYCL_TLA_DIR}") + add_subdirectory(${SYCL_TLA_DIR} ${CMAKE_BINARY_DIR}/_sycl_tla EXCLUDE_FROM_ALL) +else() + message(STATUS "xe-fuse: fetching sycl-tla (tag=${SYCL_TLA_GIT_TAG}) via FetchContent") + include(FetchContent) + FetchContent_Declare(sycl_tla + GIT_REPOSITORY https://github.com/intel/sycl-tla.git + GIT_TAG ${SYCL_TLA_GIT_TAG} + GIT_SHALLOW ON) + FetchContent_MakeAvailable(sycl_tla) + set(SYCL_TLA_DIR ${sycl_tla_SOURCE_DIR} CACHE PATH "" FORCE) + message(STATUS "xe-fuse: sycl-tla fetched to ${SYCL_TLA_DIR}") +endif() + +# ── DPCPP::DPCPP (normally created by sycl-tla; be defensive) ───────────────── +if(NOT TARGET DPCPP::DPCPP) + list(APPEND CMAKE_MODULE_PATH "${SYCL_TLA_DIR}/cmake") + find_package(DPCPP REQUIRED) +endif() + +# ── xe_fuse INTERFACE library ───────────────────────────────────────────────── +add_library(xe_fuse INTERFACE) +add_library(xe_fuse::xe_fuse ALIAS xe_fuse) + +target_include_directories(xe_fuse INTERFACE + $ + $) + +target_compile_definitions(xe_fuse INTERFACE + CUTLASS_ENABLE_SYCL + SYCL_INTEL_TARGET) + +target_link_libraries(xe_fuse INTERFACE + CUTLASS + DPCPP::DPCPP) + +# ── Internal harness target (tests + examples only, NOT installed) ───────────── +# Provides sycl_common.hpp / helper.h from sycl-tla examples/common and +# the CUTLASS util headers (tools/util/include + MKL). +add_library(_xe_fuse_test_harness INTERFACE) +target_include_directories(_xe_fuse_test_harness INTERFACE + ${SYCL_TLA_DIR}/examples/common) +target_link_libraries(_xe_fuse_test_harness INTERFACE + xe_fuse + nvidia::cutlass::tools::util) + +# ── Helper function ──────────────────────────────────────────────────────────── +include(cmake/XeFuseHelpers.cmake) + +# ── Subdirectories ───────────────────────────────────────────────────────────── +if(XE_FUSE_BUILD_TESTS) + add_subdirectory(tests) +endif() +if(XE_FUSE_BUILD_EXAMPLES) + add_subdirectory(examples) +endif() + +# ── Install rules ────────────────────────────────────────────────────────────── +include(GNUInstallDirs) +include(CMakePackageConfigHelpers) + +install(TARGETS xe_fuse + EXPORT XeFuseTargets + INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +install(DIRECTORY include/ + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +install(EXPORT XeFuseTargets + FILE XeFuseTargets.cmake + NAMESPACE xe_fuse:: + DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/XeFuse) +configure_package_config_file( + cmake/XeFuseConfig.cmake.in + ${CMAKE_CURRENT_BINARY_DIR}/XeFuseConfig.cmake + INSTALL_DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/XeFuse) +write_basic_package_version_file( + ${CMAKE_CURRENT_BINARY_DIR}/XeFuseConfigVersion.cmake + VERSION ${PROJECT_VERSION} + COMPATIBILITY AnyNewerVersion) +install(FILES + ${CMAKE_CURRENT_BINARY_DIR}/XeFuseConfig.cmake + ${CMAKE_CURRENT_BINARY_DIR}/XeFuseConfigVersion.cmake + DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/XeFuse) diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..e988b42 --- /dev/null +++ b/Makefile @@ -0,0 +1,74 @@ +# xe-fuse wrapper Makefile +# +# Usage: +# make # configure + build (auto-fetches sycl-tla) +# make SYCL_TLA_DIR=/path/to/sycl-tla # use existing sycl-tla checkout +# make SYCL_TARGET=intel_gpu_pvc # different GPU target +# make BUILD_TYPE=Debug # debug build +# make tests # build only tests +# make examples # build only examples +# make install PREFIX=/opt/xe-fuse # install headers + cmake package +# make clean # remove build dir +# make help # show this help + +BUILD_DIR ?= build +BUILD_TYPE ?= Release +SYCL_TARGET ?= intel_gpu_bmg_g31 +SYCL_TLA_DIR ?= +PREFIX ?= $(abspath install) +NPROC := $(shell nproc 2>/dev/null || echo 8) + +CMAKE_ARGS := \ + -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) \ + -DDPCPP_SYCL_TARGET=$(SYCL_TARGET) \ + -DXE_FUSE_BUILD_TESTS=ON \ + -DXE_FUSE_BUILD_EXAMPLES=ON + +ifneq ($(SYCL_TLA_DIR),) + CMAKE_ARGS += -DSYCL_TLA_DIR=$(abspath $(SYCL_TLA_DIR)) +endif + +.PHONY: all configure build tests examples install clean help + +all: build + +configure: + cmake -B $(BUILD_DIR) \ + -DCMAKE_TOOLCHAIN_FILE=$(abspath cmake/toolchain-intel-xe.cmake) \ + $(CMAKE_ARGS) \ + $(CURDIR) + +build: configure + cmake --build $(BUILD_DIR) --parallel $(NPROC) + +tests: configure + cmake --build $(BUILD_DIR) --parallel $(NPROC) --target $(shell \ + cmake --build $(BUILD_DIR) --target help 2>/dev/null | grep '^test_\|^tile_sweep\|^gemm_\|^moe_tile' | awk '{print $$1}' | tr '\n' ' ') + @echo "" + @echo "Test binaries are in $(BUILD_DIR)/tests/" + @echo "Submit GPU jobs with: sbatch tests/run_tests.sh" + +examples: configure + cmake --build $(BUILD_DIR) --parallel $(NPROC) --target moe_expert_builder moe_expert_fused + @echo "Example binaries are in $(BUILD_DIR)/examples/" + +install: build + cmake --install $(BUILD_DIR) --prefix $(PREFIX) + +clean: + rm -rf $(BUILD_DIR) + +help: + @echo "Targets : all configure build tests examples install clean" + @echo "" + @echo "Variables:" + @echo " BUILD_DIR = $(BUILD_DIR) (output directory)" + @echo " BUILD_TYPE = $(BUILD_TYPE) (Release | Debug)" + @echo " SYCL_TARGET = $(SYCL_TARGET) (e.g. intel_gpu_bmg_g31, intel_gpu_pvc)" + @echo " SYCL_TLA_DIR = $(if $(SYCL_TLA_DIR),$(SYCL_TLA_DIR),(empty — auto-fetch))" + @echo " PREFIX = $(PREFIX) (cmake --install prefix)" + @echo "" + @echo "Examples:" + @echo " make SYCL_TLA_DIR=~/sycl-tla" + @echo " make SYCL_TARGET=intel_gpu_pvc BUILD_TYPE=Debug" + @echo " make install PREFIX=/opt/xe-fuse" diff --git a/README.md b/README.md index d62e6dd..864caa5 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,333 @@ -# il-opensource-template -![GitHub License](https://img.shields.io/github/license/IntelLabs/il-opensource-template) -[![OpenSSF Scorecard](https://api.scorecard.dev/projects/github.com/IntelLabs/il-opensource-template/badge)](https://scorecard.dev/viewer/?uri=github.com/IntelLabs/il-opensource-template) - +# xe-fuse + +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. + +## Quick Start + +### Using a preset kernel + +```bash +# Generate a fused GEMM + RMSNorm + SwiGLU kernel +python3 autotune/generate_kernel.py --preset k2 -o /tmp/kernel.cpp + +# Compile +icpx -fsycl -DCUTLASS_ENABLE_SYCL -DSYCL_INTEL_TARGET \ + -I $SYCL_TLA_DIR/include -I $SYCL_TLA_DIR/tools/util/include \ + -I $SYCL_TLA_DIR/examples/common \ + -I $XE_FUSE_DIR/include \ + -O2 -std=c++17 -fsycl-targets=spir64_gen \ + -o /tmp/kernel /tmp/kernel.cpp + +# Run +/tmp/kernel --m=4096 --n=4096 --k=4096 --iterations=200 +``` + +### Available kernel presets + +``` +python3 autotune/generate_kernel.py --list-presets + +GEMM epilogue presets: + k1 D = acc * R[m] (RMSNorm) + k0a D = gamma[n] * (acc + residual) (residual + weight) + k2 D = SwiGLU(acc * R[m]) (RMSNorm + SwiGLU) + k2_geglu D = GeGLU(acc * R[m]) (RMSNorm + GeGLU) + k3 D = RoPE(acc, cos_sin) (positional encoding) + k4 D = RoPE(acc * R[m], cos_sin) (RMSNorm + RoPE, composed tree) + k4v2 D = RoPE(acc * R[m], cos_sin) (RMSNorm + RoPE, merged) + +Merged visitors (flat tree): + k1v2 D = acc * R[m] (merged ScaleRows) + k2v2 D = SwiGLU(acc * R[m]) (merged scale + SwiGLU) + k2v2_geglu D = GeGLU(acc * R[m]) (merged scale + GeGLU) + +INT8 quantization epilogues: + w8a8_dequant D_bf16 = int32_acc * scale_token[m] * scale_channel[n] + w8a8_dequant_biased D_bf16 = ... + bias[n] + +Standalone presets (unfused baselines): + sa_scale_rows D[m,n] *= scale[m] + sa_residual_gamma D = gamma[n] * (D + residual) + sa_swiglu D = SwiGLU(D) + sa_rope_scaled D = RoPE(D * scale[m], cos_sin) +``` + +### Available model presets + +``` +python3 autotune/generate_pipeline.py --list-presets + +Available model presets: + Name Model H H_kv I FFN RoPE + llama3_8b LLaMA 3 8B 4096 1024 14336 swiglu yes + llama2_7b LLaMA 2 7B 4096 4096 11008 swiglu yes + llama3_70b LLaMA 3 70B 8192 1024 28672 swiglu yes + gemma2_9b Gemma 2 9B 3584 2048 14336 geglu yes + gemma2_27b Gemma 2 27B 4608 2048 36864 geglu yes + mistral_7b Mistral 7B 4096 1024 14336 swiglu yes + qwen25_7b Qwen 2.5 7B 3584 512 18944 swiglu yes + qwen25_72b Qwen 2.5 72B 8192 1024 29568 swiglu yes + phi3_mini Phi-3 Mini 3.8B 3072 3072 8192 swiglu yes + phi3_medium Phi-3 Medium 14B 5120 5120 17920 swiglu yes +``` + +Each preset defines H (hidden dim), H_kv (KV head dim for GQA), I (intermediate/FFN dim), +activation type, and whether RoPE is used. The pipeline generator maps these to the correct +kernel variants automatically. + +### Writing a kernel in C++ + +```cpp +#include "xe-fuse/builder/epilogue_builder.hpp" +namespace b = xe_fuse::builder; + +using TileShape = cute::Shape; + +// Pick your epilogue — just compose the ops you need: +using EVT = b::GeLU>; // GEMM + RMSNorm + GeLU + +// Build the full GEMM kernel: +using Kernel = b::MakeGemm; +using Gemm = typename Kernel::Gemm; +``` + +## Builder API Reference + +### Data Sources + +| Alias | Description | +|-------|-------------| +| `Acc` | GEMM accumulator (frg_acc) | +| `AuxLoad` | Load M×N auxiliary tensor via Block2D | +| `ColBroadcast` | Per-row vector: `scale[m]` broadcast across columns | +| `RowBroadcast` | Per-column vector: `scale[n]` broadcast across rows | + +### Binary Ops + +| Alias | Formula | +|-------|---------| +| `Mul` | `A * B` element-wise | +| `Add` | `A + B` element-wise | +| `ScaleRows` | `Input * scale[m]` — per-row scaling (RMSNorm) | +| `ScaleCols` | `gamma[n] * Input` — per-column scaling | +| `AddResidual` | `acc + AuxLoad(residual)` | +| `BiasAdd` | `acc + bias[n]` — per-column bias | + +### Activation Functions + +| Alias | Formula | Used By | +|-------|---------|---------| +| `GeLU` | `x * 0.5 * (1 + erf(x/√2))` | BERT, GPT-2/3/4, Gemma | +| `GeLUTanh` | tanh approximation of GeLU | Common fast variant | +| `SiLU` | `x * sigmoid(x)` (Swish) | LLaMA, Mistral | +| `ReLU` | `max(0, x)` | Older models | +| `Sigmoid` | `1 / (1 + exp(-x))` | General | + +### Pairwise Ops (lane shuffle) + +| Alias | Formula | Notes | +|-------|---------|-------| +| `SwiGLU` | `silu(gate) * up` on adjacent pairs | LLaMA, Mistral, Qwen | +| `GeGLU` | `gelu(gate) * up` on adjacent pairs | Gemma 2 | +| `RoPE` | Rotary position embedding on acc | Direct on accumulator | +| `RoPEComposed` | RoPE on pre-processed input | Two-child visitor | +| `RoPEScaled` | Merged scale + RoPE | Flat tree, fewer dispatches | + +### Quantization (INT8/W8A8) + +| Alias | Formula | Notes | +|-------|---------|-------| +| `DequantW8A8` | `int32_acc * scale_token[m] * scale_channel[n]` | INT8 GEMM → bf16 output | +| `DequantW8A8Biased` | `... + bias[n]` | Same with per-channel bias | + +INT8 GEMM uses `int8_t` A/B inputs, `int32_t` accumulator, `bf16` output, and `AlignmentAB=32` +for 256-bit INT8 loads via the `XE_8x16x32_S32S8S8S32_TT` MMA atom. + +### Merged Visitors + +| Alias | Formula | Notes | +|-------|---------|-------| +| `ScaleRowsMerged` | `acc * R[m]` in one visitor | Flat tree: no AccFetch/MulCompute nodes | +| `SwiGLUScaled` | `SwiGLU(acc * R[m])` in one visitor | Scale + shuffle + silu(gate)*up | +| `GeGLUScaled` | `GeGLU(acc * R[m])` in one visitor | Same with GeLU activation | + +These read `frg_acc` directly and do all math in one `visit()` call. + +### Patterns + +| Alias | Description | +|-------|-------------| +| `DualOutput` | Split-tree: evaluate Input once, store to aux buffer, apply Output for primary D | + +### Kernel Builder + +```cpp +// MakeGemm +using K = b::MakeGemm; +using Gemm = typename K::Gemm; // ready to instantiate and run + +// INT8 GEMM with custom alignment: +using K = b::MakeGemm; +``` + +## Composition Examples + +```cpp +// BERT: GEMM + bias + GeLU +using EVT = b::GeLU>; + +// LLaMA / Mistral / Qwen: GEMM + RMSNorm + SwiGLU (pairwise) +using EVT = b::SwiGLU>; + +// Gemma 2: GEMM + RMSNorm + GeGLU (pairwise) +using EVT = b::GeGLU>; + +// GPT-2: GEMM + bias + GeLU_tanh +using EVT = b::GeLUTanh>; + +// K0 dual output: store raw sum for rstd + gamma-weighted primary output +using Input = b::AddResidual; +using Output = b::ScaleCols; +using EVT = b::DualOutput; + +// Custom: GEMM + residual + RMSNorm + RoPE +using Step1 = b::AddResidual; +using Step2 = b::ScaleRows; +using EVT = b::RoPEComposed; + +// W8A8 INT8 dequantization: int8×int8 GEMM → bf16 with per-token/per-channel scales +using EVT = b::DequantW8A8; +using K = b::MakeGemm; + +// Merged visitors (flat tree): +using EVT = b::ScaleRowsMerged; // same as ScaleRows but flat +using EVT = b::SwiGLUScaled; // same as SwiGLU> but flat +``` + +## Standalone Ops + +For unfused baselines or non-GEMM use: + +```cpp +#include "xe-fuse/standalone/ops.hpp" + +auto q = compat::get_default_queue(); +xe_fuse::standalone::scale_rows(q, data, scale, M, N, L); +xe_fuse::standalone::gelu(q, data, M, N, L); +xe_fuse::standalone::swiglu(q, data, M, N, L); +xe_fuse::standalone::geglu(q, data, M, N, L); +xe_fuse::standalone::rope_scaled(q, data, tmp, scale, cos_sin, M, N, L); +``` + +## Code Generation + +The `autotune/` directory contains Python tooling for generating kernel and pipeline benchmarks: + +| Tool | Description | +|------|-------------| +| `generate_kernel.py` | Single-kernel C++ from preset or JSON spec | +| `generate_pipeline.py` | Full model pipeline C++ from architecture preset | +| `model_presets.py` | Model architecture configs (LLaMA, Gemma, Mistral, Qwen, Phi-3) | +| `pipeline_template.cpp.j2` | Jinja2 template for pipeline benchmarks | +| `kernel_template.cpp.j2` | Jinja2 template for single-kernel benchmarks | +| `tile_selector.py` | Tile shape selector | +| `run_kernel.sh` | Compile + benchmark a generated kernel with structured output | + +The pipeline generator maps model architectures to kernel variants: +- SwiGLU models (LLaMA, Mistral, Qwen) use K2 with `b::SwiGLU<>` +- GeGLU models (Gemma 2) use K2 with `b::GeGLU<>` +- GQA models get correct H_kv dimensions for V/K projections +- RoPE is conditionally included based on architecture + +## MoE Expert Batched GEMM + +For Mixture-of-Experts models (LLaMA 4, Mixtral, DeepSeek-V3, DBRX), xe-fuse provides +batched expert kernels that process all routed experts in a single launch using the +CUTLASS L (batch) dimension: + +```cpp +#include "xe-fuse/kernels/gemm_moe_expert.hpp" + +// SwiGLU variant — one kernel launch for all experts +using Config = xe_fuse::MoEExpertSwiGLU; +using Gemm = Config::Gemm; + +// Launch: {M_per_expert, N=2*I, K=H, L=num_experts} +// Each batch slice gets its own fused RMSNorm + SwiGLU epilogue +``` + +Instead of 16 separate GEMM + SwiGLU kernel launches (one per expert), a single +batched launch processes all experts. The fused epilogue runs per-batch-slice, so +each expert gets its own scale vector and SwiGLU activation without extra launches. + +### MoE architectures supported + +| Model | H | Expert I | Experts | top_k | Expert GEMM (N×K) | +|-------|---|----------|---------|-------|--------------------| +| LLaMA 4 Scout | 5120 | 8192 | 16 | 1 | 16384×5120 | +| LLaMA 4 Maverick | 5120 | 8192 | 128 | 1 | 16384×5120 | +| Mixtral 8x7B | 4096 | 14336 | 8 | 2 | 28672×4096 | +| Mixtral 8x22B | 6144 | 16384 | 8 | 2 | 32768×6144 | +| DeepSeek-V3 | 7168 | 2048 | 256 | 8 | 4096×7168 | +| DBRX | 6144 | 10752 | 16 | 4 | 21504×6144 | + +## Autotune Tile Selection + +xe-fuse includes a tile shape auto-selector based on empirical sweep data (792 configurations +across 7 tiles × 5 kernels × 32 GEMM shapes from 10+ model architectures). + +```bash +# Generate pipeline with autotuned tiles for target sequence length +python3 autotune/generate_pipeline.py --preset llama3_8b --autotune --seq-len 128 -o pipeline.cpp + +# Or override tile for a single kernel +python3 autotune/generate_kernel.py --preset k2 --tile 64x128x32 -o kernel.cpp +python3 autotune/generate_kernel.py --preset k1 --tile auto --m 128 --n 4096 --k 4096 -o kernel.cpp +``` + +The tile selector (`autotune/tile_selector.py`) uses empirically-derived heuristics: +tile_M tracks the problem M (64 for M≤64, 128 for M≤128, 256 for M≥256), +tile_N adapts to the output dimension, and K is always 32 (XMX constraint). + +## File Layout + +``` +xe-fuse/ +├── include/xe-fuse/ +│ ├── builder/epilogue_builder.hpp — Builder API (start here) +│ ├── visitors/ +│ │ ├── xe_elementwise_compute.hpp — Unary ops: GeLU, SiLU, ReLU, Sigmoid +│ │ ├── xe_pairwise_compute.hpp — Pairwise: SwiGLU, GeGLU, generic pairs +│ │ ├── xe_rope_compute.hpp — RoPE: 3 visitor variants +│ │ └── xe_scalerows_compute.hpp — Merged: ScaleRows, SwiGLUScaled, GeGLUScaled +│ ├── kernels/ +│ │ ├── gemm_rmsnorm.hpp — K1: GEMM + RMSNorm +│ │ ├── gemm_rmsnorm_swiglu.hpp — K2: GEMM + RMSNorm + SwiGLU +│ │ ├── gemm_rmsnorm_rope.hpp — K4: GEMM + RMSNorm + RoPE +│ │ ├── gemm_residual_gamma.hpp — K0a: GEMM + residual + gamma +│ │ ├── gemm_dual_output.hpp — K0: split-tree dual output +│ │ ├── gemm_moe_expert.hpp — MoE: batched expert GEMM + SwiGLU/GeGLU +│ │ └── ... +│ └── standalone/ops.hpp — Element-wise SYCL baselines +├── autotune/ +│ ├── generate_kernel.py — Single-kernel code generator (--tile auto) +│ ├── generate_pipeline.py — Model pipeline code generator (--autotune) +│ ├── tile_selector.py — Tile shape heuristics for BMG-G31 +│ ├── model_presets.py — Architecture configs (10+ models) +│ ├── pipeline_template.cpp.j2 — Pipeline Jinja2 template +│ ├── kernel_template.cpp.j2 — Kernel Jinja2 template +│ └── run_kernel.sh — Compile + benchmark a generated kernel +├── examples/ +│ ├── moe_expert_builder.cpp — MoE via builder API (L=num_experts) +│ └── moe_expert_fused.cpp — MoE via kernel template +└── tests/ + ├── run_vllm_comparison.sh — xe-fuse vs vllm-xpu three-way comparison + ├── run_bench_e2e.sh — Full pipeline comparison (requires torch + vllm-xpu) + └── run_bench_vllm_real.sh — Real vllm-xpu standalone ops benchmark +``` + diff --git a/autotune/generate_kernel.py b/autotune/generate_kernel.py new file mode 100644 index 0000000..18f73c4 --- /dev/null +++ b/autotune/generate_kernel.py @@ -0,0 +1,788 @@ +#!/usr/bin/env python3 +""" +xe-fuse kernel generator — produces C++ benchmark files from kernel specs. + +This is the foundation for LLM-driven kernel optimization (Phase 3). +An LLM agent describes a kernel as a Python dict, this script generates +the C++ source, and the runner compiles + benchmarks it on GPU. + +Usage: + uv run generate_kernel.py --spec spec.json --output kernel.cpp + uv run generate_kernel.py --preset k1 --output kernel.cpp + +Kernel spec format (JSON): +{ + "name": "K4v2_merged", + "evt_description": "D = RoPE(acc * R[m], cos_sin) via merged visitor", + "tile_shape": "_256, _256, _32", + "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} + ], + "evt_args": [ + "// child 0: ColBroadcast", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "", + "// child 1: AuxLoad", + "typename b::AuxLoad::Arguments cs_args;", + "cs_args.ptr_aux = block_cos_sin.get();", + "cs_args.null_default = float(0);", + "cs_args.dAux = cutlass::make_cute_packed_stride(", + " cute::Stride, int64_t>{}, make_shape(M, N, L));", + "", + "// root: XeRoPEScaledCompute", + "typename xe_fuse::XeRoPEScaledCompute::Arguments rope_args{};", + "", + "typename EVT::Arguments evt_args{scale_args, cs_args, rope_args};" + ] +} +""" + +import argparse +import json +import os +import sys +from datetime import datetime +from pathlib import Path + +PRESETS = { + "k1": { + "name": "K1_RmsNorm", + "evt_description": "D = acc * R[m]", + "tile_shape": "_256, _256, _32", + "evt_typedefs": "using EVT = b::ScaleRows;", + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "evt_args": [ + "typename b::Acc::Arguments accum_args{};", + "", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "", + "typename b::MulOp<>::Arguments mul_args{};", + "", + "typename EVT::Arguments evt_args{accum_args, scale_args, mul_args};" + ] + }, + "k3": { + "name": "K3_RoPE", + "evt_description": "D = RoPE(acc, cos_sin)", + "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} + ], + "evt_args": [ + "typename b::AuxLoad::Arguments cs_args;", + "cs_args.ptr_aux = block_cos_sin.get();", + "cs_args.null_default = float(0);", + "cs_args.dAux = cutlass::make_cute_packed_stride(", + " cute::Stride, int64_t>{}, make_shape(M, N, L));", + "", + "typename xe_fuse::XeRoPECompute::Arguments rope_args{};", + "", + "typename EVT::Arguments evt_args{cs_args, rope_args};" + ] + }, + "k4": { + "name": "K4_RmsNormRoPE", + "evt_description": "D = RoPE(acc * R[m], cos_sin) via composed tree", + "tile_shape": "_256, _256, _32", + "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} + ], + "evt_args": [ + "// child 0: ScaleRows (inner tree)", + "typename b::Acc::Arguments accum_args{};", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "typename b::MulOp::Arguments mul_args{};", + "typename b::ScaleRows::Arguments rms_args{accum_args, scale_args, mul_args};", + "", + "// child 1: AuxLoad", + "typename b::AuxLoad::Arguments cs_args;", + "cs_args.ptr_aux = block_cos_sin.get();", + "cs_args.null_default = float(0);", + "cs_args.dAux = cutlass::make_cute_packed_stride(", + " 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};" + ] + }, + "k4v2": { + "name": "K4v2_RoPEScaled", + "evt_description": "D = RoPE(acc * R[m], cos_sin) via merged visitor (flat tree)", + "tile_shape": "_256, _256, _32", + "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} + ], + "evt_args": [ + "// child 0: ColBroadcast", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "", + "// child 1: AuxLoad", + "typename b::AuxLoad::Arguments cs_args;", + "cs_args.ptr_aux = block_cos_sin.get();", + "cs_args.null_default = float(0);", + "cs_args.dAux = cutlass::make_cute_packed_stride(", + " cute::Stride, int64_t>{}, make_shape(M, N, L));", + "", + "// root: XeRoPEScaledCompute (merged)", + "typename xe_fuse::XeRoPEScaledCompute::Arguments rope_args{};", + "", + "typename EVT::Arguments evt_args{scale_args, cs_args, rope_args};" + ] + }, + "k0a": { + "name": "K0a_ResidualGamma", + "evt_description": "D = gamma[n] * (acc + residual)", + "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} + ], + "evt_args": [ + "// RowBroadcast (child 0 of outer Mul)", + "typename b::RowBroadcast<0, TileShape, float>::Arguments gamma_args;", + "gamma_args.ptr_row = block_gamma.get();", + "gamma_args.null_default = float(1);", + "gamma_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)};", + "", + "// Inner Add tree (child 1 of outer Mul): Acc + AuxLoad", + "typename b::Acc::Arguments accum_args{};", + "typename b::AuxLoad::Arguments res_args;", + "res_args.ptr_aux = block_residual.get();", + "res_args.null_default = bf16(0);", + "res_args.dAux = cutlass::make_cute_packed_stride(", + " cute::Stride, int64_t>{}, make_shape(M, N, L));", + "typename b::AddOp<>::Arguments add_args{};", + "typename b::AddResidual::Arguments inner_args{accum_args, res_args, add_args};", + "", + "// Outer Mul", + "typename b::MulOp<>::Arguments mul_args{};", + "typename EVT::Arguments evt_args{gamma_args, inner_args, mul_args};" + ] + }, + "k2": { + "name": "K2_RmsNormSwiGLU", + "evt_description": "D = SwiGLU(acc * R[m])", + "tile_shape": "_256, _256, _32", + "evt_typedefs": "using EVT = b::SwiGLU>;", + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "evt_args": [ + "// Inner: ScaleRows", + "typename b::Acc::Arguments accum_args{};", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "typename b::MulOp::Arguments mul_args{};", + "typename b::ScaleRows::Arguments rms_args{accum_args, scale_args, mul_args};", + "", + "// Outer: SwiGLU", + "typename xe_fuse::XePairwiseCompute::Arguments swiglu_args{};", + "typename EVT::Arguments evt_args{rms_args, swiglu_args};" + ] + }, + "k2_geglu": { + "name": "K2_RmsNormGeGLU", + "evt_description": "D = GeGLU(acc * R[m])", + "tile_shape": "_256, _256, _32", + "evt_typedefs": "using EVT = b::GeGLU>;", + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "evt_args": [ + "// Inner: ScaleRows", + "typename b::Acc::Arguments accum_args{};", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "typename b::MulOp::Arguments mul_args{};", + "typename b::ScaleRows::Arguments rms_args{accum_args, scale_args, mul_args};", + "", + "// Outer: GeGLU", + "typename xe_fuse::XePairwiseCompute::Arguments geglu_args{};", + "typename EVT::Arguments evt_args{rms_args, geglu_args};" + ] + }, + "k1v2": { + "name": "K1v2_ScaleRowsMerged", + "evt_description": "D = acc * R[m] via merged visitor (flat tree)", + "tile_shape": "_256, _256, _32", + "evt_typedefs": "using EVT = b::ScaleRowsMerged;", + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "evt_args": [ + "// child 0: ColBroadcast", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "", + "// root: XeScaleRowsCompute (merged)", + "typename xe_fuse::XeScaleRowsCompute::Arguments visitor_args{};", + "", + "typename EVT::Arguments evt_args{scale_args, visitor_args};" + ] + }, + "k2v2": { + "name": "K2v2_SwiGLUScaled", + "evt_description": "D = SwiGLU(acc * R[m]) via merged visitor (flat tree)", + "tile_shape": "_256, _256, _32", + "evt_typedefs": "using EVT = b::SwiGLUScaled;", + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "evt_args": [ + "// child 0: ColBroadcast", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "", + "// root: XeScaleRowsSwiGLUCompute (merged)", + "typename xe_fuse::XeScaleRowsSwiGLUCompute::Arguments visitor_args{};", + "", + "typename EVT::Arguments evt_args{scale_args, visitor_args};" + ] + }, + "k2v2_geglu": { + "name": "K2v2_GeGLUScaled", + "evt_description": "D = GeGLU(acc * R[m]) via merged visitor (flat tree)", + "tile_shape": "_256, _256, _32", + "evt_typedefs": "using EVT = b::GeGLUScaled;", + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "evt_args": [ + "// child 0: ColBroadcast", + "typename b::ColBroadcast<0, TileShape, float>::Arguments scale_args;", + "scale_args.ptr_col = block_scale.get();", + "scale_args.null_default = float(1);", + "scale_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "", + "// root: XeScaleRowsGeGLUCompute (merged)", + "typename xe_fuse::XeScaleRowsGeGLUCompute::Arguments visitor_args{};", + "", + "typename EVT::Arguments evt_args{scale_args, visitor_args};" + ] + }, + "w8a8_dequant": { + "name": "W8A8_Dequant", + "evt_description": "D_bf16 = int32_acc * scale_token[m] * scale_channel[n]", + "tile_shape": "_256, _256, _32", + "element_a": "int8_t", + "element_b": "int8_t", + "element_d": "bf16", + "element_acc": "int32_t", + "element_compute": "float", + "alignment_ab": 32, + "alignment_cd": 8, + "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} + ], + "evt_args": [ + "// Inner Mul: Acc * scale_token[m] (ColBroadcast, Idx=0)", + "typename b::Acc::Arguments accum_args{};", + "typename b::ColBroadcast<0, TileShape, float>::Arguments token_args;", + "token_args.ptr_col = block_scale_token.get();", + "token_args.null_default = float(1);", + "token_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "typename b::MulOp<>::Arguments inner_mul_args{};", + "", + "// Outer Mul: (acc * scale_token) * scale_channel[n] (RowBroadcast, Idx=1)", + "typename b::RowBroadcast<1, TileShape, float>::Arguments channel_args;", + "channel_args.ptr_row = block_scale_channel.get();", + "channel_args.null_default = float(1);", + "channel_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)};", + "typename b::MulOp<>::Arguments outer_mul_args{};", + "", + "// 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};" + ] + }, + "w8a8_dequant_biased": { + "name": "W8A8_DequantBiased", + "evt_description": "D_bf16 = int32_acc * scale_token[m] * scale_channel[n] + bias[n]", + "tile_shape": "_256, _256, _32", + "element_a": "int8_t", + "element_b": "int8_t", + "element_d": "bf16", + "element_acc": "int32_t", + "element_compute": "float", + "alignment_ab": 32, + "alignment_cd": 8, + "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} + ], + "evt_args": [ + "// Inner Mul: Acc * scale_token[m] (ColBroadcast, Idx=0)", + "typename b::Acc::Arguments accum_args{};", + "typename b::ColBroadcast<0, TileShape, float>::Arguments token_args;", + "token_args.ptr_col = block_scale_token.get();", + "token_args.null_default = float(1);", + "token_args.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)};", + "typename b::MulOp<>::Arguments inner_mul_args{};", + "", + "// Middle Mul: (acc * scale_token) * scale_channel[n] (RowBroadcast, Idx=1)", + "typename b::RowBroadcast<1, TileShape, float>::Arguments channel_args;", + "channel_args.ptr_row = block_scale_channel.get();", + "channel_args.null_default = float(1);", + "channel_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)};", + "typename b::MulOp<>::Arguments mid_mul_args{};", + "", + "// Outer Add: dequant + bias[n] (RowBroadcast, Idx=2)", + "typename b::RowBroadcast<2, TileShape, float>::Arguments bias_args;", + "bias_args.ptr_row = block_bias.get();", + "bias_args.null_default = float(0);", + "bias_args.dRow = {cute::Int<0>{}, cute::Int<1>{}, static_cast(N)};", + "typename b::AddOp<>::Arguments add_args{};", + "", + "// Compose: EVT = Add(Mul(Mul(Acc, ColBcast), RowBcast), RowBcast_bias)", + "using InnerMul = b::Mul>;", + "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};" + ] + } +} + + +def generate_aux_allocations(aux_data: list[dict]) -> str: + lines = [] + type_map = {"float": "float", "bf16": "bf16", "int8": "int8_t", "int32": "int32_t"} + for aux in aux_data: + ctype = type_map.get(aux["type"], aux["type"]) + 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" initialize_block(block_{name}, {seed});") + return "\n".join(lines) + + +def generate_cpp(spec: dict, defaults: dict | None = None) -> str: + defaults = defaults or {} + template_path = Path(__file__).parent / "kernel_template.cpp.j2" + + if template_path.exists(): + try: + from jinja2 import Template + with open(template_path) as f: + tmpl = Template(f.read()) + + elem_a = spec.get("element_a", "bf16") + elem_b = spec.get("element_b", "bf16") + elem_d = spec.get("element_d", "bf16") + elem_acc = spec.get("element_acc", "float") + elem_compute = spec.get("element_compute", "float") + align_ab = spec.get("alignment_ab", 8) + 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}") + + return tmpl.render( + kernel_name=spec["name"], + timestamp=datetime.now().isoformat(), + evt_description=spec["evt_description"], + tile_shape=spec.get("tile_shape", "_256, _256, _32"), + evt_typedefs=spec["evt_typedefs"], + aux_allocations=generate_aux_allocations(spec.get("aux_data", [])), + evt_args_construction=" " + "\n ".join(spec.get("evt_args", [])), + default_m=defaults.get("m", 4096), + default_n=defaults.get("n", 4096), + default_k=defaults.get("k", 4096), + default_iterations=defaults.get("iterations", 200), + default_verify=defaults.get("verify", 0), + has_verify=False, + element_a=elem_a, + element_b=elem_b, + element_d=elem_d, + element_acc=elem_acc, + element_compute=elem_compute, + make_gemm_extra=make_gemm_extra, + ) + except ImportError: + pass + + # Fallback: inline template (no jinja2 dependency) + return generate_cpp_inline(spec, defaults) + + +def generate_cpp_inline(spec: dict, defaults: dict | None = None) -> str: + defaults = defaults or {} + m = defaults.get("m", 4096) + n = defaults.get("n", 4096) + k = defaults.get("k", 4096) + iters = defaults.get("iterations", 200) + verify = defaults.get("verify", 0) + + aux_alloc = generate_aux_allocations(spec.get("aux_data", [])) + evt_args_code = " " + "\n ".join(spec.get("evt_args", [])) + + # Element types and alignments (with backward-compatible defaults) + elem_a = spec.get("element_a", "bf16") + elem_b = spec.get("element_b", "bf16") + elem_d = spec.get("element_d", "bf16") + elem_acc = spec.get("element_acc", "float") + elem_compute = spec.get("element_compute", "float") + align_ab = spec.get("alignment_ab", 8) + align_cd = spec.get("alignment_cd", 8) + + # 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}") + + return f"""\ +// Auto-generated xe-fuse kernel benchmark +// Kernel: {spec["name"]} +// Generated: {datetime.now().isoformat()} +// EVT tree: {spec["evt_description"]} + +#include "xe-fuse/builder/epilogue_builder.hpp" + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/packed_stride.hpp" + +#include "sycl_common.hpp" +#include "helper.h" + +#include + +using namespace cute; +namespace b = xe_fuse::builder; +using bf16 = cutlass::bfloat16_t; + +using TileShape = Shape<{spec.get("tile_shape", "_256, _256, _32")}>; + +// ---- EVT tree ---- +{spec["evt_typedefs"]} + +using KernelConfig = b::MakeGemm; +using GemmOp = typename KernelConfig::Gemm; + +struct Options {{ + int m = {m}, n = {n}, k = {k}, l = 1; + int iterations = {iters}; + int verify = {verify}; + + void parse(int argc, char const** args) {{ + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, {m}); + cmd.get_cmd_line_argument("n", n, {n}); + cmd.get_cmd_line_argument("k", k, {k}); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, {iters}); + cmd.get_cmd_line_argument("verify", verify, {verify}); + }} +}}; + +int main(int argc, const char** argv) {{ + Options opts; + opts.parse(argc, argv); + + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + int M = opts.m, N = opts.n, K = opts.k, L = opts.l; + + using StrideA = typename GemmOp::GemmKernel::StrideA; + using StrideB = typename GemmOp::GemmKernel::StrideB; + + auto stride_A = cutlass::make_cute_packed_stride(StrideA{{}}, make_shape(M, K, L)); + auto stride_B = cutlass::make_cute_packed_stride(StrideB{{}}, make_shape(N, K, L)); + auto stride_C = cutlass::make_cute_packed_stride(KernelConfig::StrideC{{}}, make_shape(M, N, L)); + auto stride_D = cutlass::make_cute_packed_stride(KernelConfig::StrideD{{}}, make_shape(M, N, L)); + + cutlass::DeviceAllocation<{elem_a}> block_A(static_cast(M) * K * L); + cutlass::DeviceAllocation<{elem_b}> block_B(static_cast(K) * N * L); + cutlass::DeviceAllocation<{elem_d}> block_D(static_cast(M) * N * L); + + initialize_block(block_A, 2023); + initialize_block(block_B, 2022); + +{aux_alloc} + +{evt_args_code} + + typename GemmOp::GemmKernel::EpilogueArguments epilogue_args{{ + evt_args, nullptr, stride_C, block_D.get(), stride_D + }}; + typename GemmOp::GemmKernel::Arguments arguments{{ + cutlass::gemm::GemmUniversalMode::kGemm, + {{M, N, K, L}}, + {{block_A.get(), stride_A, block_B.get(), stride_B}}, + epilogue_args, hw_info + }}; + + GemmOp gemm_op; + size_t ws_size = GemmOp::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(ws_size); + + auto status = gemm_op.can_implement(arguments); + if (status != cutlass::Status::kSuccess) {{ + std::cerr << "can_implement failed" << std::endl; + return 1; + }} + CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get())); + CUTLASS_CHECK(gemm_op.run()); + compat::wait(); + + std::cout << "Disposition: launched" << std::endl; + + if (opts.iterations > 0) {{ + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) gemm_op.run(); + compat::wait(); + + float time_s = timer.seconds() / opts.iterations; + double tflops = (2.0 * M * N * K * L) * 1e-12; + std::cout << "Problem Size: " << M << 'x' << N << 'x' << K << 'x' << L << std::endl; + printf("{spec["name"]}: [%4.3f]TFlop/s (%6.4f)ms\\n", tflops / time_s, time_s * 1000); + }} + + return 0; +}} +""" + + +def generate_standalone_cpp(spec: dict, defaults: dict | None = None) -> str: + """Generate a standalone (non-GEMM) kernel benchmark.""" + defaults = defaults or {} + m = defaults.get("m", 4096) + n = defaults.get("n", 4096) + iters = defaults.get("iterations", 200) + + aux_alloc = generate_aux_allocations(spec.get("aux_data", [])) + op_code = "\n ".join(spec.get("op_code", [])) + + return f"""\ +// Auto-generated xe-fuse standalone kernel benchmark +// Op: {spec["name"]} +// Generated: {datetime.now().isoformat()} +// Description: {spec["evt_description"]} + +#include "xe-fuse/standalone/ops.hpp" + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/command_line.h" + +#include "sycl_common.hpp" +#include "helper.h" + +#include + +using namespace cute; +using bf16 = cutlass::bfloat16_t; + +struct Options {{ + int m = {m}, n = {n}, l = 1; + int iterations = {iters}; + + void parse(int argc, char const** args) {{ + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, {m}); + cmd.get_cmd_line_argument("n", n, {n}); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, {iters}); + }} +}}; + +int main(int argc, const char** argv) {{ + Options opts; + opts.parse(argc, argv); + + int M = opts.m, N = opts.n, L = opts.l; + auto q = compat::get_default_queue(); + + cutlass::DeviceAllocation block_D(static_cast(M) * N * L); + initialize_block(block_D, 2023); + +{aux_alloc} + + // Warmup + for (int i = 0; i < 5; ++i) {{ + {op_code} + q.wait(); + }} + + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) {{ + {op_code} + q.wait(); + }} + + float time_s = timer.seconds() / opts.iterations; + size_t bytes_rw = static_cast(M) * N * L * sizeof(bf16) * 2; + double bw_gb = bytes_rw / (time_s * 1e9); + + std::cout << "Problem Size: " << M << 'x' << N << 'x' << L << std::endl; + printf("{spec["name"]}: %6.4f ms %.1f GB/s (%.1f MB r+w)\\n", + time_s * 1000, bw_gb, bytes_rw / (1024.0 * 1024.0)); + + return 0; +}} +""" + + +STANDALONE_PRESETS = { + "sa_scale_rows": { + "name": "SA_ScaleRows", + "evt_description": "D[m,n] *= scale[m] (standalone RMSNorm scaling)", + "standalone": True, + "aux_data": [ + {"name": "scale", "type": "float", "shape": "M * L", "init_seed": 42} + ], + "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} + ], + "op_code": [ + "xe_fuse::standalone::residual_gamma(q, block_D.get(), block_residual.get(),", + " 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);" + ] + }, + "sa_rope_scaled": { + "name": "SA_RoPEScaled", + "evt_description": "D = RoPE(D * scale[m], cos_sin) (standalone)", + "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} + ], + "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);" + ] + } +} + + +def main(): + parser = argparse.ArgumentParser(description="xe-fuse kernel generator") + 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("--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") + args = parser.parse_args() + + if args.list_presets: + print("GEMM epilogue presets:") + for name, spec in PRESETS.items(): + print(f" {name:16s} {spec['evt_description']}") + print("\nStandalone presets:") + for name, spec in STANDALONE_PRESETS.items(): + print(f" {name:16s} {spec['evt_description']}") + return + + if not args.output and not args.list_presets: + parser.error("--output is required unless using --list-presets") + return + + if args.preset: + spec = all_presets[args.preset] + elif args.spec: + with open(args.spec) as f: + spec = json.load(f) + else: + parser.error("Provide --preset or --spec") + return + + # Tile shape override: explicit, auto, or default from preset + 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: + parts = args.tile.replace("x", ", _").lstrip("_") + spec["tile_shape"] = f"_{parts}" + + defaults = {"m": args.m, "n": args.n, "k": args.k, "iterations": args.iterations} + if spec.get("standalone"): + code = generate_standalone_cpp(spec, defaults) + else: + code = generate_cpp(spec, defaults) + + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + with open(args.output, "w") as f: + f.write(code) + + print(f"Generated: {args.output}") + print(f" Kernel: {spec['name']}") + print(f" Tile: {spec.get('tile_shape', '_256, _256, _32')}") + print(f" EVT: {spec['evt_description']}") + + +if __name__ == "__main__": + main() diff --git a/autotune/generate_pipeline.py b/autotune/generate_pipeline.py new file mode 100644 index 0000000..27cdf5e --- /dev/null +++ b/autotune/generate_pipeline.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +""" +xe-fuse pipeline generator — produces model-specific C++ pipeline benchmarks. + +Generates a complete C++ test binary that: + 1. Runs a fused GEMM+epilogue pipeline (rstd -> Q -> V -> K0 -> rstd -> FFN) + 2. Verifies correctness against FP32 references + 3. Benchmarks fused vs unfused (bare GEMM + standalone ops) + 4. Prints structured output for LLM agent parsing + +Usage: + uv run generate_pipeline.py --preset llama3_8b --output test_llama3.cpp + uv run generate_pipeline.py --config custom.json --output test_custom.cpp + uv run generate_pipeline.py --list-presets +""" + +import argparse +import json +import sys +from datetime import datetime +from pathlib import Path + +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: + template_path = Path(__file__).parent / "pipeline_template.cpp.j2" + + if not template_path.exists(): + print(f"ERROR: Template not found at {template_path}", file=sys.stderr) + sys.exit(1) + + try: + from jinja2 import Template + except ImportError: + print("ERROR: jinja2 required. Install with: uv pip install jinja2", file=sys.stderr) + sys.exit(1) + + with open(template_path) as f: + tmpl = Template(f.read()) + + pipeline_parts = [] + pipeline_parts.append(f"Q({'K4' if config['use_rope'] else 'K1'})") + pipeline_parts.append("V(K1)") + pipeline_parts.append("O(K0)") + pipeline_parts.append(f"FFN({config['ffn_activation']})") + + # Tile selection — single tile for the whole pipeline (safest) + # Picks based on the most constrained kernel (K4 RoPE at the Q projection shape) + H = config["H"] + H_kv = config["H_kv"] + N_ffn = 2 * config["I"] if config["gated_ffn"] else config["I"] + tile_vars = {} + + 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 + + return tmpl.render( + model_name=config["name"], + preset_name=preset_name, + timestamp=datetime.now().isoformat(), + default_H=config["H"], + default_H_kv=config["H_kv"], + default_I=config["I"], + use_rope=config["use_rope"], + ffn_activation=config["ffn_activation"], + gated_ffn=config["gated_ffn"], + pipeline_desc=" -> ".join(pipeline_parts), + **tile_vars, + ) + + +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("--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") + args = parser.parse_args() + + if args.list_presets: + list_presets() + return + + if not args.output: + parser.error("--output is required") + + if args.preset: + config = MODEL_PRESETS[args.preset] + preset_name = args.preset + elif args.config: + with open(args.config) as f: + config = json.load(f) + preset_name = Path(args.config).stem + else: + parser.error("Provide --preset or --config") + return + + 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: + f.write(code) + + 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" 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}") + + +if __name__ == "__main__": + main() diff --git a/autotune/kernel_template.cpp.j2 b/autotune/kernel_template.cpp.j2 new file mode 100644 index 0000000..1a367ff --- /dev/null +++ b/autotune/kernel_template.cpp.j2 @@ -0,0 +1,125 @@ +// Auto-generated xe-fuse kernel benchmark +// Kernel: {{ kernel_name }} +// Generated: {{ timestamp }} +// +// EVT tree: {{ evt_description }} + +#include "xe-fuse/builder/epilogue_builder.hpp" + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#include "sycl_common.hpp" +#include "helper.h" + +#include +#include +#include + +using namespace cute; +namespace b = xe_fuse::builder; +using bf16 = cutlass::bfloat16_t; + +using TileShape = Shape<{{ tile_shape }}>; + +// ---- EVT tree definition ---- +{{ evt_typedefs }} + +// ---- Kernel assembly ---- +using KernelConfig = b::MakeGemm; +using GemmOp = typename KernelConfig::Gemm; + +struct Options { + int m = {{ default_m }}, n = {{ default_n }}, k = {{ default_k }}, l = 1; + int iterations = {{ default_iterations }}; + int verify = {{ default_verify }}; + + void parse(int argc, char const** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, {{ default_m }}); + cmd.get_cmd_line_argument("n", n, {{ default_n }}); + cmd.get_cmd_line_argument("k", k, {{ default_k }}); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, {{ default_iterations }}); + cmd.get_cmd_line_argument("verify", verify, {{ default_verify }}); + } +}; + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + int M = opts.m, N = opts.n, K = opts.k, L = opts.l; + + using StrideA = typename GemmOp::GemmKernel::StrideA; + using StrideB = typename GemmOp::GemmKernel::StrideB; + + auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, make_shape(M, K, L)); + auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, make_shape(N, K, L)); + auto stride_C = cutlass::make_cute_packed_stride(KernelConfig::StrideC{}, make_shape(M, N, L)); + auto stride_D = cutlass::make_cute_packed_stride(KernelConfig::StrideD{}, make_shape(M, N, L)); + + // Allocate matrices + cutlass::DeviceAllocation<{{ element_a }}> block_A(static_cast(M) * K * L); + cutlass::DeviceAllocation<{{ element_b }}> block_B(static_cast(K) * N * L); + cutlass::DeviceAllocation<{{ element_d }}> block_D(static_cast(M) * N * L); + + initialize_block(block_A, 2023); + initialize_block(block_B, 2022); + + // Allocate auxiliary data +{{ aux_allocations }} + + // Build EVT arguments +{{ evt_args_construction }} + + typename GemmOp::GemmKernel::EpilogueArguments epilogue_args{ + evt_args, nullptr, stride_C, block_D.get(), stride_D + }; + typename GemmOp::GemmKernel::Arguments arguments{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, L}, + {block_A.get(), stride_A, block_B.get(), stride_B}, + epilogue_args, hw_info + }; + + GemmOp gemm_op; + size_t ws_size = GemmOp::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(ws_size); + + CUTLASS_CHECK(gemm_op.can_implement(arguments)); + CUTLASS_CHECK(gemm_op.initialize(arguments, workspace.get())); + CUTLASS_CHECK(gemm_op.run()); + compat::wait(); + +{% if has_verify %} + if (opts.verify) { +{{ verify_code }} + } else { + std::cout << "Disposition: skipped" << std::endl; + } +{% else %} + std::cout << "Disposition: skipped (no reference)" << std::endl; +{% endif %} + + if (opts.iterations > 0) { + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) gemm_op.run(); + compat::wait(); + + float time_s = timer.seconds() / opts.iterations; + double tflops = (2.0 * M * N * K * L) * 1e-12; + std::cout << "Problem Size: " << M << 'x' << N << 'x' << K << 'x' << L << std::endl; + printf("{{ kernel_name }}: [%4.3f]TFlop/s (%6.4f)ms\n", tflops / time_s, time_s * 1000); + } + + return 0; +} diff --git a/autotune/model_presets.py b/autotune/model_presets.py new file mode 100644 index 0000000..858b683 --- /dev/null +++ b/autotune/model_presets.py @@ -0,0 +1,131 @@ +""" +xe-fuse model presets — architecture configs for pipeline code generation. + +Each preset defines the dimensions and operation types for one transformer layer. +The pipeline generator uses these to produce model-specific C++ test binaries. + +Usage: + from model_presets import MODEL_PRESETS + config = MODEL_PRESETS["llama3_8b"] + +Custom configs can also be loaded from JSON with the same schema. +""" + +MODEL_PRESETS = { + "llama3_8b": { + "name": "LLaMA 3 8B", + "H": 4096, + "H_kv": 1024, # 8 KV heads * 128 head_dim (GQA) + "I": 14336, # intermediate_size + "num_layers": 32, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "llama2_7b": { + "name": "LLaMA 2 7B", + "H": 4096, + "H_kv": 4096, # MHA (no GQA) + "I": 11008, + "num_layers": 32, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "llama3_70b": { + "name": "LLaMA 3 70B", + "H": 8192, + "H_kv": 1024, # 8 KV heads * 128 head_dim + "I": 28672, + "num_layers": 80, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "gemma2_9b": { + "name": "Gemma 2 9B", + "H": 3584, + "H_kv": 2048, # 8 KV heads * 256 head_dim + "I": 14336, + "num_layers": 42, + "use_rope": True, + "ffn_activation": "geglu", + "gated_ffn": True, + }, + "gemma2_27b": { + "name": "Gemma 2 27B", + "H": 4608, + "H_kv": 2048, # 16 KV heads * 128 head_dim + "I": 36864, + "num_layers": 46, + "use_rope": True, + "ffn_activation": "geglu", + "gated_ffn": True, + }, + "mistral_7b": { + "name": "Mistral 7B", + "H": 4096, + "H_kv": 1024, # 8 KV heads * 128 head_dim + "I": 14336, + "num_layers": 32, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "qwen25_7b": { + "name": "Qwen 2.5 7B", + "H": 3584, + "H_kv": 512, # 4 KV heads * 128 head_dim + "I": 18944, + "num_layers": 28, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "qwen25_72b": { + "name": "Qwen 2.5 72B", + "H": 8192, + "H_kv": 1024, # 8 KV heads * 128 head_dim + "I": 29568, + "num_layers": 80, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "phi3_mini": { + "name": "Phi-3 Mini 3.8B", + "H": 3072, + "H_kv": 3072, # MHA (no GQA) + "I": 8192, + "num_layers": 32, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, + "phi3_medium": { + "name": "Phi-3 Medium 14B", + "H": 5120, + "H_kv": 5120, # MHA (no GQA) + "I": 17920, + "num_layers": 40, + "use_rope": True, + "ffn_activation": "swiglu", + "gated_ffn": True, + }, +} + + +def get_preset(name: str) -> dict: + if name not in MODEL_PRESETS: + available = ", ".join(MODEL_PRESETS.keys()) + raise ValueError(f"Unknown preset '{name}'. Available: {available}") + return MODEL_PRESETS[name] + + +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(" " + "-" * 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}") diff --git a/autotune/pipeline_template.cpp.j2 b/autotune/pipeline_template.cpp.j2 new file mode 100644 index 0000000..1ec548e --- /dev/null +++ b/autotune/pipeline_template.cpp.j2 @@ -0,0 +1,827 @@ +// Auto-generated xe-fuse pipeline benchmark +// Model: {{ model_name }} +// Generated: {{ timestamp }} +// Architecture: H={{ default_H }}, H_kv={{ default_H_kv }}, I={{ default_I }} +// Q/K kernel: {{ "K4 (RMSNorm+RoPE)" if use_rope else "K1 (RMSNorm)" }} +// FFN activation: {{ ffn_activation }}{% if gated_ffn %} (gated, N_ffn=2*I){% else %} (non-gated, N_ffn=I){% endif %} +// +// Pipeline: rstd -> {{ pipeline_desc }} + +#include "xe-fuse/builder/epilogue_builder.hpp" +#include "xe-fuse/kernels/gemm_rmsnorm.hpp" +#include "xe-fuse/kernels/gemm_dual_output.hpp" +#include "xe-fuse/kernels/compute_rstd.hpp" +#include "xe-fuse/standalone/ops.hpp" +#include "xe-fuse/standalone/vllm_ops.hpp" +{% if use_rope %} +#include "xe-fuse/kernels/gemm_rmsnorm_rope.hpp" +{% endif %} + +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/packed_stride.hpp" +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/device/tensor_compare.h" + +#include "sycl_common.hpp" +#include "helper.h" + +#include +#include +#include + +using namespace cute; +using bf16 = cutlass::bfloat16_t; +namespace b = xe_fuse::builder; + +// === Kernel type definitions === + +// Tile shape — autotune-selected or default 256x256x32 +using TileShape = cute::Shape; + +{% if use_rope %} +using QK_Config = xe_fuse::GemmRmsNormRoPE< + cutlass::bfloat16_t, cutlass::bfloat16_t, cutlass::bfloat16_t, + float, float, float, float, TileShape>; +{% else %} +using QK_Config = xe_fuse::GemmRmsNorm< + cutlass::bfloat16_t, cutlass::bfloat16_t, cutlass::bfloat16_t, + float, float, float, TileShape>; +{% endif %} +using K1_Config = xe_fuse::GemmRmsNorm< + cutlass::bfloat16_t, cutlass::bfloat16_t, cutlass::bfloat16_t, + float, float, float, TileShape>; +using K0_Config = xe_fuse::GemmDualOutput< + cutlass::bfloat16_t, cutlass::bfloat16_t, cutlass::bfloat16_t, + cutlass::bfloat16_t, float, float, float, TileShape>; + +{% if ffn_activation == "swiglu" %} +using FFN_EVT = b::SwiGLU>; +{% elif ffn_activation == "geglu" %} +using FFN_EVT = b::GeGLU>; +{% elif ffn_activation == "gelu" %} +using FFN_EVT = b::GeLU>; +{% endif %} +using FFN_Config = b::MakeGemm; +using FFN_Gemm = typename FFN_Config::Gemm; + +struct BareGemm { + using ElementA = bf16; + using ElementB = bf16; + using ElementD = bf16; + using ElementAcc = float; + using ElementCompute = float; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::RowMajor; + using BTS = TileShape; + using StrideC = cute::Stride, int64_t>; + using StrideD = cute::Stride, int64_t>; + + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + BTS, cute::Shape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAcc, ElementCompute, + ElementD, StrideC, 8, ElementD, StrideD, 8, + cutlass::epilogue::collective::EpilogueScheduleAuto + >::CollectiveOp; + + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + cutlass::arch::Xe20, cutlass::arch::OpClassTensorOp, + ElementA, LayoutA, 8, ElementB, LayoutB, 8, + ElementAcc, BTS, cute::Shape, + cutlass::gemm::collective::StageCountAuto, + cutlass::gemm::collective::KernelScheduleAuto + >::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + cute::Shape, CollectiveMainloop, CollectiveEpilogue>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; +}; + +struct Options { + int m = 2048; + int h = {{ default_H }}; + int hkv = {{ default_H_kv }}; + int inter = {{ default_I }}; + int l = 1; + int iterations = 0; + float eps = 1e-6f; + + void parse(int argc, const char** args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("m", m, 2048); + cmd.get_cmd_line_argument("h", h, {{ default_H }}); + cmd.get_cmd_line_argument("hkv", hkv, {{ default_H_kv }}); + cmd.get_cmd_line_argument("inter", inter, {{ default_I }}); + cmd.get_cmd_line_argument("l", l, 1); + cmd.get_cmd_line_argument("iterations", iterations, 0); + } +}; + +bool check(const char* name, bf16 const* ref, bf16 const* actual, size_t count, + float rtol = 0.05f) { + bool ok = cutlass::reference::device::BlockCompareRelativelyEqual( + ref, actual, count, static_cast(rtol), static_cast(rtol)); + std::cout << " " << name << ": " << (ok ? "Passed" : "FAILED") << std::endl; + return ok; +} + +int main(int argc, const char** argv) { + Options opts; + opts.parse(argc, argv); + + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + int M = opts.m, H = opts.h, H_kv = opts.hkv, I = opts.inter, L = opts.l; +{% if gated_ffn %} + int N_ffn = 2 * I; +{% else %} + int N_ffn = I; +{% endif %} + auto q = compat::get_default_queue(); + + std::cout << "============================================================" << std::endl; + std::cout << "xe-fuse: {{ model_name }} Pipeline Test" << std::endl; + std::cout << " M=" << M << " H=" << H << " H_kv=" << H_kv + << " I=" << I << " (N_ffn=" << N_ffn << ")" << std::endl; + std::cout << "============================================================" << std::endl; + + // === Allocations === + size_t mh = static_cast(M) * H * L; + size_t mhkv = static_cast(M) * H_kv * L; + size_t hh = static_cast(H) * H * L; + size_t h_hkv = static_cast(H) * H_kv * L; + size_t h_nffn = static_cast(H) * N_ffn * L; + size_t m_nffn = static_cast(M) * N_ffn * L; + + cutlass::DeviceAllocation x(mh); + cutlass::DeviceAllocation W_q(hh); + cutlass::DeviceAllocation W_v(h_hkv); + cutlass::DeviceAllocation W_o(hh); + cutlass::DeviceAllocation W_ffn(h_nffn); + cutlass::DeviceAllocation gamma(static_cast(H) * L); +{% if use_rope %} + cutlass::DeviceAllocation cos_sin(mh); +{% endif %} + + cutlass::DeviceAllocation R1(static_cast(M) * L); + cutlass::DeviceAllocation Q_out(mh); + cutlass::DeviceAllocation V_out(mhkv); + cutlass::DeviceAllocation attn_out(mh); + cutlass::DeviceAllocation residual1(mh); + cutlass::DeviceAllocation raw_sum(mh); + cutlass::DeviceAllocation R2(static_cast(M) * L); + cutlass::DeviceAllocation ffn_out(m_nffn); + + size_t max_mn = std::max({mh, mhkv, m_nffn}); + cutlass::DeviceAllocation ref_f32(max_mn); + cutlass::DeviceAllocation ref_bf16(max_mn); + cutlass::DeviceAllocation ref_bf16_aux(mh); + cutlass::DeviceAllocation tmp_buf(mh); + + initialize_block(x, 2023); + initialize_block(W_q, 2024); + initialize_block(W_v, 2025); + initialize_block(W_o, 2026); + initialize_block(W_ffn, 2027); + initialize_block(attn_out, 2028); + + { + std::vector h_gamma(static_cast(H) * L); + std::mt19937 rng(42); + std::uniform_real_distribution dist(0.5f, 1.5f); + for (auto& v : h_gamma) v = dist(rng); + q.memcpy(gamma.get(), h_gamma.data(), h_gamma.size() * sizeof(float)); + } +{% if use_rope %} + { + std::vector h_cs(mh); + std::mt19937 rng(2024); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (auto& v : h_cs) v = dist(rng); + q.memcpy(cos_sin.get(), h_cs.data(), h_cs.size() * sizeof(float)); + } +{% endif %} + compat::wait(); + + // === Setup fused GEMM ops === + + auto stride_A_mh = cutlass::make_cute_packed_stride( + typename QK_Config::Gemm::GemmKernel::StrideA{}, make_shape(M, H, L)); + auto stride_B_hh = cutlass::make_cute_packed_stride( + typename QK_Config::Gemm::GemmKernel::StrideB{}, make_shape(H, H, L)); + auto stride_D_mh = cutlass::make_cute_packed_stride( + QK_Config::StrideD{}, make_shape(M, H, L)); +{% if use_rope %} + auto stride_cs = cutlass::make_cute_packed_stride( + QK_Config::StrideCosSin{}, make_shape(M, H, L)); +{% endif %} + + // Q projection + typename QK_Config::Gemm qk_op; +{% if use_rope %} + auto qk_evt = QK_Config::make_evt_args(R1.get(), M, cos_sin.get(), stride_cs); +{% else %} + auto qk_evt = QK_Config::make_evt_args(R1.get(), M); +{% endif %} + typename QK_Config::Gemm::GemmKernel::EpilogueArguments qk_epi{ + qk_evt, nullptr, stride_D_mh, Q_out.get(), stride_D_mh}; + typename QK_Config::Gemm::GemmKernel::Arguments qk_args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, H, H, L}, + {x.get(), stride_A_mh, W_q.get(), stride_B_hh}, + qk_epi, hw_info}; + cutlass::device_memory::allocation qk_ws(QK_Config::Gemm::get_workspace_size(qk_args)); + CUTLASS_CHECK(qk_op.can_implement(qk_args)); + CUTLASS_CHECK(qk_op.initialize(qk_args, qk_ws.get())); + + // V projection + auto stride_A_v = cutlass::make_cute_packed_stride( + typename K1_Config::Gemm::GemmKernel::StrideA{}, make_shape(M, H, L)); + auto stride_B_hkv = cutlass::make_cute_packed_stride( + typename K1_Config::Gemm::GemmKernel::StrideB{}, make_shape(H_kv, H, L)); + auto stride_D_hkv = cutlass::make_cute_packed_stride( + K1_Config::StrideD{}, make_shape(M, H_kv, L)); + + typename K1_Config::Gemm k1_op; + auto k1_evt = K1_Config::make_evt_args(R1.get(), M); + typename K1_Config::Gemm::GemmKernel::EpilogueArguments k1_epi{ + k1_evt, nullptr, stride_D_hkv, V_out.get(), stride_D_hkv}; + typename K1_Config::Gemm::GemmKernel::Arguments k1_args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, H_kv, H, L}, + {x.get(), stride_A_v, W_v.get(), stride_B_hkv}, + k1_epi, hw_info}; + cutlass::device_memory::allocation k1_ws(K1_Config::Gemm::get_workspace_size(k1_args)); + CUTLASS_CHECK(k1_op.can_implement(k1_args)); + CUTLASS_CHECK(k1_op.initialize(k1_args, k1_ws.get())); + + // O projection (K0) + auto stride_A_o = cutlass::make_cute_packed_stride( + typename K0_Config::Gemm::GemmKernel::StrideA{}, make_shape(M, H, L)); + auto stride_B_o = cutlass::make_cute_packed_stride( + typename K0_Config::Gemm::GemmKernel::StrideB{}, make_shape(H, H, L)); + auto stride_aux = cutlass::make_cute_packed_stride( + K0_Config::StrideAux{}, make_shape(M, H, L)); + + typename K0_Config::Gemm k0_op; + auto k0_evt = K0_Config::make_evt_args( + x.get(), stride_D_mh, gamma.get(), H, raw_sum.get(), stride_aux); + typename K0_Config::Gemm::GemmKernel::EpilogueArguments k0_epi{ + k0_evt, nullptr, stride_D_mh, residual1.get(), stride_D_mh}; + typename K0_Config::Gemm::GemmKernel::Arguments k0_args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, H, H, L}, + {attn_out.get(), stride_A_o, W_o.get(), stride_B_o}, + k0_epi, hw_info}; + cutlass::device_memory::allocation k0_ws(K0_Config::Gemm::get_workspace_size(k0_args)); + CUTLASS_CHECK(k0_op.can_implement(k0_args)); + CUTLASS_CHECK(k0_op.initialize(k0_args, k0_ws.get())); + + // FFN projection + auto stride_A_ffn = cutlass::make_cute_packed_stride( + typename FFN_Gemm::GemmKernel::StrideA{}, make_shape(M, H, L)); + auto stride_B_ffn = cutlass::make_cute_packed_stride( + typename FFN_Gemm::GemmKernel::StrideB{}, make_shape(N_ffn, H, L)); + auto stride_D_ffn = cutlass::make_cute_packed_stride( + FFN_Config::StrideD{}, make_shape(M, N_ffn, L)); + + typename b::Acc::Arguments ffn_acc{}; + typename b::ColBroadcast<0, TileShape, float>::Arguments ffn_scale; + ffn_scale.ptr_col = R2.get(); + ffn_scale.null_default = float(1); + ffn_scale.dCol = {cute::Int<1>{}, cute::Int<0>{}, static_cast(M)}; + typename b::MulOp::Arguments ffn_mul{}; + typename b::ScaleRows::Arguments ffn_rms{ffn_acc, ffn_scale, ffn_mul}; +{% if ffn_activation == "swiglu" %} + typename xe_fuse::XePairwiseCompute::Arguments ffn_act{}; +{% elif ffn_activation == "geglu" %} + typename xe_fuse::XePairwiseCompute::Arguments ffn_act{}; +{% elif ffn_activation == "gelu" %} + typename xe_fuse::XeElementwiseCompute::Arguments ffn_act{}; +{% endif %} + typename FFN_EVT::Arguments ffn_evt{ffn_rms, ffn_act}; + + FFN_Gemm ffn_op; + typename FFN_Gemm::GemmKernel::EpilogueArguments ffn_epi{ + ffn_evt, nullptr, stride_D_ffn, ffn_out.get(), stride_D_ffn}; + typename FFN_Gemm::GemmKernel::Arguments ffn_args{ + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N_ffn, H, L}, + {residual1.get(), stride_A_ffn, W_ffn.get(), stride_B_ffn}, + ffn_epi, hw_info}; + cutlass::device_memory::allocation ffn_ws(FFN_Gemm::get_workspace_size(ffn_args)); + CUTLASS_CHECK(ffn_op.can_implement(ffn_args)); + CUTLASS_CHECK(ffn_op.initialize(ffn_args, ffn_ws.get())); + + // === Run fused pipeline === + std::cout << "\n=== Running fused pipeline ===" << std::endl; + + xe_fuse::launch_compute_rstd(q, x.get(), R1.get(), M, H, L, opts.eps); + compat::wait(); + CUTLASS_CHECK(qk_op.run()); + CUTLASS_CHECK(k1_op.run()); + compat::wait(); + CUTLASS_CHECK(k0_op.run()); + compat::wait(); + xe_fuse::launch_compute_rstd(q, raw_sum.get(), R2.get(), M, H, L, opts.eps); + compat::wait(); + CUTLASS_CHECK(ffn_op.run()); + compat::wait(); + + std::cout << "Pipeline complete.\n" << std::endl; + + // === Verify each step === + bool all_passed = true; + +{% if use_rope %} + std::cout << "--- Q = K4(x @ W_q, R1, cos_sin) ---" << std::endl; +{% else %} + std::cout << "--- Q = K1(x @ W_q, R1) ---" << std::endl; +{% endif %} + { + cutlass::TensorRef rA(x.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::TensorRef rB(W_q.get(), cutlass::layout::RowMajor::packed({H, H})); + cutlass::TensorRef rC(ref_f32.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::TensorRef rD(ref_f32.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::reference::device::GemmComplex( + {M, H, H}, 1.0f, rA, cutlass::ComplexTransform::kNone, + rB, cutlass::ComplexTransform::kNone, 0.0f, rC, rD, 0.0f, + L, M*H, H*H, M*H, M*H); + compat::wait(); + + auto* gemm = ref_f32.get(); + auto* r1 = R1.get(); + auto* out = ref_bf16.get(); + int h_val = H; +{% if use_rope %} + auto* cs = cos_sin.get(); + q.parallel_for(sycl::range<1>(mh), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int row = static_cast(i / h_val); + int col = static_cast(i % h_val); + int pair_base = (col / 2) * 2; + int64_t even_idx = static_cast(row) * h_val + pair_base; + int64_t odd_idx = even_idx + 1; + float even_scaled = gemm[even_idx] * r1[row]; + float odd_scaled = gemm[odd_idx] * r1[row]; + float cos_v = cs[even_idx]; + float sin_v = cs[odd_idx]; + float result = (col % 2 == 0) + ? ( even_scaled * cos_v + odd_scaled * sin_v) + : (-even_scaled * sin_v + odd_scaled * cos_v); + out[i] = static_cast(result); + }); +{% else %} + q.parallel_for(sycl::range<1>(mh), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int row = static_cast(i / h_val); + out[i] = static_cast(gemm[i] * r1[row]); + }); +{% endif %} + compat::wait(); + all_passed &= check("{{ 'K4 (GEMM+RMSNorm+RoPE)' if use_rope else 'K1 (GEMM+RMSNorm)' }} -> Q", + ref_bf16.get(), Q_out.get(), mh); + } + + std::cout << "--- V = K1(x @ W_v, R1) ---" << std::endl; + { + cutlass::TensorRef rA(x.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::TensorRef rB(W_v.get(), cutlass::layout::RowMajor::packed({H, H_kv})); + cutlass::TensorRef rC(ref_f32.get(), cutlass::layout::RowMajor::packed({M, H_kv})); + cutlass::TensorRef rD(ref_f32.get(), cutlass::layout::RowMajor::packed({M, H_kv})); + cutlass::reference::device::GemmComplex( + {M, H_kv, H}, 1.0f, rA, cutlass::ComplexTransform::kNone, + rB, cutlass::ComplexTransform::kNone, 0.0f, rC, rD, 0.0f, + L, M*H, H*H_kv, M*H_kv, M*H_kv); + compat::wait(); + + auto* gemm = ref_f32.get(); + auto* r1 = R1.get(); + auto* out = ref_bf16.get(); + int n_val = H_kv; + q.parallel_for(sycl::range<1>(mhkv), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int row = static_cast(i / n_val); + out[i] = static_cast(gemm[i] * r1[row]); + }); + compat::wait(); + all_passed &= check("K1 (GEMM+RMSNorm) -> V", ref_bf16.get(), V_out.get(), mhkv); + } + + std::cout << "--- K0(attn_out @ W_o + x) -> residual1, raw_sum ---" << std::endl; + { + cutlass::TensorRef rA(attn_out.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::TensorRef rB(W_o.get(), cutlass::layout::RowMajor::packed({H, H})); + cutlass::TensorRef rC(ref_f32.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::TensorRef rD(ref_f32.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::reference::device::GemmComplex( + {M, H, H}, 1.0f, rA, cutlass::ComplexTransform::kNone, + rB, cutlass::ComplexTransform::kNone, 0.0f, rC, rD, 0.0f, + L, M*H, H*H, M*H, M*H); + compat::wait(); + + auto* gemm = ref_f32.get(); + auto* x_ptr = x.get(); + auto* gamma_ptr = gamma.get(); + auto* ref_d = ref_bf16.get(); + auto* ref_a = ref_bf16_aux.get(); + int h_val = H; + q.parallel_for(sycl::range<1>(mh), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int col = static_cast(i % h_val); + float sum = gemm[i] + static_cast(x_ptr[i]); + ref_a[i] = static_cast(sum); + ref_d[i] = static_cast(sum * gamma_ptr[col]); + }); + compat::wait(); + all_passed &= check("K0 primary D (gamma * sum)", ref_bf16.get(), residual1.get(), mh); + all_passed &= check("K0 aux (raw sum)", ref_bf16_aux.get(), raw_sum.get(), mh); + } + + std::cout << "--- FFN = K2(residual1 @ W_ffn, R2) [{{ ffn_activation }}] ---" << std::endl; + { + cutlass::TensorRef rA(residual1.get(), cutlass::layout::RowMajor::packed({M, H})); + cutlass::TensorRef rB(W_ffn.get(), cutlass::layout::RowMajor::packed({H, N_ffn})); + cutlass::TensorRef rC(ref_f32.get(), cutlass::layout::RowMajor::packed({M, N_ffn})); + cutlass::TensorRef rD(ref_f32.get(), cutlass::layout::RowMajor::packed({M, N_ffn})); + cutlass::reference::device::GemmComplex( + {M, N_ffn, H}, 1.0f, rA, cutlass::ComplexTransform::kNone, + rB, cutlass::ComplexTransform::kNone, 0.0f, rC, rD, 0.0f, + L, M*H, H*N_ffn, M*N_ffn, M*N_ffn); + compat::wait(); + + auto* gemm = ref_f32.get(); + auto* r2 = R2.get(); + auto* out = ref_bf16.get(); + int n_val = N_ffn; +{% if gated_ffn %} + q.parallel_for(sycl::range<1>(m_nffn), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int row = static_cast(i / n_val); + int col = static_cast(i % n_val); + int pair_base = (col / 2) * 2; + int64_t gate_idx = static_cast(row) * n_val + pair_base; + int64_t up_idx = gate_idx + 1; + float gate = gemm[gate_idx] * r2[row]; + float up = gemm[up_idx] * r2[row]; +{% if ffn_activation == "swiglu" %} + float act = gate / (1.0f + sycl::exp(-gate)); +{% elif ffn_activation == "geglu" %} + float act = gate * 0.5f * (1.0f + sycl::erf(gate * 0.7071067811865475f)); +{% endif %} + out[i] = static_cast(act * up); + }); +{% else %} + q.parallel_for(sycl::range<1>(m_nffn), [=](sycl::id<1> idx) { + int64_t i = idx[0]; + int row = static_cast(i / n_val); + float val = gemm[i] * r2[row]; + out[i] = static_cast(val * 0.5f * (1.0f + sycl::erf(val * 0.7071067811865475f))); + }); +{% endif %} + compat::wait(); + all_passed &= check("FFN ({{ ffn_activation }}) -> ffn_out", ref_bf16.get(), ffn_out.get(), m_nffn); + } + + std::cout << "\n============================================================" << std::endl; + std::cout << "Pipeline: " << (all_passed ? "ALL PASSED" : "SOME FAILED") << std::endl; + std::cout << "============================================================" << std::endl; + + // === Benchmark === + if (opts.iterations > 0) { + std::cout << "\n=== Benchmark: " << opts.iterations << " iterations ===" << std::endl; + std::cout << "Pipeline: " << M << "x" << H << " hidden, " << I << " intermediate\n" << std::endl; + + double flops = 2.0 * M * L * + (static_cast(H) * H + + static_cast(H_kv) * H + + static_cast(H) * H + + static_cast(N_ffn) * H); + + // Fused warmup + xe_fuse::launch_compute_rstd(q, x.get(), R1.get(), M, H, L, opts.eps); + qk_op.run(); k1_op.run(); k0_op.run(); + xe_fuse::launch_compute_rstd(q, raw_sum.get(), R2.get(), M, H, L, opts.eps); + ffn_op.run(); + compat::wait(); + + GPU_Clock timer; + timer.start(); + for (int i = 0; i < opts.iterations; ++i) { + xe_fuse::launch_compute_rstd(q, x.get(), R1.get(), M, H, L, opts.eps); + qk_op.run(); k1_op.run(); k0_op.run(); + xe_fuse::launch_compute_rstd(q, raw_sum.get(), R2.get(), M, H, L, opts.eps); + ffn_op.run(); + } + compat::wait(); + float fused_s = timer.seconds() / opts.iterations; + + printf("Fused pipeline (GEMM + epilogue):\n"); + printf(" Time: %.4f ms\n", fused_s * 1000); + printf(" Aggregate: %.3f TFlop/s\n", flops * 1e-12 / fused_s); + + // --- vllm-equiv pipeline --- + // Uses the same bare GEMM from sycl-tla but with vllm-style fused + // standalone kernels between GEMMs (fused_add_rms_norm, silu_and_mul, etc.) + // This isolates the fusion strategy difference: INTO epilogue vs BETWEEN GEMMs. + float vllm_s = 0.0f; + { + using BareOp = typename BareGemm::Gemm; + using BareKernel = typename BareGemm::Gemm::GemmKernel; + + auto bare_sA_mh = cutlass::make_cute_packed_stride( + typename BareKernel::StrideA{}, make_shape(M, H, L)); + auto bare_sB_hh = cutlass::make_cute_packed_stride( + typename BareKernel::StrideB{}, make_shape(H, H, L)); + auto bare_sD_mh = cutlass::make_cute_packed_stride( + BareGemm::StrideD{}, make_shape(M, H, L)); + auto bare_sB_hkv = cutlass::make_cute_packed_stride( + typename BareKernel::StrideB{}, make_shape(H_kv, H, L)); + auto bare_sD_hkv = cutlass::make_cute_packed_stride( + BareGemm::StrideD{}, make_shape(M, H_kv, L)); + auto bare_sB_ffn = cutlass::make_cute_packed_stride( + typename BareKernel::StrideB{}, make_shape(N_ffn, H, L)); + auto bare_sD_ffn = cutlass::make_cute_packed_stride( + BareGemm::StrideD{}, make_shape(M, N_ffn, L)); + + // vllm needs RMSNorm weight vectors (bf16, matches vllm convention) + cutlass::DeviceAllocation norm_weight1(static_cast(H)); + cutlass::DeviceAllocation norm_weight2(static_cast(H)); + { + std::vector hw(H); + std::mt19937 rng(99); + std::uniform_real_distribution dist(0.8f, 1.2f); + for (auto& v : hw) v = static_cast(dist(rng)); + q.memcpy(norm_weight1.get(), hw.data(), hw.size() * sizeof(bf16)); + q.memcpy(norm_weight2.get(), hw.data(), hw.size() * sizeof(bf16)); + compat::wait(); + } + + // Q GEMM: x @ W_q -> tmp_buf [M, H] + typename BareKernel::Arguments vq_args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H, H, L}, + {x.get(), bare_sA_mh, W_q.get(), bare_sB_hh}, + { {1.0f, 0.0f}, nullptr, bare_sD_mh, tmp_buf.get(), bare_sD_mh}, + hw_info}; + // V GEMM: x @ W_v -> V_out [M, H_kv] + typename BareKernel::Arguments vv_args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H_kv, H, L}, + {x.get(), bare_sA_mh, W_v.get(), bare_sB_hkv}, + { {1.0f, 0.0f}, nullptr, bare_sD_hkv, V_out.get(), bare_sD_hkv}, + hw_info}; + // O GEMM: attn_out @ W_o -> residual1 [M, H] + typename BareKernel::Arguments vo_args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H, H, L}, + {attn_out.get(), bare_sA_mh, W_o.get(), bare_sB_hh}, + { {1.0f, 0.0f}, nullptr, bare_sD_mh, residual1.get(), bare_sD_mh}, + hw_info}; + // FFN GEMM: input @ W_ffn -> ffn_out [M, N_ffn] + typename BareKernel::Arguments vffn_args{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N_ffn, H, L}, + {residual1.get(), bare_sA_mh, W_ffn.get(), bare_sB_ffn}, + { {1.0f, 0.0f}, nullptr, bare_sD_ffn, ffn_out.get(), bare_sD_ffn}, + hw_info}; + + BareOp vq_op, vv_op, vo_op, vffn_op; + cutlass::device_memory::allocation ws_vq(BareOp::get_workspace_size(vq_args)); + cutlass::device_memory::allocation ws_vv(BareOp::get_workspace_size(vv_args)); + cutlass::device_memory::allocation ws_vo(BareOp::get_workspace_size(vo_args)); + cutlass::device_memory::allocation ws_vffn(BareOp::get_workspace_size(vffn_args)); + CUTLASS_CHECK(vq_op.can_implement(vq_args)); + CUTLASS_CHECK(vq_op.initialize(vq_args, ws_vq.get())); + CUTLASS_CHECK(vv_op.can_implement(vv_args)); + CUTLASS_CHECK(vv_op.initialize(vv_args, ws_vv.get())); + CUTLASS_CHECK(vo_op.can_implement(vo_args)); + CUTLASS_CHECK(vo_op.initialize(vo_args, ws_vo.get())); + CUTLASS_CHECK(vffn_op.can_implement(vffn_args)); + CUTLASS_CHECK(vffn_op.initialize(vffn_args, ws_vffn.get())); + + // cos_sin cache for vllm RoPE: [M, rot_dim] with cos in first half, sin in second +{% if use_rope %} + int head_dim = 128; + int num_q_heads = H / head_dim; + int num_kv_heads_rope = H_kv / head_dim; + int rot_dim = head_dim; + cutlass::DeviceAllocation vllm_cos_sin(static_cast(M) * rot_dim); + { + std::vector hcs(static_cast(M) * rot_dim); + std::mt19937 rng(2024); + std::uniform_real_distribution dist(-1.0f, 1.0f); + for (auto& v : hcs) v = dist(rng); + q.memcpy(vllm_cos_sin.get(), hcs.data(), hcs.size() * sizeof(float)); + compat::wait(); + } +{% endif %} + + // Warmup + xe_fuse::vllm_equiv::rms_norm(q, tmp_buf.get(), x.get(), norm_weight1.get(), M, H, opts.eps); + compat::wait(); + CUTLASS_CHECK(vq_op.run()); + compat::wait(); +{% if use_rope %} + xe_fuse::vllm_equiv::rotary_embedding(q, tmp_buf.get(), nullptr, + vllm_cos_sin.get(), head_dim, num_q_heads, num_kv_heads_rope, rot_dim, M); +{% endif %} + CUTLASS_CHECK(vv_op.run()); + compat::wait(); + xe_fuse::vllm_equiv::rms_norm(q, V_out.get(), V_out.get(), norm_weight1.get(), M, H_kv, opts.eps); + compat::wait(); + CUTLASS_CHECK(vo_op.run()); + compat::wait(); + xe_fuse::vllm_equiv::fused_add_rms_norm(q, residual1.get(), x.get(), + norm_weight2.get(), M, H, opts.eps); + compat::wait(); + CUTLASS_CHECK(vffn_op.run()); + compat::wait(); +{% if ffn_activation == "swiglu" %} + xe_fuse::vllm_equiv::silu_and_mul(q, ffn_out.get(), ffn_out.get(), I, M); +{% elif ffn_activation == "geglu" %} + xe_fuse::vllm_equiv::gelu_and_mul(q, ffn_out.get(), ffn_out.get(), I, M); +{% endif %} + compat::wait(); + + // Timed loop + timer.start(); + for (int i = 0; i < opts.iterations; ++i) { + // RMSNorm on input (pre-Q/V) + xe_fuse::vllm_equiv::rms_norm(q, tmp_buf.get(), x.get(), norm_weight1.get(), M, H, opts.eps); + // Q projection + CUTLASS_CHECK(vq_op.run()); +{% if use_rope %} + // RoPE on Q (NeoX-style, head_dim=128) + xe_fuse::vllm_equiv::rotary_embedding(q, tmp_buf.get(), nullptr, + vllm_cos_sin.get(), head_dim, num_q_heads, num_kv_heads_rope, rot_dim, M); +{% endif %} + // V projection + CUTLASS_CHECK(vv_op.run()); + // RMSNorm on V + xe_fuse::vllm_equiv::rms_norm(q, V_out.get(), V_out.get(), norm_weight1.get(), M, H_kv, opts.eps); + // O projection + CUTLASS_CHECK(vo_op.run()); + // fused_add_rms_norm: residual add + RMSNorm for FFN input + xe_fuse::vllm_equiv::fused_add_rms_norm(q, residual1.get(), x.get(), + norm_weight2.get(), M, H, opts.eps); + // FFN GEMM + CUTLASS_CHECK(vffn_op.run()); + // silu_and_mul / gelu_and_mul (contiguous halves -> half-width output) +{% if ffn_activation == "swiglu" %} + xe_fuse::vllm_equiv::silu_and_mul(q, ffn_out.get(), ffn_out.get(), I, M); +{% elif ffn_activation == "geglu" %} + xe_fuse::vllm_equiv::gelu_and_mul(q, ffn_out.get(), ffn_out.get(), I, M); +{% elif ffn_activation == "gelu" %} + xe_fuse::standalone::gelu(q, ffn_out.get(), M, N_ffn, L); +{% endif %} + } + compat::wait(); + vllm_s = timer.seconds() / opts.iterations; + + printf("\nvllm-equiv pipeline (bare GEMM + fused standalone ops):\n"); + printf(" Time: %.4f ms\n", vllm_s * 1000); + printf(" Aggregate: %.3f TFlop/s\n", flops * 1e-12 / vllm_s); + printf(" Ops: rms_norm + GEMM + RoPE + GEMM + rms_norm + GEMM + fused_add_rms_norm + GEMM + {{ ffn_activation }}\n"); + } + + // --- Unfused pipeline --- + using BareOp = typename BareGemm::Gemm; + using BareKernel = typename BareGemm::Gemm::GemmKernel; + + auto bare_sA_mh = cutlass::make_cute_packed_stride( + typename BareKernel::StrideA{}, make_shape(M, H, L)); + auto bare_sB_hh = cutlass::make_cute_packed_stride( + typename BareKernel::StrideB{}, make_shape(H, H, L)); + auto bare_sD_mh = cutlass::make_cute_packed_stride( + BareGemm::StrideD{}, make_shape(M, H, L)); + auto bare_sB_hkv = cutlass::make_cute_packed_stride( + typename BareKernel::StrideB{}, make_shape(H_kv, H, L)); + auto bare_sD_hkv = cutlass::make_cute_packed_stride( + BareGemm::StrideD{}, make_shape(M, H_kv, L)); + auto bare_sB_ffn = cutlass::make_cute_packed_stride( + typename BareKernel::StrideB{}, make_shape(N_ffn, H, L)); + auto bare_sD_ffn = cutlass::make_cute_packed_stride( + BareGemm::StrideD{}, make_shape(M, N_ffn, L)); + + typename BareKernel::Arguments args_bq{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H, H, L}, + {x.get(), bare_sA_mh, W_q.get(), bare_sB_hh}, + { {1.0f, 0.0f}, nullptr, bare_sD_mh, tmp_buf.get(), bare_sD_mh}, + hw_info}; + typename BareKernel::Arguments args_bv{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H_kv, H, L}, + {x.get(), bare_sA_mh, W_v.get(), bare_sB_hkv}, + { {1.0f, 0.0f}, nullptr, bare_sD_hkv, V_out.get(), bare_sD_hkv}, + hw_info}; + typename BareKernel::Arguments args_bo{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, H, H, L}, + {attn_out.get(), bare_sA_mh, W_o.get(), bare_sB_hh}, + { {1.0f, 0.0f}, nullptr, bare_sD_mh, residual1.get(), bare_sD_mh}, + hw_info}; + typename BareKernel::Arguments args_bffn{ + cutlass::gemm::GemmUniversalMode::kGemm, {M, N_ffn, H, L}, + {residual1.get(), bare_sA_mh, W_ffn.get(), bare_sB_ffn}, + { {1.0f, 0.0f}, nullptr, bare_sD_ffn, ffn_out.get(), bare_sD_ffn}, + hw_info}; + + BareOp bare_q_op, bare_v_op, bare_o_op, bare_ffn_op; + cutlass::device_memory::allocation ws_bq(BareOp::get_workspace_size(args_bq)); + cutlass::device_memory::allocation ws_bv(BareOp::get_workspace_size(args_bv)); + cutlass::device_memory::allocation ws_bo(BareOp::get_workspace_size(args_bo)); + cutlass::device_memory::allocation ws_bffn(BareOp::get_workspace_size(args_bffn)); + CUTLASS_CHECK(bare_q_op.can_implement(args_bq)); + CUTLASS_CHECK(bare_q_op.initialize(args_bq, ws_bq.get())); + CUTLASS_CHECK(bare_v_op.can_implement(args_bv)); + CUTLASS_CHECK(bare_v_op.initialize(args_bv, ws_bv.get())); + CUTLASS_CHECK(bare_o_op.can_implement(args_bo)); + CUTLASS_CHECK(bare_o_op.initialize(args_bo, ws_bo.get())); + CUTLASS_CHECK(bare_ffn_op.can_implement(args_bffn)); + CUTLASS_CHECK(bare_ffn_op.initialize(args_bffn, ws_bffn.get())); + + size_t mh_bytes = mh * sizeof(bf16); + + // Unfused warmup + xe_fuse::launch_compute_rstd(q, x.get(), R1.get(), M, H, L, opts.eps); + bare_q_op.run(); + xe_fuse::standalone::scale_rows(q, tmp_buf.get(), R1.get(), M, H, L); +{% if use_rope %} + xe_fuse::standalone::rope(q, Q_out.get(), tmp_buf.get(), cos_sin.get(), M, H, L); +{% else %} + q.memcpy(Q_out.get(), tmp_buf.get(), mh_bytes); +{% endif %} + bare_v_op.run(); + xe_fuse::standalone::scale_rows(q, V_out.get(), R1.get(), M, H_kv, L); + bare_o_op.run(); + xe_fuse::standalone::add_residual(q, residual1.get(), x.get(), M, H, L); + q.memcpy(raw_sum.get(), residual1.get(), mh_bytes); + xe_fuse::standalone::scale_cols(q, residual1.get(), gamma.get(), M, H, L); + xe_fuse::launch_compute_rstd(q, raw_sum.get(), R2.get(), M, H, L, opts.eps); + bare_ffn_op.run(); + xe_fuse::standalone::scale_rows(q, ffn_out.get(), R2.get(), M, N_ffn, L); +{% if ffn_activation == "swiglu" %} + xe_fuse::standalone::swiglu(q, ffn_out.get(), M, N_ffn, L); +{% elif ffn_activation == "geglu" %} + xe_fuse::standalone::geglu(q, ffn_out.get(), M, N_ffn, L); +{% elif ffn_activation == "gelu" %} + xe_fuse::standalone::gelu(q, ffn_out.get(), M, N_ffn, L); +{% endif %} + compat::wait(); + + timer.start(); + for (int i = 0; i < opts.iterations; ++i) { + xe_fuse::launch_compute_rstd(q, x.get(), R1.get(), M, H, L, opts.eps); + bare_q_op.run(); + xe_fuse::standalone::scale_rows(q, tmp_buf.get(), R1.get(), M, H, L); +{% if use_rope %} + xe_fuse::standalone::rope(q, Q_out.get(), tmp_buf.get(), cos_sin.get(), M, H, L); +{% else %} + q.memcpy(Q_out.get(), tmp_buf.get(), mh_bytes); +{% endif %} + bare_v_op.run(); + xe_fuse::standalone::scale_rows(q, V_out.get(), R1.get(), M, H_kv, L); + bare_o_op.run(); + xe_fuse::standalone::add_residual(q, residual1.get(), x.get(), M, H, L); + q.memcpy(raw_sum.get(), residual1.get(), mh_bytes); + xe_fuse::standalone::scale_cols(q, residual1.get(), gamma.get(), M, H, L); + xe_fuse::launch_compute_rstd(q, raw_sum.get(), R2.get(), M, H, L, opts.eps); + bare_ffn_op.run(); + xe_fuse::standalone::scale_rows(q, ffn_out.get(), R2.get(), M, N_ffn, L); +{% if ffn_activation == "swiglu" %} + xe_fuse::standalone::swiglu(q, ffn_out.get(), M, N_ffn, L); +{% elif ffn_activation == "geglu" %} + xe_fuse::standalone::geglu(q, ffn_out.get(), M, N_ffn, L); +{% elif ffn_activation == "gelu" %} + xe_fuse::standalone::gelu(q, ffn_out.get(), M, N_ffn, L); +{% endif %} + } + compat::wait(); + float unfused_s = timer.seconds() / opts.iterations; + + printf("\nUnfused pipeline (bare GEMM + standalone ops):\n"); + printf(" Time: %.4f ms\n", unfused_s * 1000); + printf(" Aggregate: %.3f TFlop/s\n", flops * 1e-12 / unfused_s); + + float speedup = unfused_s / fused_s; + float savings = (1.0f - fused_s / unfused_s) * 100.0f; + printf("\nFusion speedup vs naive: %.2fx (fused saves %.1f%% of pipeline time)\n", speedup, savings); + std::cout << "Note: K projection omitted (same kernel as Q)" << std::endl; + + // Structured output for LLM agent parsing + std::cout << "\n=== STRUCTURED OUTPUT ===" << std::endl; + std::cout << "PIPELINE: {{ preset_name }}" << std::endl; + std::cout << "MODEL: {{ model_name }}" << std::endl; + printf("DIMS: M=%d H=%d H_kv=%d I=%d N_ffn=%d\n", M, H, H_kv, I, N_ffn); + std::cout << "CORRECTNESS: " << (all_passed ? "ALL_PASSED" : "SOME_FAILED") << std::endl; + printf("FUSED: %.4f ms %.3f TFlop/s\n", fused_s * 1000, flops * 1e-12 / fused_s); + printf("VLLM_EQUIV: %.4f ms %.3f TFlop/s\n", vllm_s * 1000, flops * 1e-12 / vllm_s); + printf("UNFUSED: %.4f ms %.3f TFlop/s\n", unfused_s * 1000, flops * 1e-12 / unfused_s); + printf("XE_VS_VLLM: %.2fx\n", vllm_s / fused_s); + printf("XE_VS_NAIVE: %.2fx\n", speedup); + printf("VLLM_VS_NAIVE: %.2fx\n", unfused_s / vllm_s); + std::cout << "GEMMS: Q({{ 'K4' if use_rope else 'K1' }}) V(K1) O(K0) FFN({{ ffn_activation }})" << std::endl; + std::cout << "NOTE: K projection omitted (same kernel as Q)" << std::endl; + } + + return all_passed ? 0 : 1; +} diff --git a/autotune/run_kernel.sh b/autotune/run_kernel.sh new file mode 100644 index 0000000..a3eb9fb --- /dev/null +++ b/autotune/run_kernel.sh @@ -0,0 +1,89 @@ +#!/bin/bash +# xe-fuse: compile and benchmark a generated kernel +# +# Prerequisites: source your oneAPI environment before running, e.g.: +# source /opt/intel/oneapi/setvars.sh +# +# Usage: +# ./run_kernel.sh [extra_args...] +# +# Output is structured for machine parsing: +# BUILD: OK|FAILED +# SPILLS: |none +# : []TFlop/s (