Opt the moe performance

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/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp

@@ -129,6 +129,7 @@ class MoEUpcastInt2SymToInt8KernelFast;
 // coalesced burst per row for FP16/BF16 writes, which matches the L1
 // cache-line size on the target XPUs.
 constexpr int WG_N = 32;
+constexpr int WG_K_INT4_FAST = 8;
 constexpr int PACK_K_FP = 4;
 constexpr int PACK_K_INT8 = 4;
 constexpr int PACK_K_INT4 = 2;
@@ -277,15 +278,17 @@ void launch_dequant_int4(sycl::queue* q, const uint8_t* weights_NKp, const Scala
   // relative to the byte-per-work-item path.
   if ((K % PACK_K_INT4_FAST) == 0 && (group_size % PACK_K_INT4_FAST) == 0) {
     const int k_words = K / PACK_K_INT4_FAST;  // == k_packed / 4
-    sycl::range<3> global_fast{static_cast<size_t>(E), static_cast<size_t>(k_words),
+    const int k_words_padded = ((k_words + WG_K_INT4_FAST - 1) / WG_K_INT4_FAST) * WG_K_INT4_FAST;
+    sycl::range<3> global_fast{static_cast<size_t>(E), static_cast<size_t>(k_words_padded),
                                static_cast<size_t>((N + WG_N - 1) / WG_N) * WG_N};
-    sycl::range<3> local_fast{1, 1, static_cast<size_t>(WG_N)};
+    sycl::range<3> local_fast{1, static_cast<size_t>(WG_K_INT4_FAST), static_cast<size_t>(WG_N)};
 
     q->parallel_for<MoEDequantKernelInt4Fast<ScalarT, Asym>>(
         sycl::nd_range<3>(global_fast, local_fast), [=](sycl::nd_item<3> it) {
           const int e = static_cast<int>(it.get_global_id(0));
           if (num_tokens_per_expert != nullptr && num_tokens_per_expert[e] == 0) return;
           const int kw = static_cast<int>(it.get_global_id(1));
+          if (kw >= k_words) return;
           const int n = static_cast<int>(it.get_global_id(2));
           if (n >= N) return;
           const int k_base = kw * PACK_K_INT4_FAST;

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

@@ -944,7 +944,7 @@ void moe_prefill_fp8_dpas_per_tensor_dispatch(
 
   if (A_avg_M <= 8) {
     ARK_DPAS_PT_LAUNCH(dpas_w8a16_policy_m_16);
-  } else if (A_avg_M <= 32) {
+  } else if (A_avg_M <= 512) {
     ARK_DPAS_PT_LAUNCH(dpas_w8a16_policy_m_32);
   } else {
     ARK_DPAS_PT_LAUNCH(dpas_w8a16_policy);
@@ -1002,7 +1002,7 @@ void moe_prefill_fp8_dpas_per_group_dispatch(
 
   if (A_avg_M <= 8) {
     ARK_DPAS_PG_LAUNCH(dpas_w8a16_policy_m_16);
-  } else if (A_avg_M <= 32) {
+  } else if (A_avg_M <= 512) {
     ARK_DPAS_PG_LAUNCH(dpas_w8a16_policy_m_32);
   } else {
     ARK_DPAS_PG_LAUNCH(dpas_w8a16_policy);

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

@@ -794,7 +794,7 @@ void moe_prefill_int_dpas_per_tensor_dispatch(
 
   if (A_avg_M <= 8) {
     ARK_DPAS_INT_PT_LAUNCH(dpas_w8a16_policy_m_16);
-  } else if (A_avg_M <= 32) {
+  } else if (A_avg_M <= 512) {
     ARK_DPAS_INT_PT_LAUNCH(dpas_w8a16_policy_m_32);
   } else {
     ARK_DPAS_INT_PT_LAUNCH(dpas_w8a16_policy);
@@ -857,7 +857,7 @@ void moe_prefill_int_dpas_per_group_dispatch(
 
   if (A_avg_M <= 8) {
     ARK_DPAS_INT_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16);
-  } else if (A_avg_M <= 32) {
+  } else if (A_avg_M <= 512) {
     ARK_DPAS_INT_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32);
   } else {
     ARK_DPAS_INT_PG_LAUNCH_SYM(dpas_w8a16_policy);