feat: W4A8 ARK XPU MoE kernel (int4 weight / int8 compute) with prefill + decode

auto_round_extension/ark/auto_round_kernel/ark.cpp

@@ -582,6 +582,42 @@ static void moe_gemm_prefill_int_dpas_wrapper(torch_ptr stream, torch_ptr activa
                                  (int*)num_tokens_per_expert, num_experts, total_tokens);
 }
 
+// W4A8 MoE: int4 weights re-scaled to int8 (AUTO_S8), int8 DPAS compute,
+// per-token dynamically quantized int8 activations. `weights_s8` is
+// [E, N, K] int8 and `wscales` is [E, N, K/rescale_block] FP32, both produced
+// by `moe_w4a8_prepack`. STATUS: NEEDS-HARDWARE-VALIDATION.
+static void moe_w4a8_prepack_wrapper(torch_ptr stream, torch_ptr weights_s4, torch_ptr scales,
+                                     torch_ptr weights_s8, torch_ptr wscales, int act_dtype, int num_experts,
+                                     int N, int K, int group_size, int rescale_group_size) {
+  ark::moe_w4a8_prepack((sycl::queue*)stream, (void*)weights_s4, (void*)scales, (void*)weights_s8,
+                        (void*)wscales, (BTLA_DTYPE)(act_dtype), num_experts, N, K, group_size,
+                        rescale_group_size);
+}
+
+// `qact`/`ascale` are optional pre-quantized activations (0 = quantize in the
+// call); `row_to_token`/`routing_weights`/`fused_out` are the optional fused
+// top-k reduction (0 = write the unreduced `[T, N]` output).
+static void moe_gemm_w4a8_wrapper(torch_ptr stream, torch_ptr activations, torch_ptr weights_s8,
+                                  torch_ptr wscales, torch_ptr outputs, int act_dtype, int N, int K,
+                                  int rescale_block_size, torch_ptr num_tokens_per_expert, int num_experts,
+                                  int total_tokens, int phase, torch_ptr qact, torch_ptr ascale,
+                                  torch_ptr row_to_token, torch_ptr routing_weights, torch_ptr fused_out,
+                                  int fused_batch) {
+  ark::moe_gemm_w4a8((sycl::queue*)stream, (void*)activations, (void*)weights_s8, (void*)wscales, (void*)outputs,
+                     (BTLA_DTYPE)(act_dtype), N, K, rescale_block_size, (int*)num_tokens_per_expert, num_experts,
+                     total_tokens, phase, (const void*)qact, (const float*)ascale, (const int*)row_to_token,
+                     (const float*)routing_weights, (float*)fused_out, fused_batch);
+}
+
+// Standalone per-token activation quantization -- the pass `moe_gemm_w4a8`
+// runs internally, exposed so it can be timed on its own and so a caller can
+// pre-quantize once and hand the result back as `qact`/`ascale`.
+static void moe_w4a8_quant_act_wrapper(torch_ptr stream, torch_ptr activations, torch_ptr qact, torch_ptr ascale,
+                                       int act_dtype, int total_tokens, int K) {
+  ark::moe_w4a8_quant_act((sycl::queue*)stream, (const void*)activations, (void*)qact, (void*)ascale,
+                          (BTLA_DTYPE)(act_dtype), total_tokens, K);
+}
+
 static void sage_dynamic_quant(torch_ptr stream, torch_ptr input, torch_ptr bias, torch_ptr output, torch_ptr scale_out,
                                int num_rows, int head_dim, int block_size) {
   auto* q = (sycl::queue*)stream;
@@ -1430,6 +1466,11 @@ PYBIND11_MODULE(PY_NAME, m) {
   m.def("moe_gemm_prefill", &ark::moe_gemm_prefill_wrapper);
   m.def("moe_gemm_prefill_fp8_dpas", &ark::moe_gemm_prefill_fp8_dpas_wrapper);
   m.def("moe_gemm_prefill_int_dpas", &ark::moe_gemm_prefill_int_dpas_wrapper);
+  m.def("moe_w4a8_prepack", &ark::moe_w4a8_prepack_wrapper);
+  m.def("moe_gemm_w4a8", &ark::moe_gemm_w4a8_wrapper);
+  m.def("moe_w4a8_quant_act", &ark::moe_w4a8_quant_act_wrapper);
+  m.def("moe_w4a8_rescale_block_size", &ark::moe_w4a8_rescale_block_size);
+  m.def("moe_w4a8_release_scratch", &ark::moe_w4a8_release_scratch);
   m.def("matmul_sycl_tla", &ark::matmul_sycl_tla);
 #endif  // ARK_SYCL_TLA
 #if !defined(ARK_XPU)

auto_round_extension/ark/auto_round_kernel/sycl_tla_generation.cmake

@@ -448,6 +448,68 @@ set(MOE_SOURCE_MODE 16)
 generate_sycl_tla_source(sycl_tla_moe.cpp.in sycl_tla_moe_prefill_s4.cpp)
 set(MOE_SOURCE_MODE 17)
 generate_sycl_tla_source(sycl_tla_moe.cpp.in sycl_tla_moe_prefill_native_fp8.cpp)
+# W4A8 (int4 weights re-scaled to int8, int8 activations) prepack + GEMM.
+#
+# This used to be a single `sycl_tla_moe_w4a8.cpp`, which instantiated all 52
+# W4A8 kernels -- 12 grouped DPAS GEMMs, 14 decode GEMVs, 22 activation-quant
+# and 4 prepack kernels -- and peaked at ~4.2 GB of compiler RSS. It is now
+# fanned out the same way the S4 prefill is (see MOE_SOURCE_MODE 8): one TU per
+# variant, with the dispatcher left holding no kernels at all.
+#
+#   mode 21  dispatcher            0 kernels, cutlass-free
+#   mode 22  prefill, dtype x tile 1 DPAS kernel each  (12 TUs)
+#   mode 23  decode, dtype         7 kernels each      (2 TUs, cutlass-free)
+#   mode 24  act quant, dtype      11 kernels each     (2 TUs, cutlass-free)
+#   mode 25  prepack, dtype        2 kernels each      (2 TUs, cutlass-free)
+#
+# Only the 12 prefill TUs include the CuTe stack; the other seven take
+# `sycl_tla_moe_w4a8_kernels.hpp` / `..._helpers.hpp`, which stop at plain SYCL.
+set(MOE_SOURCE_MODE 21)
+generate_sycl_tla_source(sycl_tla_moe.cpp.in sycl_tla_moe_w4a8.cpp)
+
+# Prefill tile ladder: the policy names must match `W4A8PrefillTile` /
+# `moe_w4a8_prefill_select_tile` in `sycl_tla_moe_w4a8_helpers.hpp`, and the
+# function names the `prefill_*` declarations there.
+set(_w4a8_tiles m8 m64 m128 m128n256 m256n128 large)
+foreach(_w4a8_dtype IN ITEMS f16 bf16)
+  if(_w4a8_dtype STREQUAL "f16")
+    set(MOE_SCALAR_TYPE sycl::half)
+  else()
+    set(MOE_SCALAR_TYPE sycl::ext::oneapi::bfloat16)
+  endif()
+
+  set(MOE_SOURCE_MODE 22)
+  foreach(_w4a8_tile IN LISTS _w4a8_tiles)
+    if(_w4a8_tile STREQUAL "m8")
+      set(MOE_POLICY_NAME w4a8_policy_m_8)
+    elseif(_w4a8_tile STREQUAL "m64")
+      set(MOE_POLICY_NAME w4a8_policy_m_64)
+    elseif(_w4a8_tile STREQUAL "m128")
+      set(MOE_POLICY_NAME w4a8_policy_m_128)
+    elseif(_w4a8_tile STREQUAL "m128n256")
+      set(MOE_POLICY_NAME w4a8_policy_m_128_n256)
+    elseif(_w4a8_tile STREQUAL "m256n128")
+      set(MOE_POLICY_NAME w4a8_policy_m_256_n128)
+    else()
+      set(MOE_POLICY_NAME w4a8_policy_large)
+    endif()
+    set(MOE_FUNCTION_NAME prefill_${_w4a8_dtype}_${_w4a8_tile})
+    generate_sycl_tla_source(sycl_tla_moe.cpp.in
+                             sycl_tla_moe_w4a8_prefill_${_w4a8_dtype}_${_w4a8_tile}.cpp)
+  endforeach()
+
+  set(MOE_SOURCE_MODE 23)
+  set(MOE_FUNCTION_NAME decode_${_w4a8_dtype})
+  generate_sycl_tla_source(sycl_tla_moe.cpp.in sycl_tla_moe_w4a8_decode_${_w4a8_dtype}.cpp)
+
+  set(MOE_SOURCE_MODE 24)
+  set(MOE_FUNCTION_NAME quant_${_w4a8_dtype})
+  generate_sycl_tla_source(sycl_tla_moe.cpp.in sycl_tla_moe_w4a8_quant_${_w4a8_dtype}.cpp)
+
+  set(MOE_SOURCE_MODE 25)
+  set(MOE_FUNCTION_NAME prepack_${_w4a8_dtype})
+  generate_sycl_tla_source(sycl_tla_moe.cpp.in sycl_tla_moe_w4a8_prepack_${_w4a8_dtype}.cpp)
+endforeach()
 
 foreach(_moe_dtype IN ITEMS f16 bf16)
   if(_moe_dtype STREQUAL "f16")

auto_round_extension/ark/auto_round_kernel/sycl_tla_moe.cpp.in

@@ -1040,6 +1040,116 @@ void dispatch_@MOE_DISPATCH_DTYPE@_@MOE_TILE_SUFFIX@(sycl::queue* q, const void*
 }  // namespace moe_gemm_detail
 }  // namespace ark
 
+#elif @MOE_SOURCE_MODE@ == 21
+
+// W4A8 dispatcher. Deliberately includes only the cutlass-free front end: the
+// tile choice and the public entry points are host code, and every kernel
+// lives in one of the per-variant TUs generated by modes 22-25. This TU
+// therefore instantiates no kernels at all.
+#include "sycl_tla_common.hpp"
+#include "sycl_tla_moe_w4a8_helpers.hpp"
+
+namespace ark {
+
+void moe_w4a8_prepack(sycl::queue* q, void* weights_s4, void* scales, void* weights_s8, void* wscales,
+                      BTLA_DTYPE act_dtype, int num_experts, int N, int K, int group_size,
+                      int rescale_group_size) {
+  moe_w4a8_detail::moe_w4a8_prepack(q, weights_s4, scales, weights_s8, wscales, act_dtype, num_experts, N, K,
+                                    group_size, rescale_group_size);
+}
+
+void moe_gemm_w4a8(sycl::queue* q, void* activations, void* weights_s8, void* wscales, void* outputs,
+                   BTLA_DTYPE act_dtype, int N, int K, int rescale_block_size, int* num_tokens_per_expert,
+                   int num_experts, int total_tokens, int phase, const void* qact_in, const float* ascale_in,
+                   const int* row_to_token, const float* routing_weights, float* fused_out, int fused_batch) {
+  moe_w4a8_detail::moe_gemm_w4a8(q, activations, weights_s8, wscales, outputs, act_dtype, N, K, rescale_block_size,
+                                 num_tokens_per_expert, num_experts, total_tokens, phase, qact_in, ascale_in,
+                                 row_to_token, routing_weights, fused_out, fused_batch);
+}
+
+int moe_w4a8_rescale_block_size(int K, int group_size, int rescale_group_size) {
+  return moe_w4a8_detail::moe_w4a8_rescale_block_size(K, group_size, rescale_group_size);
+}
+
+void moe_w4a8_quant_act(sycl::queue* q, const void* activations, void* qact, void* ascale, BTLA_DTYPE act_dtype,
+                        int total_tokens, int K) {
+  moe_w4a8_detail::moe_w4a8_quant_act(q, activations, qact, ascale, act_dtype, total_tokens, K);
+}
+
+void moe_w4a8_release_scratch() { moe_w4a8_detail::moe_w4a8_release_scratch(); }
+
+}  // namespace ark
+
+#elif @MOE_SOURCE_MODE@ == 22
+
+// W4A8 prefill, one DPAS policy per translation unit.
+//
+// The single TU this replaced instantiated 52 kernels -- 12 of them grouped
+// DPAS GEMMs -- and peaked at ~4.2 GB of compiler RSS. Fanning the tile ladder
+// out over (dtype x tile) leaves exactly one DPAS kernel here, mirroring how
+// `sycl_tla_moe_prefill_s4_*.cpp` splits the S4 prefill.
+#include "sycl_tla_moe_w4a8.hpp"
+
+namespace ark {
+namespace moe_w4a8_detail {
+
+void @MOE_FUNCTION_NAME@(const W4A8PrefillParams& params) {
+  moe_w4a8::moe_w4a8_prefill_launch<moe_w4a8::@MOE_POLICY_NAME@, @MOE_SCALAR_TYPE@>(params);
+}
+
+}  // namespace moe_w4a8_detail
+}  // namespace ark
+
+#elif @MOE_SOURCE_MODE@ == 23
+
+// W4A8 decode GEMV, one dtype per translation unit (7 kernels: the legacy
+// GEMV plus the six K-split (NCOLS, CH) specializations). Takes the
+// cutlass-free kernel header, so this TU never parses CuTe.
+#include "sycl_tla_moe_w4a8_kernels.hpp"
+
+namespace ark {
+namespace moe_w4a8_detail {
+
+void @MOE_FUNCTION_NAME@(const W4A8DecodeParams& params) {
+  moe_w4a8::moe_w4a8_decode_launch<@MOE_SCALAR_TYPE@>(params);
+}
+
+}  // namespace moe_w4a8_detail
+}  // namespace ark
+
+#elif @MOE_SOURCE_MODE@ == 24
+
+// W4A8 per-token activation quantization, one dtype per translation unit
+// (11 kernels: the message-width x unroll grid, the single-pass variants and
+// the scalar fallback). Cutlass-free.
+#include "sycl_tla_moe_w4a8_kernels.hpp"
+
+namespace ark {
+namespace moe_w4a8_detail {
+
+void @MOE_FUNCTION_NAME@(const W4A8QuantParams& params) {
+  moe_w4a8::moe_w4a8_quant_launch<@MOE_SCALAR_TYPE@>(params);
+}
+
+}  // namespace moe_w4a8_detail
+}  // namespace ark
+
+#elif @MOE_SOURCE_MODE@ == 25
+
+// W4A8 one-shot AUTO_S8 prepack, one dtype per translation unit (2 kernels:
+// the per-block scale reduction and the int4 -> int8 re-scale). Cutlass-free.
+#include "sycl_tla_moe_w4a8_kernels.hpp"
+
+namespace ark {
+namespace moe_w4a8_detail {
+
+void @MOE_FUNCTION_NAME@(const W4A8PrepackParams& params) {
+  moe_w4a8::moe_w4a8_prepack_launch<@MOE_SCALAR_TYPE@>(params);
+}
+
+}  // namespace moe_w4a8_detail
+}  // namespace ark
+
 #endif
 
 #endif  // ARK_XPU && ARK_SYCL_TLA

auto_round_extension/ark/auto_round_kernel/sycl_tla_moe_w4a8_scratch.cpp

@@ -0,0 +1,125 @@
+// Device scratch management for the W4A8 MoE path.
+//
+// The bookkeeping below is deliberately defined here rather than in a header so
+// that the module holds exactly one instance of it, and so that `utils.hpp`
+// (and the bestla JIT headers behind it) stays out of the cutlass-free W4A8
+// translation units. See `sycl_tla_moe_w4a8_scratch.hpp` for the rationale.
+//
+// Copyright (C) 2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include "sycl_tla_moe_w4a8_scratch.hpp"
+
+#if defined(ARK_XPU) && defined(ARK_SYCL_TLA)
+
+#include <map>
+#include <mutex>
+#include <stdexcept>
+#include <utility>
+#include <vector>
+
+#include "utils.hpp"
+
+namespace ark {
+namespace moe_w4a8 {
+
+namespace {
+
+struct W4A8ScratchState {
+  std::mutex mu;
+  // Device key (`DeviceMemoryPool::get_device_key`) -> a queue handle on that
+  // device, held *by value*: a `sycl::queue` is a reference-counted handle, so
+  // keeping a copy guarantees the queue outlives the memory allocated against
+  // it.
+  std::map<size_t, sycl::queue> queues;
+};
+
+// Intentionally leaked, see the header.
+W4A8ScratchState& w4a8_scratch_state() {
+  static W4A8ScratchState* s = new W4A8ScratchState();
+  return *s;
+}
+
+// Acquire a slab from the shared pool, synchronizing first when the request
+// grows it: `DeviceMemoryPool` frees the old pointer in place when it grows a
+// slot, and in-flight kernels may still be reading the old slab, so the wait
+// has to happen before the call rather than after.
+//
+// The caller must hold `W4A8ScratchState::mu`.
+void* acquire_w4a8_slab(sycl::queue* q, size_t bytes, size_t buf_loc) {
+  auto* pool = DeviceMemoryPool::Instance();
+  const size_t held = pool->get_scratch_size(buf_loc, q);
+  if (held != 0 && held < bytes) {
+    q->wait();
+  }
+  void* ptr = pool->get_scratch_mem(bytes, buf_loc, q);
+  if (ptr == nullptr) {
+    // The pool records the slot before checking the result, so a failed
+    // allocation leaves a {bytes, nullptr} entry behind that would satisfy
+    // every later request of this size or smaller without ever retrying.
+    // Drop it so the next call allocates again.
+    pool->detach_scratch_mem(buf_loc, q);
+    throw std::runtime_error("moe_gemm_w4a8: failed to allocate device scratch buffer");
+  }
+  auto& st = w4a8_scratch_state();
+  const size_t key = pool->get_device_key(q);
+  if (st.queues.find(key) == st.queues.end()) {
+    st.queues.emplace(key, *q);
+  }
+  return ptr;
+}
+
+}  // namespace
+
+uint8_t* acquire_qact_scratch(sycl::queue* q, size_t bytes) {
+  if (q == nullptr) {
+    throw std::invalid_argument("moe_gemm_w4a8: device scratch requires a non-null SYCL queue");
+  }
+  if (bytes == 0) return nullptr;
+  auto& st = w4a8_scratch_state();
+  std::lock_guard<std::mutex> lock(st.mu);
+  return static_cast<uint8_t*>(acquire_w4a8_slab(q, bytes, kW4A8QactScratchLoc));
+}
+
+int* acquire_expert_map_scratch(sycl::queue* q, size_t bytes) {
+  if (q == nullptr) {
+    throw std::invalid_argument("moe_gemm_w4a8: device scratch requires a non-null SYCL queue");
+  }
+  if (bytes == 0) return nullptr;
+  auto& st = w4a8_scratch_state();
+  std::lock_guard<std::mutex> lock(st.mu);
+  return static_cast<int*>(acquire_w4a8_slab(q, bytes, kW4A8ExpertMapScratchLoc));
+}
+
+void moe_w4a8_release_scratch() {
+  auto& st = w4a8_scratch_state();
+
+  // Detach everything under the lock, then drop the lock before the device sync
+  // and the frees: `wait()` blocks for an unbounded time and must not be held
+  // across. Because the slabs are already out of the pool's tables, an acquire
+  // that races in behind us allocates fresh ones instead of handing back a
+  // pointer we are about to free.
+  std::vector<std::pair<sycl::queue, void*>> pending;
+  {
+    std::lock_guard<std::mutex> lock(st.mu);
+    auto* pool = DeviceMemoryPool::Instance();
+    for (auto& kv : st.queues) {
+      sycl::queue q = kv.second;
+      for (size_t loc : {kW4A8QactScratchLoc, kW4A8ExpertMapScratchLoc}) {
+        void* ptr = pool->detach_scratch_mem(loc, &q);
+        if (ptr != nullptr) pending.emplace_back(q, ptr);
+      }
+    }
+    st.queues.clear();
+  }
+
+  for (auto& item : pending) {
+    item.first.wait();
+    sycl::free(item.second, item.first);
+  }
+}
+
+}  // namespace moe_w4a8
+}  // namespace ark
+
+#endif  // ARK_XPU && ARK_SYCL_TLA