perf: N-based tile selection for XPU bf16/fp16 MoE grouped GEMM

auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp

@@ -6,8 +6,10 @@
 
 #pragma once
 
+#include <cctype>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <stdexcept>
 #include <string>
 
@@ -53,27 +55,32 @@ namespace moe_detail {
 using namespace cute;
 using namespace MoE;
 
-// Helper to choose TiledMMA based on element types
-template <class TA, class TB>
-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 <class WGTile, class SGLayout, class TA, class TB>
+auto choose_tiled_mma() {
   using TA_non_CV = cutlass::platform::remove_cv_t<TA>;
   using TB_non_CV = cutlass::platform::remove_cv_t<TB>;
   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<Shape<_8, _2, _1>, Stride<_2, _1, _0>>;  // 8x2 SG tiling, n-major
-
   using MMA = typename TiledMMAHelper<MMA_Atom<decltype(op)>, Layout<WGTile>, SGLayout>::TiledMMA;
 
   return MMA{};
 }
 
-// Unique kernel name tag
-template <typename EA, typename EB, typename ED, char layoutA, char layoutB>
+// 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 <typename EA, typename EB, typename ED, char layoutA, char layoutB, class WGTile, class SGLayout>
 class MoEGemmKernel;
 
 // MOE GEMM launcher template
-template <char layoutA, char layoutB, typename ElementA, typename ElementB, typename ElementS, typename ElementD>
+template <char layoutA, char layoutB, class WGTile, class SGLayout, typename ElementA, typename ElementB,
+          typename ElementS, typename ElementD>
 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<Shape<int, int, int>>{1, &dummy_problem_shape, nullptr};
 
-  using TileShape = Shape<_256, _128, _32>;
+  using TileShape = WGTile;
   using ClusterShape = Shape<_1, _1, _1>;
 
   auto scheduler_params = PersistentTileSchedulerXeMoE<ProblemShape>::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<ProblemShape>::Arguments{1, RasterOrderOptions::AlongN});
 
-  auto mma = choose_tiled_mma(activations, weights);
+  auto mma = choose_tiled_mma<WGTile, SGLayout, ElementA, ElementB>();
   auto MaxThreadsPerWorkgroup = size(mma);
   dim3 local_range{static_cast<unsigned int>(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<MoEGemmKernel<ElementA, ElementB, ElementD, layoutA, layoutB>>(
+  auto event = q->parallel_for<MoEGemmKernel<ElementA, ElementB, ElementD, layoutA, layoutB, WGTile, SGLayout>>(
       sycl::nd_range<3>(global, local), kernel_props, [=](auto) {
         MoE::MoEGEMM<XE_LOAD_2D<16, 32, 32, 16>, 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<char>(std::tolower(static_cast<unsigned char>(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 <typename Element>
+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<Shape<_8, _1, _1>, Stride<_1, _1, _0>>;
+  using N128 = Shape<_256, _128, _32>;
+  using SG128 = Layout<Shape<_8, _2, _1>, Stride<_2, _1, _0>>;
+  using N256 = Shape<_256, _256, _32>;
+  using SG256 = Layout<Shape<_8, _4, _1>, 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,15 +195,15 @@ 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<Element>(
           q, static_cast<const Element*>(activations), static_cast<const Element*>(weights),
           static_cast<const Element*>(scales), static_cast<Element*>(outputs), N, K, num_tokens_per_expert,
           num_experts);
       break;
     }
     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<Element>(
           q, static_cast<const Element*>(activations), static_cast<const Element*>(weights),
           static_cast<const Element*>(scales), static_cast<Element*>(outputs), N, K, num_tokens_per_expert,
           num_experts);

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:

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 的正确性测试