From 8469d1eb8e9006065ea345b6da7f777f182cd667 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 23 Jul 2026 02:30:37 +0000 Subject: [PATCH] perf: N-based tile selection for XPU bf16/fp16 MoE grouped GEMM --- .../wrapper/include/sycl_tla_moe.hpp | 94 ++++++++++++++++--- .../ark/test/README_MOE_PREFILL_PERF.md | 23 +++++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 22 +++++ 3 files changed, 125 insertions(+), 14 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp index cff5626de8..e5a2b59df9 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp @@ -6,8 +6,10 @@ #pragma once +#include #include #include +#include #include #include @@ -53,27 +55,32 @@ namespace moe_detail { using namespace cute; using namespace MoE; -// Helper to choose TiledMMA based on element types -template -auto choose_tiled_mma(TA* A, TB* B) { +// Helper to choose TiledMMA for a given work-group tile / sub-group layout. +// +// The MMA atom (``XE_DPAS_TT<8, float, ...>``) is fixed; only the work-group +// tile (``WGTile``) and the sub-group tiling (``SGLayout``) vary between tile +// policies. Because every bf16/fp16 policy below keeps the same number of +// sub-group rows in M (8) the per-sub-group tile stays 32x64x32, so the same +// 2D block copy atoms remain valid across all of them. +template +auto choose_tiled_mma() { using TA_non_CV = cutlass::platform::remove_cv_t; using TB_non_CV = cutlass::platform::remove_cv_t; auto op = XE_DPAS_TT<8, float, TA_non_CV, TB_non_CV>{}; - using WGTile = Shape<_256, _128, _32>; // 256x128 WG tile size - using SGLayout = Layout, Stride<_2, _1, _0>>; // 8x2 SG tiling, n-major - using MMA = typename TiledMMAHelper, Layout, SGLayout>::TiledMMA; return MMA{}; } -// Unique kernel name tag -template +// Unique kernel name tag. The tile policy (WGTile / SGLayout) is part of the +// tag so each policy specialization produces a distinct SYCL kernel name. +template class MoEGemmKernel; // MOE GEMM launcher template -template +template void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const ElementB* weights, const ElementS* scales, ElementD* outputs, const int gemm_n, const int gemm_k, int* num_rows_per_expert_device, const int num_experts) { @@ -86,7 +93,7 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen auto dummy_group_problem_shape = cutlass::gemm::GroupProblemShape>{1, &dummy_problem_shape, nullptr}; - using TileShape = Shape<_256, _128, _32>; + using TileShape = WGTile; using ClusterShape = Shape<_1, _1, _1>; auto scheduler_params = PersistentTileSchedulerXeMoE::to_underlying_arguments( @@ -97,7 +104,7 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen scheduler_params, dummy_group_problem_shape, TileShape{}, ClusterShape{}, hw_info, PersistentTileSchedulerXeMoE::Arguments{1, RasterOrderOptions::AlongN}); - auto mma = choose_tiled_mma(activations, weights); + auto mma = choose_tiled_mma(); auto MaxThreadsPerWorkgroup = size(mma); dim3 local_range{static_cast(MaxThreadsPerWorkgroup), 1, 1}; @@ -110,7 +117,7 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen syclex::properties kernel_props{syclex::sub_group_size<16>, intelex::grf_size<256>}; - auto event = q->parallel_for>( + auto event = q->parallel_for>( sycl::nd_range<3>(global, local), kernel_props, [=](auto) { MoE::MoEGEMM, XE_LOAD_2D_VNNI<16, 32, 16, 16>, XE_STORE_2D<16, 8, 32>, 'R', 'R', 'R'>(activations, weights, scales, outputs, mma, num_rows_per_expert_device, num_experts, gemm_n, @@ -121,6 +128,65 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen event.wait(); } +// Whether the N-based tile-policy heuristic is enabled (default on). +// +// Set ``ARK_MOE_GEMM_FIXED_TILE`` to a truthy value ("1"/"true"/"on"/"yes") +// to always use the historical fixed 256x128 (8x2) tile regardless of N. +// This provides an escape hatch should a specific device regress with the +// wider tiles. +inline bool moe_gemm_fixed_tile() { + const char* env = std::getenv("ARK_MOE_GEMM_FIXED_TILE"); + if (env == nullptr) { + return false; + } + std::string v(env); + for (auto& c : v) { + c = static_cast(std::tolower(static_cast(c))); + } + return !(v == "0" || v == "false" || v == "off" || v == "no" || v.empty()); +} + +// Select the work-group tile policy from the output width ``N`` and dispatch, +// mirroring the ``w16a16`` large-M heuristic in vllm-xpu-kernels grouped GEMM: +// +// * N <= 64 -> 256x64x32, SGLayout 8x1 +// * N <= 512 -> 256x128x32, SGLayout 8x2 (historical default) +// * N > 512 -> 256x256x32, SGLayout 8x4 +// +// Prefill routes many tokens per expert (large M), so the taller/wider N tile +// increases sub-group utilization and reduces the number of work-group tiles +// launched for the large-N up/down projections. All three policies share the +// same per-sub-group tile (32x64x32), so the copy atoms in +// ``moe_gemm_launcher`` remain valid. +template +void moe_gemm_dispatch(sycl::queue* q, const Element* activations, const Element* weights, const Element* scales, + Element* outputs, const int gemm_n, const int gemm_k, int* num_rows_per_expert_device, + const int num_experts) { + using N64 = Shape<_256, _64, _32>; + using SG64 = Layout, Stride<_1, _1, _0>>; + using N128 = Shape<_256, _128, _32>; + using SG128 = Layout, Stride<_2, _1, _0>>; + using N256 = Shape<_256, _256, _32>; + using SG256 = Layout, Stride<_4, _1, _0>>; + + if (moe_gemm_fixed_tile()) { + moe_gemm_launcher<'R', 'R', N128, SG128, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + return; + } + + if (gemm_n <= 64) { + moe_gemm_launcher<'R', 'R', N64, SG64, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + } else if (gemm_n <= 512) { + moe_gemm_launcher<'R', 'R', N128, SG128, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + } else { + moe_gemm_launcher<'R', 'R', N256, SG256, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + } +} + } // namespace moe_detail // Public MOE GEMM API @@ -129,7 +195,7 @@ inline void moe_gemm(sycl::queue* q, void* activations, void* weights, void* sca switch (dtype) { case BTLA_DTYPE::BF16: { using Element = cutlass::bfloat16_t; - moe_detail::moe_gemm_launcher<'R', 'R', Element, Element, Element, Element>( + moe_detail::moe_gemm_dispatch( q, static_cast(activations), static_cast(weights), static_cast(scales), static_cast(outputs), N, K, num_tokens_per_expert, num_experts); @@ -137,7 +203,7 @@ inline void moe_gemm(sycl::queue* q, void* activations, void* weights, void* sca } case BTLA_DTYPE::F16: { using Element = cutlass::half_t; - moe_detail::moe_gemm_launcher<'R', 'R', Element, Element, Element, Element>( + moe_detail::moe_gemm_dispatch( q, static_cast(activations), static_cast(weights), static_cast(scales), static_cast(outputs), N, K, num_tokens_per_expert, num_experts); diff --git a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md index ff97deaa13..296701271d 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -137,6 +137,29 @@ uneven E=8 8 4096 4096 610 28.9012 10.1234 2. **Speedup**: Higher is better - shows performance gain over baseline 3. **Latency (ms)**: Lower is better - actual kernel execution time +## BF16/FP16 Prefill Tile Selection + +The BF16/FP16 grouped GEMM (`ark.moe_gemm`, measured in `test_perf_fp` and +used as the accumulation path by the quantized prefill kernels) selects its +work-group tile from the output width `N`, mirroring the `w16a16` large-M +policy heuristic in +[`vllm-project/vllm-xpu-kernels`](https://github.com/vllm-project/vllm-xpu-kernels): + +| Condition | Work-group tile | Sub-group layout | +| ----------- | --------------- | ---------------- | +| `N <= 64` | `256x64x32` | `8x1` | +| `N <= 512` | `256x128x32` | `8x2` (historical default) | +| `N > 512` | `256x256x32` | `8x4` | + +Prefill routes many tokens per expert (large M), so the wider `256x256` tile +for large-`N` up/down projections raises sub-group utilization and lowers the +number of work-group tiles launched. All three policies keep the same +per-sub-group tile (`32x64x32`), so the 2D block copy atoms are shared. +Implemented in `sycl_tla_moe.hpp`. + +Set `ARK_MOE_GEMM_FIXED_TILE=1` to force the historical fixed `256x128` (`8x2`) +tile regardless of `N` (escape hatch for per-device tuning/regressions). + ## Integration with CI/CD This test can be integrated into performance regression testing: diff --git a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md index 732cfec70c..9bf41b13c1 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF_CN.md @@ -84,6 +84,28 @@ test_moe_prefill_perf.py 2. **Speedup**: 越大越好 — 表示相对基线的性能提升 3. **Latency (ms)**: 越小越好 — 实际 kernel 执行时间 +## BF16/FP16 Prefill Tile 选择 + +BF16/FP16 分组 GEMM (`ark.moe_gemm`，在 `test_perf_fp` 中测量，同时被量化 +prefill kernel 用作累加路径) 会根据输出宽度 `N` 选择 work-group tile， +对齐 +[`vllm-project/vllm-xpu-kernels`](https://github.com/vllm-project/vllm-xpu-kernels) +中 `w16a16` 的 large-M 策略启发式： + +| 条件 | Work-group tile | Sub-group layout | +| ----------- | --------------- | ---------------- | +| `N <= 64` | `256x64x32` | `8x1` | +| `N <= 512` | `256x128x32` | `8x2`（历史默认） | +| `N > 512` | `256x256x32` | `8x4` | + +Prefill 每个专家会路由多个 token（M 较大），因此对 large-`N` 的 up/down 投影 +使用更宽的 `256x256` tile 可提升 sub-group 利用率并减少启动的 work-group tile +数量。三种策略保持相同的 per-sub-group tile (`32x64x32`)，因此共享同一套 2D +block copy atom。实现见 `sycl_tla_moe.hpp`。 + +设置 `ARK_MOE_GEMM_FIXED_TILE=1` 可强制无视 `N` 始终使用历史固定的 `256x128` +(`8x2`) tile（用于逐设备调优/回归的应急开关）。 + ## 相关文件 - `test_moe.py`: MoE GEMM 的正确性测试