From 8be551d35f265250dcaf14f7069d9411d0965dcc Mon Sep 17 00:00:00 2001 From: "Dong, Bo1" Date: Mon, 27 Jul 2026 10:27:55 +0800 Subject: [PATCH 01/47] update the code Signed-off-by: Dong, Bo1 --- .../wrapper/include/sycl_tla_moe_mixed.hpp | 7 +- .../include/sycl_tla_moe_prefill_fp8_dpas.hpp | 4 +- .../include/sycl_tla_moe_prefill_int_dpas.hpp | 4 +- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 269 ++++++++++++++---- 4 files changed, 217 insertions(+), 67 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp index e6a3ececcd..2312702563 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp +++ b/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(E), static_cast(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(E), static_cast(k_words_padded), static_cast((N + WG_N - 1) / WG_N) * WG_N}; - sycl::range<3> local_fast{1, 1, static_cast(WG_N)}; + sycl::range<3> local_fast{1, static_cast(WG_K_INT4_FAST), static_cast(WG_N)}; q->parallel_for>( sycl::nd_range<3>(global_fast, local_fast), [=](sycl::nd_item<3> it) { const int e = static_cast(it.get_global_id(0)); if (num_tokens_per_expert != nullptr && num_tokens_per_expert[e] == 0) return; const int kw = static_cast(it.get_global_id(1)); + if (kw >= k_words) return; const int n = static_cast(it.get_global_id(2)); if (n >= N) return; const int k_base = kw * PACK_K_INT4_FAST; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index 1638bf7af0..35db93cac2 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/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); diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_int_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_int_dpas.hpp index 2e05dfa442..8f7260bd90 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_int_dpas.hpp +++ b/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); diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index fad8d87eb7..f002a092d1 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -91,6 +91,7 @@ #pragma once #include +#include #include #include #include @@ -134,6 +135,70 @@ using ::ark::moe_dpas_fp8::cute_scalar; using ::ark::moe_dpas_fp8::cute_scalar_t; using ::ark::moe_dpas_fp8::make_moe_tensor; +inline int _moe_prefill_s4_env_int(const char* name, int default_v, + int min_v, int max_v) { + const char* env = std::getenv(name); + if (env == nullptr) return default_v; + char* end = nullptr; + long v = std::strtol(env, &end, 10); + if (end == env || *end != '\0') return default_v; + if (v < static_cast(min_v)) return min_v; + if (v > static_cast(max_v)) return max_v; + return static_cast(v); +} + +inline int moe_prefill_dpas_s4_prefetch_dist() { + // Runtime-tunable A/B prefetch depth for S4 per-group mainloop. + // Default keeps current behavior. + return _moe_prefill_s4_env_int("ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST", + 3, 1, 16); +} + +inline int moe_prefill_dpas_s4_prefetch_dist_scale() { + // Runtime-tunable scale prefetch lookahead (in groups). + // Default keeps current behavior. + return _moe_prefill_s4_env_int( + "ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST_SCALE", 3, 1, 16); +} + +enum class MoePrefillS4PolicyMode { + Auto, + W8A16, + W8A16M16, + W8A16M32, + W16A16, +}; + +inline bool moe_prefill_dpas_s4_force_direct() { + const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4"); + if (env != nullptr) { + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "1" || s == "true" || s == "on" || s == "yes" || + s == "direct" || s == "force" || s == "packed") return true; + } + + const char* policy = std::getenv("ARK_MOE_PREFILL_DPAS_S4_POLICY"); + if (policy == nullptr) return false; + std::string p(policy); + for (auto& c : p) c = static_cast(std::tolower(static_cast(c))); + return p == "w16a16" || p == "w16" || p == "w8a16" || p == "w8" || + p == "w8a16_m16" || p == "w8_m16" || p == "m16" || + p == "w8a16_m32" || p == "w8_m32" || p == "m32"; +} + +inline MoePrefillS4PolicyMode moe_prefill_dpas_s4_policy_mode() { + const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4_POLICY"); + if (env == nullptr) return MoePrefillS4PolicyMode::Auto; + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "w16a16" || s == "w16") return MoePrefillS4PolicyMode::W16A16; + if (s == "w8a16_m16" || s == "w8_m16" || s == "m16") return MoePrefillS4PolicyMode::W8A16M16; + if (s == "w8a16_m32" || s == "w8_m32" || s == "m32") return MoePrefillS4PolicyMode::W8A16M32; + if (s == "w8a16" || s == "w8") return MoePrefillS4PolicyMode::W8A16; + return MoePrefillS4PolicyMode::Auto; +} + // --------------------------------------------------------------------------- // Variant B -- per-K-group S4 (sym) mainloop. // @@ -172,8 +237,9 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( const ElementBI* Bias, DTensor& C, // (M,N) -- ElementA Coord blk_coord, - TiledMMA const& mma) { - using TA = typename ATensor::element_type; + TiledMMA const& mma, + int prefetch_dist, + int prefetch_dist_scale) { using TB = typename BTensor::element_type; static_assert(std::is_same_v, "xe_gemm_s4_pergroup: ElementB must be cutlass::int4b_t (sym only)"); @@ -204,8 +270,13 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( auto thr_copy_b = copy_b.get_slice(local_id); auto thr_copy_c = copy_c.get_slice(local_id); - auto tCrA = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); - auto tCrB = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + // Ping-pong SG fragments for software pipelining: + // while tile k computes on one fragment pair, tile k+1 is unpacked + // into the other pair. + auto tCrA0 = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB0 = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + auto tCrA1 = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB1 = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_, _, 0)); auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_, _, 0)); @@ -226,11 +297,8 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( auto pAgA = thr_prefetch_A.partition_S(gA); auto pBgB = thr_prefetch_B.partition_S(gB); - // Prefetch distance mirrors `xe_gemm_int_pergroup<>` for now. - // On-hardware perf tuning may want to grow `prefetch_dist` on the - // packed path since the B stream is half the bandwidth. - const int prefetch_dist = 3; - const int prefetch_dist_scale = 3; + // Prefetch distances are host-selected and threaded into the device + // path so no env/stdlib calls are needed inside SYCL kernels. constexpr auto barrier_scope = ScopeWorkgroup; int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); int k_tile_prefetch = 0; @@ -249,11 +317,8 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( auto n_tile_start = wg_n * tile_n; auto sg_local_n_coord = cutlass::get_sub_group_id() % ATOM_N; - auto sg_local_m_coord = cutlass::get_sub_group_id() / ATOM_N; int sg_local_id = cutlass::get_sub_group_local_id(); int n_sg_start = sg_local_n_coord * SG_N; - int m_sg_start = sg_local_m_coord * SG_M; - int m_tile_start = wg_m * tile_m; int group_num = get<1>(A.shape()) / group_size; // Group-local accumulator: same fragment shape as `tCrC`, cleared at @@ -267,6 +332,16 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( // Per-SG per-N scale cache. Same layout / semantics as the INT8 // per-group path. float sg_scale[sg_n_strides]; + int sg_scale_base[sg_n_strides]; + + // Precompute per-lane base offsets into Scales so the hot loop only + // adds group_idx. + CUTLASS_PRAGMA_UNROLL + for (int sn = 0; sn < sg_n_strides; ++sn) { + int sg_local_n = sn * sg_local_range + sg_local_id; + sg_scale_base[sn] = + (n_tile_start + n_sg_start + sg_local_n) * group_num; + } CUTE_UNROLL for (; k_tile_prefetch < prefetch_dist; k_tile_prefetch++) { @@ -290,27 +365,48 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( } } + if (k_tile_count == 0) { + return; + } + + // Prime the pipeline with tile 0. + barrier_arrive(barrier_scope); + copy(copy_a, tAgA(_, _, _, 0), tArA); + copy(copy_b, tBgB(_, _, _, 0), tBrB); + if (k_tile_prefetch < k_tile_count) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + reorder(tArA, tCrA0); + reorder(tBrB, tCrB0); + barrier_wait(barrier_scope); + + int curr_buf = 0; + static constexpr int tiles_per_group = group_size / tile_k; + int group_tile_remaining = 0; + int group_idx = 0; for (int k_tile = 0; k_tile < k_tile_count; k_tile++, k_tile_prefetch++) { - barrier_arrive(barrier_scope); - - copy(copy_a, tAgA(_, _, _, k_tile), tArA); - copy(copy_b, tBgB(_, _, _, k_tile), tBrB); - - // Group-boundary scale reload. Same math as the INT8 per-group - // path: `tile_k` is expressed in element units (nibbles), not - // bytes, so `k_tile * tile_k` is the reduction position in - // *element* space and the modulo test against `group_size` matches - // the scale-tensor layout `[E, N, K/group_size]` unchanged. - if (k_tile * tile_k % group_size == 0) { - int group_idx = (k_tile * tile_k) / group_size; + const bool has_next = (k_tile + 1 < k_tile_count); + const int next_tile = k_tile + 1; + + if (has_next) { + barrier_arrive(barrier_scope); + copy(copy_a, tAgA(_, _, _, next_tile), tArA); + copy(copy_b, tBgB(_, _, _, next_tile), tBrB); + } + + // Group-boundary scale reload via a tile counter instead of per-tile + // modulo/div in the hot loop. For all supported S4 shapes in this + // file, `group_size` is a multiple of `tile_k`. + if (group_tile_remaining == 0) { + group_idx = k_tile / tiles_per_group; CUTLASS_PRAGMA_UNROLL for (int sn = 0; sn < sg_n_strides; ++sn) { - int sg_local_n = sn * sg_local_range + sg_local_id; - sg_scale[sn] = static_cast( - Scales[(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx]); + sg_scale[sn] = + static_cast(Scales[sg_scale_base[sn] + group_idx]); } - if ((group_idx + prefetch_dist_scale) * group_size < shape<1>(A)) { + if (group_idx + prefetch_dist_scale < group_num) { auto next_scales_tensor = make_tensor( make_gmem_ptr(reinterpret_cast( Scales + (n_tile_start + n_sg_start) * group_num + @@ -324,32 +420,41 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( make_identity_tensor(make_shape(Int{}, Int<1>{}))); prefetch(prefetch_scales, pSgS(_, 0, 0)); } + group_tile_remaining = tiles_per_group; } - if (k_tile_prefetch < k_tile_count) { + if (has_next && k_tile_prefetch < k_tile_count) { prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); } - // `reorder` performs the in-register `int4b_t -> ElementA` unpack - // + sign-extend + cast via `cutlass::NumericArrayConverter< - // ElementA, cutlass::int4b_t, N>`. Once `tCrB` carries bf16/fp16 - // values it is compatible with the same DPAS atom used by the FP8 - // / INT8 per-group paths. See the header preamble open-question - // (1) -- if the pinned cutlass-sycl is missing this converter - // specialisation this line is where the build fails. - reorder(tArA, tCrA); - reorder(tBrB, tCrB); + if (has_next) { + // Pipeline stage: unpack tile k+1 while tile k compute is pending. + // `reorder` performs in-register `int4b_t -> ElementA` conversion via + // `NumericArrayConverter`. + if (curr_buf == 0) { + reorder(tArA, tCrA1); + reorder(tBrB, tCrB1); + } else { + reorder(tArA, tCrA0); + reorder(tBrB, tCrB0); + } + } // HOT MAINLOOP -- MMA accumulates into `tCrC_group`. Per-N scale // is applied ONCE at the end of the group in the fold block below. - cute::gemm(mma, tCrA, tCrB, tCrC_group); + if (curr_buf == 0) { + cute::gemm(mma, tCrA0, tCrB0, tCrC_group); + } else { + cute::gemm(mma, tCrA1, tCrB1, tCrC_group); + } - // Group-boundary fold. Fires when either (a) the NEXT k_tile would - // start a new scale group, or (b) we've reached the last k_tile of - // the K reduction (tail-group protection). - const bool is_group_end = (((k_tile + 1) * tile_k) % group_size == 0) || - (k_tile + 1 == k_tile_count); + group_tile_remaining -= 1; + // Group-boundary fold. Fires when either (a) we've consumed all + // tiles in the current group, or (b) this is the last reduction tile + // (tail-group protection). + const bool is_group_end = (group_tile_remaining == 0) || + (k_tile + 1 == k_tile_count); if (is_group_end) { CUTLASS_PRAGMA_UNROLL for (int sn = 0; sn < sg_n_strides; ++sn) { @@ -363,7 +468,10 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( } } - barrier_wait(barrier_scope); + if (has_next) { + barrier_wait(barrier_scope); + curr_buf ^= 1; + } } if (Bias != nullptr) { @@ -408,7 +516,10 @@ CUTE_DEVICE void MoEGEMM_s4(const ElementA* Activations, const int* rows_per_expert, const int32_t num_experts, const int32_t group_size, const int32_t gemm_n, - const int32_t gemm_k, int32_t* atomic_buffer, + const int32_t gemm_k, + int prefetch_dist, + int prefetch_dist_scale, + int32_t* atomic_buffer, const sycl::local_accessor& slm_mem_const) { constexpr char actual_layout_of_B = LayoutKindB ^ ('R' ^ 'C'); @@ -495,7 +606,7 @@ CUTE_DEVICE void MoEGEMM_s4(const ElementA* Activations, #define ARK_MOE_DPAS_S4_GROUP_CALLER(GS) \ xe_gemm_s4_pergroup( \ A_tensor, B_tensor, ptr_Scales_curr_batch, ptr_Bias_curr_batch, \ - D_tensor, tile_coord, mma); + D_tensor, tile_coord, mma, prefetch_dist, prefetch_dist_scale); if (group_size == 32) { ARK_MOE_DPAS_S4_GROUP_CALLER(32) } else if (group_size == 64) { @@ -534,7 +645,10 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, const ElementBI* bias, ElementD* outputs, const int gemm_n, const int gemm_k, const int* rows_per_expert, const int num_experts, - const int group_size, int32_t* atomic_buffer) { + const int group_size, + int prefetch_dist, + int prefetch_dist_scale, + int32_t* atomic_buffer) { using ElementA_non_CV = cutlass::platform::remove_cv_t; // DPAS atom keeps its bf16/fp16 x bf16/fp16 -> fp32 shape; the S4 B // tensor is upcast to ElementA in `reorder(tBrB, tCrB)` in the @@ -581,7 +695,7 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, layoutB, 'R'>( activations, weights, scales, bias, outputs, mma, rows_per_expert, num_experts, group_size, gemm_n, gemm_k, - atomic_buffer, local_mem); + prefetch_dist, prefetch_dist_scale, atomic_buffer, local_mem); }); }); @@ -636,6 +750,9 @@ void moe_prefill_s4_dpas_per_group_dispatch( reinterpret_cast(weights_NKp); int A_avg_M = total_tokens / E; + const int prefetch_dist = moe_prefill_dpas_s4_prefetch_dist(); + const int prefetch_dist_scale = moe_prefill_dpas_s4_prefetch_dist_scale(); + const MoePrefillS4PolicyMode policy_mode = moe_prefill_dpas_s4_policy_mode(); int32_t* atomic_buffer = sycl::malloc_device(1, *q); if (atomic_buffer == nullptr) { @@ -647,14 +764,35 @@ void moe_prefill_s4_dpas_per_group_dispatch( MoEGEMMLauncher_s4<'R', 'C', policy>( \ *q, activations_ca, weights_i4, scales_ca, \ static_cast(nullptr), outputs_ca, N, K, \ - num_tokens_per_expert, E, group_size, atomic_buffer); + num_tokens_per_expert, E, group_size, prefetch_dist, \ + prefetch_dist_scale, atomic_buffer); - if (A_avg_M <= 8) { + if (policy_mode == MoePrefillS4PolicyMode::W16A16) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w16a16_policy); + } else if (policy_mode == MoePrefillS4PolicyMode::W8A16M16) { ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (A_avg_M <= 32) { + } else if (policy_mode == MoePrefillS4PolicyMode::W8A16M32) { ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + } else if (policy_mode == MoePrefillS4PolicyMode::W8A16) { + if (A_avg_M <= 8) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); + } else if (A_avg_M <= 32) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + } else { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + } } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + // Auto policy for the direct path. On BMG MiniMax prefill shapes the + // small-M w8a16_m32 tile is much faster than larger tiles even up to + // 8K prompts (average M ~= 341), because per-expert token counts are + // highly skewed and many experts underfill larger M tiles. + if (A_avg_M <= 8) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); + } else if (A_avg_M <= 512) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + } else { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + } } #undef ARK_DPAS_S4_PG_LAUNCH_SYM @@ -662,24 +800,33 @@ void moe_prefill_s4_dpas_per_group_dispatch( } // --------------------------------------------------------------------------- -// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default ON, semantics -// identical to `moe_prefill_dpas_int_enabled` / `moe_prefill_dpas_fp8 -// _enabled`). Decoupled from `ARK_MOE_PREFILL_DPAS_INT8` so this new -// single-pass path can be disabled in isolation if it regresses -- -// switching S4 off falls back to the S4->S8 upcast + INT8 DPAS path -// which is itself gated by `ARK_MOE_PREFILL_DPAS_INT8`. +// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4`. The packed-nibble direct +// DPAS path is opt-in and uses the auto policy above by default. Current BMG +// MiniMax measurements prefer the w8a16_m32 tile family for real prefill +// distributions; callers can still force another policy for experiments via +// `ARK_MOE_PREFILL_DPAS_S4_POLICY=w16a16|w8a16|w8a16_m16|w8a16_m32`. // // Truthy values (case-insensitive): "1", "true", "on", "yes". // Explicit "0" / "false" / "off" / "no" disable. Re-read on every // call so benchmarks / tests can toggle the path in-process. +// +// Additional runtime tuning knobs for the S4 per-group mainloop: +// - ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST (default 3, clamp [1, 16]) +// - ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST_SCALE (default 3, clamp [1, 16]) +// - ARK_MOE_PREFILL_DPAS_S4_POLICY (auto|w8a16|w8a16_m16| +// w8a16_m32|w16a16, +// default auto) +// +// These tune A/B tile prefetch depth and per-group scale prefetch lookahead +// respectively; `*_POLICY` selects the GEMM tile family. // --------------------------------------------------------------------------- inline bool moe_prefill_dpas_s4_enabled() { const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4"); - if (env == nullptr) return true; // default ON + if (env == nullptr) return false; std::string s(env); for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); if (s == "0" || s == "false" || s == "off" || s == "no") return false; - return true; + return moe_prefill_dpas_s4_force_direct(); } // --------------------------------------------------------------------------- From ff76f390ecfec7dab36313bb7155d13aa2a1beff Mon Sep 17 00:00:00 2001 From: "Dong, Bo1" Date: Thu, 30 Jul 2026 16:10:30 +0800 Subject: [PATCH 02/47] opt 16 bit group gemm Signed-off-by: Dong, Bo1 --- .../wrapper/include/sycl_tla_moe.hpp | 94 ++++++++++++++++--- 1 file changed, 80 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); From 68bb067d54ab539b90b1f4b4561dbf8bd6799ba6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 06:04:58 +0000 Subject: [PATCH 03/47] perf: dedicated w4a16 tiles for int4 moe prefill + widen decode gemv --- .../wrapper/include/sycl_tla_moe_decode.hpp | 85 ++++++++++++------- .../include/sycl_tla_moe_prefill_fp8_dpas.hpp | 29 +++++++ .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 27 ++++-- .../ark/test/README_MOE_PREFILL_PERF.md | 20 +++++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 18 ++++ 5 files changed, 143 insertions(+), 36 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 1b8a4feb0b..11535a3257 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -204,6 +204,41 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight // of byte i, the value at k = 2*i+1 is the HIGH nibble. This matches the // existing CPU/XPU `packq` layout for S4_CLIP weights. // ---------------------------------------------------------------------------- + +// Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK/2 +// packed weight bytes + a vec activation block). Templated on +// CHUNK so the caller can run a wide (32) stage first and a narrower (16) +// stage for the remainder, which keeps the fast path active for group sizes +// that are a multiple of 32 (32/64/128/256 -- the shipped quant configs) +// without regressing group_size == 16 (which drops straight to the 16-wide +// stage). The math is identical to the scalar path. +template +static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, + float& acc) { + using ActVec = sycl::vec; + using PackVec = sycl::vec; + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + const ActVec av = *reinterpret_cast(act_ptr); + const PackVec pv = *reinterpret_cast(w_ptr); +#pragma unroll + for (int b = 0; b < CHUNK / 2; ++b) { + int q0, q1; + decode_int4_pair(pv[b], q0, q1); + float w0, w1; + if constexpr (Asym) { + w0 = (static_cast(q0) - zero) * scale; + w1 = (static_cast(q1) - zero) * scale; + } else { + w0 = static_cast(q0) * scale; + w1 = static_cast(q1) * scale; + } + const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); + acc += static_cast(a0) * w0; + acc += static_cast(a1) * w1; + } +} + template void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, @@ -253,39 +288,27 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig zero = static_cast(z_row[g]); } const int k_base = g * group_size; - // Vectorized path: process 16 K-elements at a time, which is - // 8 packed weight bytes and a vec activation block. - // group_size is a multiple of 16 in every supported config - // (group_size >= 32, even); a scalar tail loop covers leftovers. - constexpr int CHUNK = 16; - using ActVec = sycl::vec; - using PackVec = sycl::vec; - static_assert(sizeof(ScalarT) == sizeof(uint16_t), - "ScalarT must be a 16-bit floating type"); - const int chunk_end = (group_size / CHUNK) * CHUNK; + // Vectorized ladder: process 32 K-elements at a time (16 packed + // weight bytes + vec activation block), then a 16-wide + // stage for the remainder, then a scalar tail. Widening the first + // stage to 32 amortizes the per-group scale load and loop overhead + // across twice as many multiply-adds for the shipped group sizes + // (32/64/128/256), while the 16-wide stage keeps group_size == 16 + // on the fast path. int kk = 0; - for (; kk < chunk_end; kk += CHUNK) { - const ActVec av = *reinterpret_cast(act_row + k_base + kk); - const PackVec pv = *reinterpret_cast(w_row + (k_base + kk) / 2); -#pragma unroll - for (int b = 0; b < CHUNK / 2; ++b) { - int q0, q1; - decode_int4_pair(pv[b], q0, q1); - float w0, w1; - if constexpr (Asym) { - w0 = (static_cast(q0) - zero) * scale; - w1 = (static_cast(q1) - zero) * scale; - } else { - w0 = static_cast(q0) * scale; - w1 = static_cast(q1) * scale; - } - const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); - acc += static_cast(a0) * w0; - acc += static_cast(a1) * w1; - } + constexpr int CHUNK32 = 32; + const int end32 = (group_size / CHUNK32) * CHUNK32; + for (; kk < end32; kk += CHUNK32) { + int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, + zero, acc); + } + constexpr int CHUNK16 = 16; + const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; + for (; kk < end16; kk += CHUNK16) { + int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, + zero, acc); } - // Scalar tail for group_size not divisible by CHUNK. + // Scalar tail for group_size not divisible by 16. for (; kk < group_size; kk += 2) { const uint8_t packed = w_row[(k_base + kk) / 2]; int q0, q1; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index 1638bf7af0..b9898f81f0 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -193,6 +193,35 @@ class dpas_w8a16_policy_m_32 : public dpas_policy_base { using SGLayout = Layout, Stride<_4, _1, _0>>; }; +// --------------------------------------------------------------------------- +// 4-bit (S4 / w4a16) tile policies. +// +// The S4 mixed-input prefill mainloop (`sycl_tla_moe_prefill_s4_dpas.hpp`) +// reads a *halved* B-side byte stream (two nibbles per byte). With the same +// M/N footprint the packed weight fits in half the L2/GRF traffic of the INT8 +// path, so a larger N tile pays off: the default large-M policy uses a +// 128x256x32 WG tile (vs. the INT8 128x128x16), matching the reference +// `w4a16_policy` in vllm-xpu-kernels `csrc/xpu/grouped_gemm/xe_2/ +// gemm_xe2_policy.hpp`. The small-M buckets (m_8/m_16/m_32) reuse the same +// 64-wide N tiles as the INT8 path -- the reference uses identical shapes +// there. Only the default (large-M) tile and the new m_8 bucket differ, so we +// define those two here and alias m_16/m_32 to the shared shapes in the S4 +// header. +// --------------------------------------------------------------------------- +class dpas_w4a16_policy : public dpas_policy_base { + public: + using WGTile = Shape<_128, _256, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; + + using GmemTiledCopyD = XE_STORE_2D<16, 8, 32>; +}; + +class dpas_w4a16_policy_m_8 : public dpas_policy_base { + public: + using WGTile = Shape<_8, _64, _32>; + using SGLayout = Layout, Stride<_4, _1, _0>>; +}; + // --------------------------------------------------------------------------- // `apply_scale` -- inline-asm per-lane multiply of a bf16/fp16 fragment by // an FP32 scalar. Copied verbatim from vllm-xpu-kernels `gemm_xe2.hpp`. diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index fad8d87eb7..693a8675e5 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -129,6 +129,15 @@ using ::ark::moe_dpas_fp8::dpas_w16a16_policy; using ::ark::moe_dpas_fp8::dpas_w8a16_policy; using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; +// Dedicated 4-bit tile policies. The default (large-M) tile is 128x256x32 +// (halved packed-B stream lets a wider N tile pay off) and the m_8 bucket is +// new; both mirror the reference `w4a16_policy*` in vllm-xpu-kernels. The +// m_16 / m_32 buckets share the INT8 64-wide N tiles, so they are aliased to +// the existing `dpas_w8a16_policy_m_16 / _m_32` shapes rather than duplicated. +using ::ark::moe_dpas_fp8::dpas_w4a16_policy; +using ::ark::moe_dpas_fp8::dpas_w4a16_policy_m_8; +using dpas_w4a16_policy_m_16 = ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; +using dpas_w4a16_policy_m_32 = ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; using ::ark::moe_dpas_fp8::ScaleMode; using ::ark::moe_dpas_fp8::cute_scalar; using ::ark::moe_dpas_fp8::cute_scalar_t; @@ -649,12 +658,20 @@ void moe_prefill_s4_dpas_per_group_dispatch( static_cast(nullptr), outputs_ca, N, K, \ num_tokens_per_expert, E, group_size, atomic_buffer); - if (A_avg_M <= 8) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (A_avg_M <= 32) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + // Four-tier bucketing on the average tokens-per-expert, matching the + // reference `w4a16` dispatch in vllm-xpu-kernels + // (`grouped_gemm_xe2_interface.hpp`): tiny M uses an 8-row tile, and the + // 32-row tile now covers M up to 128 (instead of jumping to the wide + // large-M tile at 33) so mid-size chunked-prefill batches avoid the + // padding waste of the 128-row tile. + if (A_avg_M <= 4) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_8); + } else if (A_avg_M <= 8) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_16); + } else if (A_avg_M <= 128) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_32); } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy); } #undef ARK_DPAS_S4_PG_LAUNCH_SYM 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..c1ed20ef5b 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -291,6 +291,26 @@ these fail: - `group_size ∈ {32, 64, 128, 256}` - `asym == false` (asym S4 is out of scope for both DPAS paths) +**S4 DPAS tile policies** — the single-pass mainloop (precedence 1) +now selects a dedicated 4-bit tile policy by the average tokens-per- +expert (`A_avg_M = total_tokens / E`), mirroring the reference +`w4a16` dispatch in `vllm-project/vllm-xpu-kernels` +(`grouped_gemm_xe2_interface.hpp`). Because the packed-nibble B stream +is half the byte volume of the INT8 path, the large-M tile is widened +to `128×256×32` (vs. the INT8 `128×128×16`) so the DPAS accumulators +and the halved B-side bandwidth are better utilised: + +| `A_avg_M` bucket | WG tile (M×N×K) | Policy (`sycl_tla_moe_prefill_fp8_dpas.hpp`) | +| ---------------- | --------------- | -------------------------------------------- | +| `≤ 4` | `8×64×32` | `dpas_w4a16_policy_m_8` | +| `≤ 8` | `16×64×32` | `dpas_w4a16_policy_m_16` (= `w8a16_m_16`) | +| `≤ 128` | `32×64×32` | `dpas_w4a16_policy_m_32` (= `w8a16_m_32`) | +| `> 128` | `128×256×32` | `dpas_w4a16_policy` | + +The mid-size `32×64` tile now covers `A_avg_M` up to 128 (previously it +jumped to the wide tile at 33), which avoids padding waste on the +common chunked-prefill batch sizes. + Accuracy parity is covered by `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group`, which forces `ARK_MOE_PREFILL_DPAS_S4=1` + 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..713474651d 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 @@ -222,6 +222,24 @@ S4-sym 有两条独立的 DPAS 路径;asym S4 始终回退到 dequant 路径。 - `group_size ∈ {32, 64, 128, 256}` - `asym == false`(asym S4 不在两条 DPAS 路径的支持范围内) +**S4 DPAS tile 策略** — 单遍 mainloop(优先级 1)现在按每专家平均 +token 数(`A_avg_M = total_tokens / E`)选择专用的 4-bit tile 策略, +与参考实现 `vllm-project/vllm-xpu-kernels` +(`grouped_gemm_xe2_interface.hpp`)的 `w4a16` 分派一致。由于 packed- +nibble 的 B 流字节量是 INT8 路径的一半,大 M tile 加宽到 `128×256×32` +(相比 INT8 的 `128×128×16`),以更充分利用 DPAS 累加器与减半的 B 侧 +带宽: + +| `A_avg_M` 分档 | WG tile (M×N×K) | 策略(`sycl_tla_moe_prefill_fp8_dpas.hpp`) | +| -------------- | --------------- | ------------------------------------------ | +| `≤ 4` | `8×64×32` | `dpas_w4a16_policy_m_8` | +| `≤ 8` | `16×64×32` | `dpas_w4a16_policy_m_16`(= `w8a16_m_16`) | +| `≤ 128` | `32×64×32` | `dpas_w4a16_policy_m_32`(= `w8a16_m_32`) | +| `> 128` | `128×256×32` | `dpas_w4a16_policy` | + +中等大小的 `32×64` tile 现在覆盖 `A_avg_M` 至 128(此前在 33 就跳到大 +tile),避免了常见 chunked-prefill batch 大小下的 padding 浪费。 + 精度对齐由 `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group` 覆盖,该用例强制 `ARK_MOE_PREFILL_DPAS_S4=1` + From 6cda6c8cd3e183ab84f08e7bdf769e79094fbe32 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 06:25:32 +0000 Subject: [PATCH 04/47] fix: avoid illegal sycl::vec in int4 MoE decode kernel --- .../wrapper/include/sycl_tla_moe_decode.hpp | 45 +++++++++++-------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 11535a3257..de480993b0 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -215,27 +215,36 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight template static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, float& acc) { - using ActVec = sycl::vec; - using PackVec = sycl::vec; static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); - const ActVec av = *reinterpret_cast(act_ptr); - const PackVec pv = *reinterpret_cast(w_ptr); + static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); + // sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so a single + // vec load is illegal. Process the chunk in 16-wide sub-blocks + // (16 activations + 8 packed weight bytes each), which keeps CHUNK == 32 + // valid while reusing the same code path for CHUNK == 16. + constexpr int SUB = 16; + using ActVec = sycl::vec; + using PackVec = sycl::vec; #pragma unroll - for (int b = 0; b < CHUNK / 2; ++b) { - int q0, q1; - decode_int4_pair(pv[b], q0, q1); - float w0, w1; - if constexpr (Asym) { - w0 = (static_cast(q0) - zero) * scale; - w1 = (static_cast(q1) - zero) * scale; - } else { - w0 = static_cast(q0) * scale; - w1 = static_cast(q1) * scale; + for (int s = 0; s < CHUNK / SUB; ++s) { + const ActVec av = *reinterpret_cast(act_ptr + s * SUB); + const PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 2)); +#pragma unroll + for (int b = 0; b < SUB / 2; ++b) { + int q0, q1; + decode_int4_pair(pv[b], q0, q1); + float w0, w1; + if constexpr (Asym) { + w0 = (static_cast(q0) - zero) * scale; + w1 = (static_cast(q1) - zero) * scale; + } else { + w0 = static_cast(q0) * scale; + w1 = static_cast(q1) * scale; + } + const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); + acc += static_cast(a0) * w0; + acc += static_cast(a1) * w1; } - const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); - acc += static_cast(a0) * w0; - acc += static_cast(a1) * w1; } } From ed2c78faffabd60722f4331a6be38e83160f009e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 06:48:34 +0000 Subject: [PATCH 05/47] perf: widen int8 MoE decode to 32-wide vectorized ladder --- .../wrapper/include/sycl_tla_moe_decode.hpp | 83 +++++++++++++------ 1 file changed, 56 insertions(+), 27 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index de480993b0..5bd28df973 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -349,6 +349,42 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // type is used for sym and asym; the only difference is the sign interpretation // performed at decode time. // ---------------------------------------------------------------------------- + +// Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK weight +// bytes + a vec activation block). Templated on CHUNK so the +// caller can run a wide (32) stage first and a narrower (16) stage for the +// remainder, mirroring the int4 path: widening the first stage amortizes the +// per-group scale load and loop overhead across twice as many multiply-adds for +// the shipped group sizes (32/64/128/256) without regressing group_size == 16. +// sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so CHUNK is processed +// in 16-wide sub-blocks. The math is identical to the scalar path. +template +static inline void int8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, + float& acc) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); + constexpr int SUB = 16; + using ActVec = sycl::vec; + using ByteVec = sycl::vec; +#pragma unroll + for (int s = 0; s < CHUNK / SUB; ++s) { + const ActVec av = *reinterpret_cast(act_ptr + s * SUB); + const ByteVec wv = *reinterpret_cast(w_ptr + s * SUB); +#pragma unroll + for (int u = 0; u < SUB; ++u) { + const int qv = decode_int8(wv[u]); + float w; + if constexpr (Asym) { + w = (static_cast(qv) - zero) * scale; + } else { + w = static_cast(qv) * scale; + } + const ScalarT a = sycl::bit_cast(static_cast(av[u])); + acc += static_cast(a) * w; + } + } +} + template void launch_int8(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, @@ -397,39 +433,32 @@ void launch_int8(sycl::queue* q, const ScalarT* activations, const uint8_t* weig zero = static_cast(z_row[g]); } const int k_base = g * group_size; - // Vectorized path: 16 weights (16 bytes) + 16 activations per load. - // group_size is typically 128 (mult of 16); scalar tail handles - // anything that doesn't divide evenly. - constexpr int CHUNK = 16; - using ActVec = sycl::vec; - using ByteVec = sycl::vec; - static_assert(sizeof(ScalarT) == sizeof(uint16_t), - "ScalarT must be a 16-bit floating type"); - const int chunk_end = (group_size / CHUNK) * CHUNK; + // Vectorized ladder mirroring the int4 path: process 32 K-elements + // (32 weight bytes + vec activations) at a time, then a + // 16-wide stage for the remainder, then a scalar tail. Widening the + // first stage amortizes the per-group scale load and loop overhead + // for the shipped group sizes (32/64/128/256), while the 16-wide + // stage keeps group_size == 16 on the fast path. int kk = 0; - for (; kk < chunk_end; kk += CHUNK) { - const ActVec av = *reinterpret_cast(act_row + k_base + kk); - const ByteVec wv = *reinterpret_cast(w_row + k_base + kk); -#pragma unroll - for (int u = 0; u < CHUNK; ++u) { - const int q = decode_int8(wv[u]); - float w; - if constexpr (Asym) { - w = (static_cast(q) - zero) * scale; - } else { - w = static_cast(q) * scale; - } - const ScalarT a = sycl::bit_cast(static_cast(av[u])); - acc += static_cast(a) * w; - } + constexpr int CHUNK32 = 32; + const int end32 = (group_size / CHUNK32) * CHUNK32; + for (; kk < end32; kk += CHUNK32) { + int8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, scale, zero, + acc); + } + constexpr int CHUNK16 = 16; + const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; + for (; kk < end16; kk += CHUNK16) { + int8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, scale, zero, + acc); } for (; kk < group_size; ++kk) { - const int q = decode_int8(w_row[k_base + kk]); + const int qv = decode_int8(w_row[k_base + kk]); float w; if constexpr (Asym) { - w = (static_cast(q) - zero) * scale; + w = (static_cast(qv) - zero) * scale; } else { - w = static_cast(q) * scale; + w = static_cast(qv) * scale; } acc += static_cast(act_row[k_base + kk]) * w; } From 488fa34b2fb9a14d0eb75fdc7f27345fcf43193a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 06:49:21 +0000 Subject: [PATCH 06/47] perf: widen int2 MoE decode to 32-wide vectorized ladder --- .../wrapper/include/sycl_tla_moe_decode.hpp | 104 +++++++++++------- 1 file changed, 64 insertions(+), 40 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 5bd28df973..c4ef83a8c7 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -477,6 +477,54 @@ void launch_int8(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // Asym=false: signed 2-bit value in [-2, 1]; dequant = q * scale // Asym=true : unsigned 2-bit value in [0, 3]; dequant = (q - zero) * scale // ---------------------------------------------------------------------------- + +// Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK/4 +// packed weight bytes + a vec activation block). Templated on +// CHUNK so the caller can run a wide (32) stage first and a narrower (16) stage +// for the remainder, mirroring the int4/int8 paths. sycl::vec only supports +// widths of 1, 2, 3, 4, 8 or 16, so CHUNK is processed in 16-wide sub-blocks +// (16 activations + 4 packed bytes each). The math is identical to the scalar +// path. +template +static inline void int2_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, + float& acc) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); + constexpr int SUB = 16; + using ActVec = sycl::vec; + using PackVec = sycl::vec; +#pragma unroll + for (int s = 0; s < CHUNK / SUB; ++s) { + const ActVec av = *reinterpret_cast(act_ptr + s * SUB); + const PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 4)); +#pragma unroll + for (int b = 0; b < SUB / 4; ++b) { + int qq[4]; + decode_int2_quad(pv[b], qq); + float w0, w1, w2, w3; + if constexpr (Asym) { + w0 = (static_cast(qq[0]) - zero) * scale; + w1 = (static_cast(qq[1]) - zero) * scale; + w2 = (static_cast(qq[2]) - zero) * scale; + w3 = (static_cast(qq[3]) - zero) * scale; + } else { + w0 = static_cast(qq[0]) * scale; + w1 = static_cast(qq[1]) * scale; + w2 = static_cast(qq[2]) * scale; + w3 = static_cast(qq[3]) * scale; + } + const ScalarT a0 = sycl::bit_cast(static_cast(av[4 * b + 0])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[4 * b + 1])); + const ScalarT a2 = sycl::bit_cast(static_cast(av[4 * b + 2])); + const ScalarT a3 = sycl::bit_cast(static_cast(av[4 * b + 3])); + acc += static_cast(a0) * w0; + acc += static_cast(a1) * w1; + acc += static_cast(a2) * w2; + acc += static_cast(a3) * w3; + } + } +} + template void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, @@ -530,47 +578,23 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig zero = static_cast(z_row[g]); } const int k_base = g * group_size; - // Vectorized: 16 K-elements per chunk = 4 packed bytes (4 values - // each) plus a vec activation block. group_size is a - // multiple of 4 and typically 128 (mult of 16); scalar tail covers - // any leftover. We load activations via uint16_t to stay portable - // across SYCL implementations that may not provide - // sycl::vec. - constexpr int CHUNK = 16; - using ActVec = sycl::vec; - using PackVec = sycl::vec; - static_assert(sizeof(ScalarT) == sizeof(uint16_t), - "ScalarT must be a 16-bit floating type"); - const int chunk_end = (group_size / CHUNK) * CHUNK; + // Vectorized ladder mirroring the int4/int8 paths: process 32 + // K-elements (8 packed bytes + vec activations) at a + // time, then a 16-wide stage for the remainder, then a scalar tail. + // group_size is a multiple of 4; the wide stage amortizes the + // per-group scale load for the shipped group sizes (32/64/128/256). int kk = 0; - for (; kk < chunk_end; kk += CHUNK) { - const ActVec av = *reinterpret_cast(act_row + k_base + kk); - const PackVec pv = *reinterpret_cast(w_row + (k_base + kk) / 4); -#pragma unroll - for (int b = 0; b < CHUNK / 4; ++b) { - int q[4]; - decode_int2_quad(pv[b], q); - float w0, w1, w2, w3; - if constexpr (Asym) { - w0 = (static_cast(q[0]) - zero) * scale; - w1 = (static_cast(q[1]) - zero) * scale; - w2 = (static_cast(q[2]) - zero) * scale; - w3 = (static_cast(q[3]) - zero) * scale; - } else { - w0 = static_cast(q[0]) * scale; - w1 = static_cast(q[1]) * scale; - w2 = static_cast(q[2]) * scale; - w3 = static_cast(q[3]) * scale; - } - const ScalarT a0 = sycl::bit_cast(static_cast(av[4 * b + 0])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[4 * b + 1])); - const ScalarT a2 = sycl::bit_cast(static_cast(av[4 * b + 2])); - const ScalarT a3 = sycl::bit_cast(static_cast(av[4 * b + 3])); - acc += static_cast(a0) * w0; - acc += static_cast(a1) * w1; - acc += static_cast(a2) * w2; - acc += static_cast(a3) * w3; - } + constexpr int CHUNK32 = 32; + const int end32 = (group_size / CHUNK32) * CHUNK32; + for (; kk < end32; kk += CHUNK32) { + int2_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 4, scale, + zero, acc); + } + constexpr int CHUNK16 = 16; + const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; + for (; kk < end16; kk += CHUNK16) { + int2_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 4, scale, + zero, acc); } // Scalar tail (4 values per byte). for (; kk < group_size; kk += 4) { From 2284e77c83596321499046a3a56f01b1ee92964f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 06:53:12 +0000 Subject: [PATCH 07/47] perf: widen fp8 MoE decode to 32-wide vectorized ladder --- .../wrapper/include/sycl_tla_moe_decode.hpp | 64 +++++++++++++------ 1 file changed, 44 insertions(+), 20 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index c4ef83a8c7..76366f3ec1 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -632,6 +632,33 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // inline bit-manipulation path at compile time. The choice is driven at // launch time by the env var `ARK_FP8_DECODE_USE_LUT` (default: ON). // ---------------------------------------------------------------------------- + +// Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK weight +// bytes + a vec activation block). Templated on CHUNK so the +// caller can run a wide (32) stage first and a narrower (16) stage for the +// remainder, mirroring the int4/int8 paths. sycl::vec only supports widths of +// 1, 2, 3, 4, 8 or 16, so CHUNK is processed in 16-wide sub-blocks. The math is +// identical to the scalar path. +template +static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float& acc) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); + constexpr int SUB = 16; + using ActVec = sycl::vec; + using ByteVec = sycl::vec; +#pragma unroll + for (int s = 0; s < CHUNK / SUB; ++s) { + const ActVec av = *reinterpret_cast(act_ptr + s * SUB); + const ByteVec wv = *reinterpret_cast(w_ptr + s * SUB); +#pragma unroll + for (int u = 0; u < SUB; ++u) { + const float w = decode_fp8(wv[u]) * scale; + const ScalarT a = sycl::bit_cast(static_cast(av[u])); + acc += static_cast(a) * w; + } + } +} + template void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, int K, int group_size) { @@ -669,27 +696,24 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh for (int g = 0; g < num_groups_k; ++g) { const float scale = static_cast(s_row[g]); const int k_base = g * group_size; - // Vectorized: 16 weights (16 bytes) + 16 activations per load. - // Decode each FP8 byte to float inline, then apply the per-group - // scale. group_size is typically 128 (mult of 16); scalar tail - // covers anything that doesn't divide evenly. - constexpr int CHUNK = 16; - using ActVec = sycl::vec; - using ByteVec = sycl::vec; - static_assert(sizeof(ScalarT) == sizeof(uint16_t), - "ScalarT must be a 16-bit floating type"); - const int chunk_end = (group_size / CHUNK) * CHUNK; + // Vectorized ladder mirroring the int4/int8 paths: process 32 + // K-elements (32 weight bytes + vec activations) at a + // time, then a 16-wide stage for the remainder, then a scalar tail. + // Decode each FP8 byte to float, then apply the per-group scale. + // Widening the first stage amortizes the per-group scale load for + // the shipped group sizes (32/64/128/256). int kk = 0; - for (; kk < chunk_end; kk += CHUNK) { - const ActVec av = *reinterpret_cast(act_row + k_base + kk); - const ByteVec wv = *reinterpret_cast(w_row + k_base + kk); -#pragma unroll - for (int u = 0; u < CHUNK; ++u) { - const uint8_t raw = wv[u]; - const float w = decode_fp8(raw) * scale; - const ScalarT a = sycl::bit_cast(static_cast(av[u])); - acc += static_cast(a) * w; - } + constexpr int CHUNK32 = 32; + const int end32 = (group_size / CHUNK32) * CHUNK32; + for (; kk < end32; kk += CHUNK32) { + fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, scale, + acc); + } + constexpr int CHUNK16 = 16; + const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; + for (; kk < end16; kk += CHUNK16) { + fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, scale, + acc); } for (; kk < group_size; ++kk) { const uint8_t raw = w_row[k_base + kk]; From a039ad95a88d2d1e447624e224756d37cfe143bc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 07:07:34 +0000 Subject: [PATCH 08/47] perf: hoist per-group/per-expert scale out of fp8 MoE decode inner loop --- .../wrapper/include/sycl_tla_moe_decode.hpp | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 76366f3ec1..0f03604ccf 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -637,10 +637,16 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // bytes + a vec activation block). Templated on CHUNK so the // caller can run a wide (32) stage first and a narrower (16) stage for the // remainder, mirroring the int4/int8 paths. sycl::vec only supports widths of -// 1, 2, 3, 4, 8 or 16, so CHUNK is processed in 16-wide sub-blocks. The math is -// identical to the scalar path. +// 1, 2, 3, 4, 8 or 16, so CHUNK is processed in 16-wide sub-blocks. +// +// The per-group scale is constant across the whole group, so it is NOT applied +// here: this accumulates the raw dot product (sum of act * decoded_fp8) and the +// caller multiplies the group total by the scale once (Σ a·(w·s) == s·Σ a·w). +// For the per-expert / per-tensor scale case (group_size == K, one scale per +// output row) this collapses the whole K reduction to a single scale multiply, +// removing one multiply per K element on the decode hot path. template -static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float& acc) { +static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc) { static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); constexpr int SUB = 16; @@ -652,7 +658,7 @@ static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr const ByteVec wv = *reinterpret_cast(w_ptr + s * SUB); #pragma unroll for (int u = 0; u < SUB; ++u) { - const float w = decode_fp8(wv[u]) * scale; + const float w = decode_fp8(wv[u]); const ScalarT a = sycl::bit_cast(static_cast(av[u])); acc += static_cast(a) * w; } @@ -699,27 +705,32 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // Vectorized ladder mirroring the int4/int8 paths: process 32 // K-elements (32 weight bytes + vec activations) at a // time, then a 16-wide stage for the remainder, then a scalar tail. - // Decode each FP8 byte to float, then apply the per-group scale. - // Widening the first stage amortizes the per-group scale load for - // the shipped group sizes (32/64/128/256). + // The per-group scale is constant across the group, so accumulate the + // raw dot product here and apply the scale once below (Σ a·(w·s) == + // s·Σ a·w). Widening the first stage amortizes the per-group scale + // load for the shipped group sizes (32/64/128/256); hoisting the + // scale removes one multiply per K element, which is the dominant + // cost for the per-expert / per-tensor scale case (group_size == K). + float group_acc = 0.0f; int kk = 0; constexpr int CHUNK32 = 32; const int end32 = (group_size / CHUNK32) * CHUNK32; for (; kk < end32; kk += CHUNK32) { - fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, scale, - acc); + fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, + group_acc); } constexpr int CHUNK16 = 16; const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; for (; kk < end16; kk += CHUNK16) { - fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, scale, - acc); + fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, + group_acc); } for (; kk < group_size; ++kk) { const uint8_t raw = w_row[k_base + kk]; - const float w = decode_fp8(raw) * scale; - acc += static_cast(act_row[k_base + kk]) * w; + const float w = decode_fp8(raw); + group_acc += static_cast(act_row[k_base + kk]) * w; } + acc += group_acc * scale; } outputs[static_cast(token) * N + n_global] = static_cast(acc); From f2c47391ec13670047edb036a545bfe4366ef415 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:55:53 +0000 Subject: [PATCH 09/47] perf: defer per-group scale fold in native fp8 MoE prefill sym mainloop --- .../sycl_tla_moe_prefill_fp8_native.hpp | 55 ++++++++++++++----- 1 file changed, 41 insertions(+), 14 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_native.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_native.hpp index e3bed66e04..db8835aa1e 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_native.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_native.hpp @@ -75,10 +75,12 @@ // for g in [0, K/group_size): // scale group; barrier + scale reload // stage A[BM][group_size] into SLM // ONE cooperative load + barrier // scale = scales[e, n_col, g] // per-lane scale, loaded ONCE +// group_acc[m] = 0 // per-group deferred-scale acc // for sub in [0, group_size/BK): // BK sub-tile inside the group // load BK fp8 bytes for this lane // 4-byte chunked, unrolled -// w_col[k] = decode_fp8(byte) * scale -// acc[m] += sum_k a_slm[m, sub*BK+k] * w_col[k] +// w_col[k] = decode_fp8(byte) // scale deferred, not folded here +// group_acc[m] += sum_k a_slm[m, sub*BK+k] * w_col[k] +// acc[m] += group_acc[m] * scale // fold group scale ONCE per group // // vs. the original one-level loop that reloaded A + issued a barrier // once per BK-wide K-tile and reloaded the scale on every iteration. @@ -255,11 +257,15 @@ sycl::event launch_moe_prefill_fp8_native(sycl::queue* q, const ScalarT* activat // 3. Inner loop over `sub in [0, group_size/BK)` runs BK-wide // sub-tiles fully from SLM + registers with no extra barrier: // a. Fetch BK fp8 weight bytes (4-byte chunked, unrolled). - // b. Decode + fold scale in registers into w_col[BK]. - // c. MAC into acc[m] for each of the BM output rows. + // b. Decode fp8 -> float in registers into w_col[BK]. + // c. MAC into group_acc[m] for each of the BM output rows. + // Then fold the per-group scale into acc[m] once at the group + // boundary (deferred-scale), so the per-element `* scale` is + // lifted out of the weight-decode hot loop. // - // Per-BK partial-sum accumulation order is preserved bit-for-bit, - // so the FP8 parity tests (7e-2 tolerance) remain valid. + // Per-BK partial-sum accumulation order within a group is preserved, + // and folding the constant group scale once is distributive with the + // old per-element fold, so the FP8 parity tests (7e-2 tol) hold. // ----------------------------------------------------------------- const size_t w_row_stride = static_cast(K); // [E, N, K] row-major const size_t w_expert_stride = static_cast(N) * w_row_stride; @@ -298,6 +304,18 @@ sycl::event launch_moe_prefill_fp8_native(sycl::queue* q, const ScalarT* activat // in the prior revision). const float scale = static_cast(scales[s_row_base + static_cast(g)]); + // Per-scale-group deferred-scale accumulator (one FP32 per output + // row). The group scale is constant across the whole group, so we + // accumulate the raw A*W dot product here and fold in `scale` + // ONCE per group below (sum_k a*(w*s) == s * sum_k a*w). This + // removes the per-element `* scale` multiply from the weight-decode + // stage (group_size multiplies per lane per group) in exchange for + // BM multiplies at the group boundary — the same deferred-scale + // design the DPAS variant-B mainloop uses (`tCrC_group`). + float group_acc[BM]; +#pragma unroll + for (int m = 0; m < BM; ++m) group_acc[m] = 0.0f; + // --------- 3. Inner BK-sub-tile loop ------------------------- for (int sub = 0; sub < gs_per_tile; ++sub) { const int base_k_sub = sub * BK; @@ -307,7 +325,8 @@ sycl::event launch_moe_prefill_fp8_native(sycl::queue* q, const ScalarT* activat // aligned: `weights_NK` is 4-byte aligned (tensor storage) // and the offset `w_row_base + base_gk + base_k_sub` is a // multiple of BK = 32 (K % BK == 0, base_gk multiple of - // group_size which is a multiple of BK). + // group_size which is a multiple of BK). The per-group scale + // is NOT folded here — it is deferred to the group boundary. const size_t w_off = w_row_base + static_cast(base_gk) + static_cast(base_k_sub); const uint32_t* w_u32 = reinterpret_cast(weights_NK + w_off); @@ -320,16 +339,17 @@ sycl::event launch_moe_prefill_fp8_native(sycl::queue* q, const ScalarT* activat const uint8_t b1 = static_cast((w >> 8) & 0xFFu); const uint8_t b2 = static_cast((w >> 16) & 0xFFu); const uint8_t b3 = static_cast((w >> 24) & 0xFFu); - w_col[wi * 4 + 0] = moe_dequant::decode_fp8(b0) * scale; - w_col[wi * 4 + 1] = moe_dequant::decode_fp8(b1) * scale; - w_col[wi * 4 + 2] = moe_dequant::decode_fp8(b2) * scale; - w_col[wi * 4 + 3] = moe_dequant::decode_fp8(b3) * scale; + w_col[wi * 4 + 0] = moe_dequant::decode_fp8(b0); + w_col[wi * 4 + 1] = moe_dequant::decode_fp8(b1); + w_col[wi * 4 + 2] = moe_dequant::decode_fp8(b2); + w_col[wi * 4 + 3] = moe_dequant::decode_fp8(b3); } // 3b. MAC. For each output row m in this tile, dot-product // the length-BK slice of A[m] (staged in SLM) with - // `w_col`, accumulate into `acc[m]`. Same per-BK partial- - // sum shape as the original kernel to preserve numerics. + // `w_col`, accumulate into `group_acc[m]` (raw, unscaled). + // Same per-BK partial-sum shape as before to preserve + // numerics; the group scale is folded in once below. const size_t a_col_base = static_cast(base_k_sub); #pragma unroll for (int m = 0; m < BM; ++m) { @@ -341,10 +361,17 @@ sycl::event launch_moe_prefill_fp8_native(sycl::queue* q, const ScalarT* activat const float a_f = static_cast(a_slm[a_off]); sum += a_f * w_col[k]; } - acc[m] += sum; + group_acc[m] += sum; } } + // Fold the per-group scale into the running accumulator once per + // scale group (deferred-scale). Distributive with the per-element + // fold used previously, modulo FP32-accumulator ordering (within + // the 7e-2 FP8 tolerance in test_moe_prefill_accuracy.py). +#pragma unroll + for (int m = 0; m < BM; ++m) acc[m] += group_acc[m] * scale; + // Barrier before the next scale group re-stages A[]. it.barrier(sycl::access::fence_space::local_space); } From 90ba433293f82e6f71b5896c31c40d5dc0e64d32 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:45:29 +0000 Subject: [PATCH 10/47] fix: resolve S4 MoE dispatch conflict (opt_moe env-policy + our w4a16 tiles) --- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 41 +++++++------------ 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index 3b2532ed44..ea499f7b88 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -776,22 +776,13 @@ void moe_prefill_s4_dpas_per_group_dispatch( num_tokens_per_expert, E, group_size, prefetch_dist, \ prefetch_dist_scale, atomic_buffer); -<<<<<<< HEAD - // Four-tier bucketing on the average tokens-per-expert, matching the - // reference `w4a16` dispatch in vllm-xpu-kernels - // (`grouped_gemm_xe2_interface.hpp`): tiny M uses an 8-row tile, and the - // 32-row tile now covers M up to 128 (instead of jumping to the wide - // large-M tile at 33) so mid-size chunked-prefill batches avoid the - // padding waste of the 128-row tile. - if (A_avg_M <= 4) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_8); - } else if (A_avg_M <= 8) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_16); - } else if (A_avg_M <= 128) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_32); - } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy); -======= + // Explicit env overrides (from the opt_moe_kernel policy infrastructure) + // take precedence and let benchmarks pin a specific tile family. The + // default (Auto) path uses our dedicated four-tier `w4a16` bucketing: + // the halved packed-B stream lets a wider N tile pay off, tiny M uses an + // 8-row tile, and the 32-row tile covers M up to 128 (instead of jumping + // to the wide large-M tile at 33) so mid-size chunked-prefill batches + // avoid the padding waste of the 128-row tile. if (policy_mode == MoePrefillS4PolicyMode::W16A16) { ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w16a16_policy); } else if (policy_mode == MoePrefillS4PolicyMode::W8A16M16) { @@ -807,18 +798,16 @@ void moe_prefill_s4_dpas_per_group_dispatch( ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); } } else { - // Auto policy for the direct path. On BMG MiniMax prefill shapes the - // small-M w8a16_m32 tile is much faster than larger tiles even up to - // 8K prompts (average M ~= 341), because per-expert token counts are - // highly skewed and many experts underfill larger M tiles. - if (A_avg_M <= 8) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (A_avg_M <= 512) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + // Auto policy for the direct path -- dedicated w4a16 tiles. + if (A_avg_M <= 4) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_8); + } else if (A_avg_M <= 8) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_16); + } else if (A_avg_M <= 128) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_32); } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy); } ->>>>>>> origin/opt_moe_kernel } #undef ARK_DPAS_S4_PG_LAUNCH_SYM From f4efc89518d97f8eac6163323f994972b2dc37b2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:22:28 +0000 Subject: [PATCH 11/47] fix: restore INT4 MoE perf by defaulting S4 direct DPAS path ON --- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 35 ++++++------------- 1 file changed, 10 insertions(+), 25 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index ea499f7b88..153f13802c 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -178,24 +178,6 @@ enum class MoePrefillS4PolicyMode { W16A16, }; -inline bool moe_prefill_dpas_s4_force_direct() { - const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4"); - if (env != nullptr) { - std::string s(env); - for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "1" || s == "true" || s == "on" || s == "yes" || - s == "direct" || s == "force" || s == "packed") return true; - } - - const char* policy = std::getenv("ARK_MOE_PREFILL_DPAS_S4_POLICY"); - if (policy == nullptr) return false; - std::string p(policy); - for (auto& c : p) c = static_cast(std::tolower(static_cast(c))); - return p == "w16a16" || p == "w16" || p == "w8a16" || p == "w8" || - p == "w8a16_m16" || p == "w8_m16" || p == "m16" || - p == "w8a16_m32" || p == "w8_m32" || p == "m32"; -} - inline MoePrefillS4PolicyMode moe_prefill_dpas_s4_policy_mode() { const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4_POLICY"); if (env == nullptr) return MoePrefillS4PolicyMode::Auto; @@ -815,11 +797,14 @@ void moe_prefill_s4_dpas_per_group_dispatch( } // --------------------------------------------------------------------------- -// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4`. The packed-nibble direct -// DPAS path is opt-in and uses the auto policy above by default. Current BMG -// MiniMax measurements prefer the w8a16_m32 tile family for real prefill -// distributions; callers can still force another policy for experiments via -// `ARK_MOE_PREFILL_DPAS_S4_POLICY=w16a16|w8a16|w8a16_m16|w8a16_m32`. +// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default ON). The packed-nibble +// direct DPAS path is the fast INT4 path: it reads the halved packed-B stream +// and folds the nibble->act upcast into the DPAS mainloop, so it must stay +// enabled by default to preserve INT4 MoE performance. By default it uses the +// auto policy above (our dedicated four-tier `w4a16` tiles). Callers can still +// pin another tile family for experiments via +// `ARK_MOE_PREFILL_DPAS_S4_POLICY=w16a16|w8a16|w8a16_m16|w8a16_m32`, which only +// selects the GEMM tile -- it does not gate the path. // // Truthy values (case-insensitive): "1", "true", "on", "yes". // Explicit "0" / "false" / "off" / "no" disable. Re-read on every @@ -837,11 +822,11 @@ void moe_prefill_s4_dpas_per_group_dispatch( // --------------------------------------------------------------------------- inline bool moe_prefill_dpas_s4_enabled() { const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4"); - if (env == nullptr) return false; + if (env == nullptr) return true; // default ON -- fast INT4 direct path std::string s(env); for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); if (s == "0" || s == "false" || s == "off" || s == "no") return false; - return moe_prefill_dpas_s4_force_direct(); + return true; } // --------------------------------------------------------------------------- From 5466928079d00e85c09c29551382ea0ae091f8be Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:44:21 +0000 Subject: [PATCH 12/47] revert: restore int4 MoE to original base implementation --- .../wrapper/include/sycl_tla_moe_decode.hpp | 94 ++---- .../wrapper/include/sycl_tla_moe_mixed.hpp | 7 +- .../include/sycl_tla_moe_prefill_fp8_dpas.hpp | 29 -- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 269 ++++-------------- 4 files changed, 94 insertions(+), 305 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 0f03604ccf..adfb775aa3 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -204,50 +204,6 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight // of byte i, the value at k = 2*i+1 is the HIGH nibble. This matches the // existing CPU/XPU `packq` layout for S4_CLIP weights. // ---------------------------------------------------------------------------- - -// Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK/2 -// packed weight bytes + a vec activation block). Templated on -// CHUNK so the caller can run a wide (32) stage first and a narrower (16) -// stage for the remainder, which keeps the fast path active for group sizes -// that are a multiple of 32 (32/64/128/256 -- the shipped quant configs) -// without regressing group_size == 16 (which drops straight to the 16-wide -// stage). The math is identical to the scalar path. -template -static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, - float& acc) { - static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); - static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); - // sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so a single - // vec load is illegal. Process the chunk in 16-wide sub-blocks - // (16 activations + 8 packed weight bytes each), which keeps CHUNK == 32 - // valid while reusing the same code path for CHUNK == 16. - constexpr int SUB = 16; - using ActVec = sycl::vec; - using PackVec = sycl::vec; -#pragma unroll - for (int s = 0; s < CHUNK / SUB; ++s) { - const ActVec av = *reinterpret_cast(act_ptr + s * SUB); - const PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 2)); -#pragma unroll - for (int b = 0; b < SUB / 2; ++b) { - int q0, q1; - decode_int4_pair(pv[b], q0, q1); - float w0, w1; - if constexpr (Asym) { - w0 = (static_cast(q0) - zero) * scale; - w1 = (static_cast(q1) - zero) * scale; - } else { - w0 = static_cast(q0) * scale; - w1 = static_cast(q1) * scale; - } - const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); - acc += static_cast(a0) * w0; - acc += static_cast(a1) * w1; - } - } -} - template void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, @@ -297,27 +253,39 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig zero = static_cast(z_row[g]); } const int k_base = g * group_size; - // Vectorized ladder: process 32 K-elements at a time (16 packed - // weight bytes + vec activation block), then a 16-wide - // stage for the remainder, then a scalar tail. Widening the first - // stage to 32 amortizes the per-group scale load and loop overhead - // across twice as many multiply-adds for the shipped group sizes - // (32/64/128/256), while the 16-wide stage keeps group_size == 16 - // on the fast path. + // Vectorized path: process 16 K-elements at a time, which is + // 8 packed weight bytes and a vec activation block. + // group_size is a multiple of 16 in every supported config + // (group_size >= 32, even); a scalar tail loop covers leftovers. + constexpr int CHUNK = 16; + using ActVec = sycl::vec; + using PackVec = sycl::vec; + static_assert(sizeof(ScalarT) == sizeof(uint16_t), + "ScalarT must be a 16-bit floating type"); + const int chunk_end = (group_size / CHUNK) * CHUNK; int kk = 0; - constexpr int CHUNK32 = 32; - const int end32 = (group_size / CHUNK32) * CHUNK32; - for (; kk < end32; kk += CHUNK32) { - int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, - zero, acc); - } - constexpr int CHUNK16 = 16; - const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; - for (; kk < end16; kk += CHUNK16) { - int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, - zero, acc); + for (; kk < chunk_end; kk += CHUNK) { + const ActVec av = *reinterpret_cast(act_row + k_base + kk); + const PackVec pv = *reinterpret_cast(w_row + (k_base + kk) / 2); +#pragma unroll + for (int b = 0; b < CHUNK / 2; ++b) { + int q0, q1; + decode_int4_pair(pv[b], q0, q1); + float w0, w1; + if constexpr (Asym) { + w0 = (static_cast(q0) - zero) * scale; + w1 = (static_cast(q1) - zero) * scale; + } else { + w0 = static_cast(q0) * scale; + w1 = static_cast(q1) * scale; + } + const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); + acc += static_cast(a0) * w0; + acc += static_cast(a1) * w1; + } } - // Scalar tail for group_size not divisible by 16. + // Scalar tail for group_size not divisible by CHUNK. for (; kk < group_size; kk += 2) { const uint8_t packed = w_row[(k_base + kk) / 2]; int q0, q1; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp index 2312702563..e6a3ececcd 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp @@ -129,7 +129,6 @@ 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; @@ -278,17 +277,15 @@ 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 - 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(E), static_cast(k_words_padded), + sycl::range<3> global_fast{static_cast(E), static_cast(k_words), static_cast((N + WG_N - 1) / WG_N) * WG_N}; - sycl::range<3> local_fast{1, static_cast(WG_K_INT4_FAST), static_cast(WG_N)}; + sycl::range<3> local_fast{1, 1, static_cast(WG_N)}; q->parallel_for>( sycl::nd_range<3>(global_fast, local_fast), [=](sycl::nd_item<3> it) { const int e = static_cast(it.get_global_id(0)); if (num_tokens_per_expert != nullptr && num_tokens_per_expert[e] == 0) return; const int kw = static_cast(it.get_global_id(1)); - if (kw >= k_words) return; const int n = static_cast(it.get_global_id(2)); if (n >= N) return; const int k_base = kw * PACK_K_INT4_FAST; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index 05bce99f76..35db93cac2 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -193,35 +193,6 @@ class dpas_w8a16_policy_m_32 : public dpas_policy_base { using SGLayout = Layout, Stride<_4, _1, _0>>; }; -// --------------------------------------------------------------------------- -// 4-bit (S4 / w4a16) tile policies. -// -// The S4 mixed-input prefill mainloop (`sycl_tla_moe_prefill_s4_dpas.hpp`) -// reads a *halved* B-side byte stream (two nibbles per byte). With the same -// M/N footprint the packed weight fits in half the L2/GRF traffic of the INT8 -// path, so a larger N tile pays off: the default large-M policy uses a -// 128x256x32 WG tile (vs. the INT8 128x128x16), matching the reference -// `w4a16_policy` in vllm-xpu-kernels `csrc/xpu/grouped_gemm/xe_2/ -// gemm_xe2_policy.hpp`. The small-M buckets (m_8/m_16/m_32) reuse the same -// 64-wide N tiles as the INT8 path -- the reference uses identical shapes -// there. Only the default (large-M) tile and the new m_8 bucket differ, so we -// define those two here and alias m_16/m_32 to the shared shapes in the S4 -// header. -// --------------------------------------------------------------------------- -class dpas_w4a16_policy : public dpas_policy_base { - public: - using WGTile = Shape<_128, _256, _32>; - using SGLayout = Layout, Stride<_8, _1, _0>>; - - using GmemTiledCopyD = XE_STORE_2D<16, 8, 32>; -}; - -class dpas_w4a16_policy_m_8 : public dpas_policy_base { - public: - using WGTile = Shape<_8, _64, _32>; - using SGLayout = Layout, Stride<_4, _1, _0>>; -}; - // --------------------------------------------------------------------------- // `apply_scale` -- inline-asm per-lane multiply of a bf16/fp16 fragment by // an FP32 scalar. Copied verbatim from vllm-xpu-kernels `gemm_xe2.hpp`. diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index 153f13802c..fad8d87eb7 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -91,7 +91,6 @@ #pragma once #include -#include #include #include #include @@ -130,66 +129,11 @@ using ::ark::moe_dpas_fp8::dpas_w16a16_policy; using ::ark::moe_dpas_fp8::dpas_w8a16_policy; using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; -// Dedicated 4-bit tile policies. The default (large-M) tile is 128x256x32 -// (halved packed-B stream lets a wider N tile pay off) and the m_8 bucket is -// new; both mirror the reference `w4a16_policy*` in vllm-xpu-kernels. The -// m_16 / m_32 buckets share the INT8 64-wide N tiles, so they are aliased to -// the existing `dpas_w8a16_policy_m_16 / _m_32` shapes rather than duplicated. -using ::ark::moe_dpas_fp8::dpas_w4a16_policy; -using ::ark::moe_dpas_fp8::dpas_w4a16_policy_m_8; -using dpas_w4a16_policy_m_16 = ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; -using dpas_w4a16_policy_m_32 = ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; using ::ark::moe_dpas_fp8::ScaleMode; using ::ark::moe_dpas_fp8::cute_scalar; using ::ark::moe_dpas_fp8::cute_scalar_t; using ::ark::moe_dpas_fp8::make_moe_tensor; -inline int _moe_prefill_s4_env_int(const char* name, int default_v, - int min_v, int max_v) { - const char* env = std::getenv(name); - if (env == nullptr) return default_v; - char* end = nullptr; - long v = std::strtol(env, &end, 10); - if (end == env || *end != '\0') return default_v; - if (v < static_cast(min_v)) return min_v; - if (v > static_cast(max_v)) return max_v; - return static_cast(v); -} - -inline int moe_prefill_dpas_s4_prefetch_dist() { - // Runtime-tunable A/B prefetch depth for S4 per-group mainloop. - // Default keeps current behavior. - return _moe_prefill_s4_env_int("ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST", - 3, 1, 16); -} - -inline int moe_prefill_dpas_s4_prefetch_dist_scale() { - // Runtime-tunable scale prefetch lookahead (in groups). - // Default keeps current behavior. - return _moe_prefill_s4_env_int( - "ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST_SCALE", 3, 1, 16); -} - -enum class MoePrefillS4PolicyMode { - Auto, - W8A16, - W8A16M16, - W8A16M32, - W16A16, -}; - -inline MoePrefillS4PolicyMode moe_prefill_dpas_s4_policy_mode() { - const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4_POLICY"); - if (env == nullptr) return MoePrefillS4PolicyMode::Auto; - std::string s(env); - for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "w16a16" || s == "w16") return MoePrefillS4PolicyMode::W16A16; - if (s == "w8a16_m16" || s == "w8_m16" || s == "m16") return MoePrefillS4PolicyMode::W8A16M16; - if (s == "w8a16_m32" || s == "w8_m32" || s == "m32") return MoePrefillS4PolicyMode::W8A16M32; - if (s == "w8a16" || s == "w8") return MoePrefillS4PolicyMode::W8A16; - return MoePrefillS4PolicyMode::Auto; -} - // --------------------------------------------------------------------------- // Variant B -- per-K-group S4 (sym) mainloop. // @@ -228,9 +172,8 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( const ElementBI* Bias, DTensor& C, // (M,N) -- ElementA Coord blk_coord, - TiledMMA const& mma, - int prefetch_dist, - int prefetch_dist_scale) { + TiledMMA const& mma) { + using TA = typename ATensor::element_type; using TB = typename BTensor::element_type; static_assert(std::is_same_v, "xe_gemm_s4_pergroup: ElementB must be cutlass::int4b_t (sym only)"); @@ -261,13 +204,8 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( auto thr_copy_b = copy_b.get_slice(local_id); auto thr_copy_c = copy_c.get_slice(local_id); - // Ping-pong SG fragments for software pipelining: - // while tile k computes on one fragment pair, tile k+1 is unpacked - // into the other pair. - auto tCrA0 = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); - auto tCrB0 = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); - auto tCrA1 = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); - auto tCrB1 = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + auto tCrA = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_, _, 0)); auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_, _, 0)); @@ -288,8 +226,11 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( auto pAgA = thr_prefetch_A.partition_S(gA); auto pBgB = thr_prefetch_B.partition_S(gB); - // Prefetch distances are host-selected and threaded into the device - // path so no env/stdlib calls are needed inside SYCL kernels. + // Prefetch distance mirrors `xe_gemm_int_pergroup<>` for now. + // On-hardware perf tuning may want to grow `prefetch_dist` on the + // packed path since the B stream is half the bandwidth. + const int prefetch_dist = 3; + const int prefetch_dist_scale = 3; constexpr auto barrier_scope = ScopeWorkgroup; int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); int k_tile_prefetch = 0; @@ -308,8 +249,11 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( auto n_tile_start = wg_n * tile_n; auto sg_local_n_coord = cutlass::get_sub_group_id() % ATOM_N; + auto sg_local_m_coord = cutlass::get_sub_group_id() / ATOM_N; int sg_local_id = cutlass::get_sub_group_local_id(); int n_sg_start = sg_local_n_coord * SG_N; + int m_sg_start = sg_local_m_coord * SG_M; + int m_tile_start = wg_m * tile_m; int group_num = get<1>(A.shape()) / group_size; // Group-local accumulator: same fragment shape as `tCrC`, cleared at @@ -323,16 +267,6 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( // Per-SG per-N scale cache. Same layout / semantics as the INT8 // per-group path. float sg_scale[sg_n_strides]; - int sg_scale_base[sg_n_strides]; - - // Precompute per-lane base offsets into Scales so the hot loop only - // adds group_idx. - CUTLASS_PRAGMA_UNROLL - for (int sn = 0; sn < sg_n_strides; ++sn) { - int sg_local_n = sn * sg_local_range + sg_local_id; - sg_scale_base[sn] = - (n_tile_start + n_sg_start + sg_local_n) * group_num; - } CUTE_UNROLL for (; k_tile_prefetch < prefetch_dist; k_tile_prefetch++) { @@ -356,48 +290,27 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( } } - if (k_tile_count == 0) { - return; - } - - // Prime the pipeline with tile 0. - barrier_arrive(barrier_scope); - copy(copy_a, tAgA(_, _, _, 0), tArA); - copy(copy_b, tBgB(_, _, _, 0), tBrB); - if (k_tile_prefetch < k_tile_count) { - prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); - prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); - } - reorder(tArA, tCrA0); - reorder(tBrB, tCrB0); - barrier_wait(barrier_scope); - - int curr_buf = 0; - static constexpr int tiles_per_group = group_size / tile_k; - int group_tile_remaining = 0; - int group_idx = 0; for (int k_tile = 0; k_tile < k_tile_count; k_tile++, k_tile_prefetch++) { - const bool has_next = (k_tile + 1 < k_tile_count); - const int next_tile = k_tile + 1; - - if (has_next) { - barrier_arrive(barrier_scope); - copy(copy_a, tAgA(_, _, _, next_tile), tArA); - copy(copy_b, tBgB(_, _, _, next_tile), tBrB); - } - - // Group-boundary scale reload via a tile counter instead of per-tile - // modulo/div in the hot loop. For all supported S4 shapes in this - // file, `group_size` is a multiple of `tile_k`. - if (group_tile_remaining == 0) { - group_idx = k_tile / tiles_per_group; + barrier_arrive(barrier_scope); + + copy(copy_a, tAgA(_, _, _, k_tile), tArA); + copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + + // Group-boundary scale reload. Same math as the INT8 per-group + // path: `tile_k` is expressed in element units (nibbles), not + // bytes, so `k_tile * tile_k` is the reduction position in + // *element* space and the modulo test against `group_size` matches + // the scale-tensor layout `[E, N, K/group_size]` unchanged. + if (k_tile * tile_k % group_size == 0) { + int group_idx = (k_tile * tile_k) / group_size; CUTLASS_PRAGMA_UNROLL for (int sn = 0; sn < sg_n_strides; ++sn) { - sg_scale[sn] = - static_cast(Scales[sg_scale_base[sn] + group_idx]); + int sg_local_n = sn * sg_local_range + sg_local_id; + sg_scale[sn] = static_cast( + Scales[(n_tile_start + n_sg_start + sg_local_n) * group_num + group_idx]); } - if (group_idx + prefetch_dist_scale < group_num) { + if ((group_idx + prefetch_dist_scale) * group_size < shape<1>(A)) { auto next_scales_tensor = make_tensor( make_gmem_ptr(reinterpret_cast( Scales + (n_tile_start + n_sg_start) * group_num + @@ -411,41 +324,32 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( make_identity_tensor(make_shape(Int{}, Int<1>{}))); prefetch(prefetch_scales, pSgS(_, 0, 0)); } - group_tile_remaining = tiles_per_group; } - if (has_next && k_tile_prefetch < k_tile_count) { + if (k_tile_prefetch < k_tile_count) { prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); } - if (has_next) { - // Pipeline stage: unpack tile k+1 while tile k compute is pending. - // `reorder` performs in-register `int4b_t -> ElementA` conversion via - // `NumericArrayConverter`. - if (curr_buf == 0) { - reorder(tArA, tCrA1); - reorder(tBrB, tCrB1); - } else { - reorder(tArA, tCrA0); - reorder(tBrB, tCrB0); - } - } + // `reorder` performs the in-register `int4b_t -> ElementA` unpack + // + sign-extend + cast via `cutlass::NumericArrayConverter< + // ElementA, cutlass::int4b_t, N>`. Once `tCrB` carries bf16/fp16 + // values it is compatible with the same DPAS atom used by the FP8 + // / INT8 per-group paths. See the header preamble open-question + // (1) -- if the pinned cutlass-sycl is missing this converter + // specialisation this line is where the build fails. + reorder(tArA, tCrA); + reorder(tBrB, tCrB); // HOT MAINLOOP -- MMA accumulates into `tCrC_group`. Per-N scale // is applied ONCE at the end of the group in the fold block below. - if (curr_buf == 0) { - cute::gemm(mma, tCrA0, tCrB0, tCrC_group); - } else { - cute::gemm(mma, tCrA1, tCrB1, tCrC_group); - } + cute::gemm(mma, tCrA, tCrB, tCrC_group); - group_tile_remaining -= 1; - // Group-boundary fold. Fires when either (a) we've consumed all - // tiles in the current group, or (b) this is the last reduction tile - // (tail-group protection). - const bool is_group_end = (group_tile_remaining == 0) || - (k_tile + 1 == k_tile_count); + // Group-boundary fold. Fires when either (a) the NEXT k_tile would + // start a new scale group, or (b) we've reached the last k_tile of + // the K reduction (tail-group protection). + const bool is_group_end = (((k_tile + 1) * tile_k) % group_size == 0) || + (k_tile + 1 == k_tile_count); if (is_group_end) { CUTLASS_PRAGMA_UNROLL for (int sn = 0; sn < sg_n_strides; ++sn) { @@ -459,10 +363,7 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( } } - if (has_next) { - barrier_wait(barrier_scope); - curr_buf ^= 1; - } + barrier_wait(barrier_scope); } if (Bias != nullptr) { @@ -507,10 +408,7 @@ CUTE_DEVICE void MoEGEMM_s4(const ElementA* Activations, const int* rows_per_expert, const int32_t num_experts, const int32_t group_size, const int32_t gemm_n, - const int32_t gemm_k, - int prefetch_dist, - int prefetch_dist_scale, - int32_t* atomic_buffer, + const int32_t gemm_k, int32_t* atomic_buffer, const sycl::local_accessor& slm_mem_const) { constexpr char actual_layout_of_B = LayoutKindB ^ ('R' ^ 'C'); @@ -597,7 +495,7 @@ CUTE_DEVICE void MoEGEMM_s4(const ElementA* Activations, #define ARK_MOE_DPAS_S4_GROUP_CALLER(GS) \ xe_gemm_s4_pergroup( \ A_tensor, B_tensor, ptr_Scales_curr_batch, ptr_Bias_curr_batch, \ - D_tensor, tile_coord, mma, prefetch_dist, prefetch_dist_scale); + D_tensor, tile_coord, mma); if (group_size == 32) { ARK_MOE_DPAS_S4_GROUP_CALLER(32) } else if (group_size == 64) { @@ -636,10 +534,7 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, const ElementBI* bias, ElementD* outputs, const int gemm_n, const int gemm_k, const int* rows_per_expert, const int num_experts, - const int group_size, - int prefetch_dist, - int prefetch_dist_scale, - int32_t* atomic_buffer) { + const int group_size, int32_t* atomic_buffer) { using ElementA_non_CV = cutlass::platform::remove_cv_t; // DPAS atom keeps its bf16/fp16 x bf16/fp16 -> fp32 shape; the S4 B // tensor is upcast to ElementA in `reorder(tBrB, tCrB)` in the @@ -686,7 +581,7 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, layoutB, 'R'>( activations, weights, scales, bias, outputs, mma, rows_per_expert, num_experts, group_size, gemm_n, gemm_k, - prefetch_dist, prefetch_dist_scale, atomic_buffer, local_mem); + atomic_buffer, local_mem); }); }); @@ -741,9 +636,6 @@ void moe_prefill_s4_dpas_per_group_dispatch( reinterpret_cast(weights_NKp); int A_avg_M = total_tokens / E; - const int prefetch_dist = moe_prefill_dpas_s4_prefetch_dist(); - const int prefetch_dist_scale = moe_prefill_dpas_s4_prefetch_dist_scale(); - const MoePrefillS4PolicyMode policy_mode = moe_prefill_dpas_s4_policy_mode(); int32_t* atomic_buffer = sycl::malloc_device(1, *q); if (atomic_buffer == nullptr) { @@ -755,41 +647,14 @@ void moe_prefill_s4_dpas_per_group_dispatch( MoEGEMMLauncher_s4<'R', 'C', policy>( \ *q, activations_ca, weights_i4, scales_ca, \ static_cast(nullptr), outputs_ca, N, K, \ - num_tokens_per_expert, E, group_size, prefetch_dist, \ - prefetch_dist_scale, atomic_buffer); - - // Explicit env overrides (from the opt_moe_kernel policy infrastructure) - // take precedence and let benchmarks pin a specific tile family. The - // default (Auto) path uses our dedicated four-tier `w4a16` bucketing: - // the halved packed-B stream lets a wider N tile pay off, tiny M uses an - // 8-row tile, and the 32-row tile covers M up to 128 (instead of jumping - // to the wide large-M tile at 33) so mid-size chunked-prefill batches - // avoid the padding waste of the 128-row tile. - if (policy_mode == MoePrefillS4PolicyMode::W16A16) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w16a16_policy); - } else if (policy_mode == MoePrefillS4PolicyMode::W8A16M16) { + num_tokens_per_expert, E, group_size, atomic_buffer); + + if (A_avg_M <= 8) { ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (policy_mode == MoePrefillS4PolicyMode::W8A16M32) { + } else if (A_avg_M <= 32) { ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); - } else if (policy_mode == MoePrefillS4PolicyMode::W8A16) { - if (A_avg_M <= 8) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (A_avg_M <= 32) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); - } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); - } } else { - // Auto policy for the direct path -- dedicated w4a16 tiles. - if (A_avg_M <= 4) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_8); - } else if (A_avg_M <= 8) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_16); - } else if (A_avg_M <= 128) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_32); - } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy); - } + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); } #undef ARK_DPAS_S4_PG_LAUNCH_SYM @@ -797,32 +662,20 @@ void moe_prefill_s4_dpas_per_group_dispatch( } // --------------------------------------------------------------------------- -// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default ON). The packed-nibble -// direct DPAS path is the fast INT4 path: it reads the halved packed-B stream -// and folds the nibble->act upcast into the DPAS mainloop, so it must stay -// enabled by default to preserve INT4 MoE performance. By default it uses the -// auto policy above (our dedicated four-tier `w4a16` tiles). Callers can still -// pin another tile family for experiments via -// `ARK_MOE_PREFILL_DPAS_S4_POLICY=w16a16|w8a16|w8a16_m16|w8a16_m32`, which only -// selects the GEMM tile -- it does not gate the path. +// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default ON, semantics +// identical to `moe_prefill_dpas_int_enabled` / `moe_prefill_dpas_fp8 +// _enabled`). Decoupled from `ARK_MOE_PREFILL_DPAS_INT8` so this new +// single-pass path can be disabled in isolation if it regresses -- +// switching S4 off falls back to the S4->S8 upcast + INT8 DPAS path +// which is itself gated by `ARK_MOE_PREFILL_DPAS_INT8`. // // Truthy values (case-insensitive): "1", "true", "on", "yes". // Explicit "0" / "false" / "off" / "no" disable. Re-read on every // call so benchmarks / tests can toggle the path in-process. -// -// Additional runtime tuning knobs for the S4 per-group mainloop: -// - ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST (default 3, clamp [1, 16]) -// - ARK_MOE_PREFILL_DPAS_S4_PREFETCH_DIST_SCALE (default 3, clamp [1, 16]) -// - ARK_MOE_PREFILL_DPAS_S4_POLICY (auto|w8a16|w8a16_m16| -// w8a16_m32|w16a16, -// default auto) -// -// These tune A/B tile prefetch depth and per-group scale prefetch lookahead -// respectively; `*_POLICY` selects the GEMM tile family. // --------------------------------------------------------------------------- inline bool moe_prefill_dpas_s4_enabled() { const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4"); - if (env == nullptr) return true; // default ON -- fast INT4 direct path + if (env == nullptr) return true; // default ON std::string s(env); for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); if (s == "0" || s == "false" || s == "off" || s == "no") return false; From 515f839a0d2d377a4eddbacd0a3d0f39abb86485 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:57:12 +0000 Subject: [PATCH 13/47] Changes before error encountered Agent-Logs-Url: https://github.com/intel/auto-round/sessions/9f204578-f116-4469-ac32-efbb1cdcfadc --- .../wrapper/include/sycl_tla_moe_decode.hpp | 94 +++++++++++++------ .../include/sycl_tla_moe_prefill_fp8_dpas.hpp | 29 ++++++ .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 27 +++++- 3 files changed, 114 insertions(+), 36 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index adfb775aa3..0f03604ccf 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -204,6 +204,50 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight // of byte i, the value at k = 2*i+1 is the HIGH nibble. This matches the // existing CPU/XPU `packq` layout for S4_CLIP weights. // ---------------------------------------------------------------------------- + +// Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK/2 +// packed weight bytes + a vec activation block). Templated on +// CHUNK so the caller can run a wide (32) stage first and a narrower (16) +// stage for the remainder, which keeps the fast path active for group sizes +// that are a multiple of 32 (32/64/128/256 -- the shipped quant configs) +// without regressing group_size == 16 (which drops straight to the 16-wide +// stage). The math is identical to the scalar path. +template +static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, + float& acc) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); + // sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so a single + // vec load is illegal. Process the chunk in 16-wide sub-blocks + // (16 activations + 8 packed weight bytes each), which keeps CHUNK == 32 + // valid while reusing the same code path for CHUNK == 16. + constexpr int SUB = 16; + using ActVec = sycl::vec; + using PackVec = sycl::vec; +#pragma unroll + for (int s = 0; s < CHUNK / SUB; ++s) { + const ActVec av = *reinterpret_cast(act_ptr + s * SUB); + const PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 2)); +#pragma unroll + for (int b = 0; b < SUB / 2; ++b) { + int q0, q1; + decode_int4_pair(pv[b], q0, q1); + float w0, w1; + if constexpr (Asym) { + w0 = (static_cast(q0) - zero) * scale; + w1 = (static_cast(q1) - zero) * scale; + } else { + w0 = static_cast(q0) * scale; + w1 = static_cast(q1) * scale; + } + const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); + acc += static_cast(a0) * w0; + acc += static_cast(a1) * w1; + } + } +} + template void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, @@ -253,39 +297,27 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig zero = static_cast(z_row[g]); } const int k_base = g * group_size; - // Vectorized path: process 16 K-elements at a time, which is - // 8 packed weight bytes and a vec activation block. - // group_size is a multiple of 16 in every supported config - // (group_size >= 32, even); a scalar tail loop covers leftovers. - constexpr int CHUNK = 16; - using ActVec = sycl::vec; - using PackVec = sycl::vec; - static_assert(sizeof(ScalarT) == sizeof(uint16_t), - "ScalarT must be a 16-bit floating type"); - const int chunk_end = (group_size / CHUNK) * CHUNK; + // Vectorized ladder: process 32 K-elements at a time (16 packed + // weight bytes + vec activation block), then a 16-wide + // stage for the remainder, then a scalar tail. Widening the first + // stage to 32 amortizes the per-group scale load and loop overhead + // across twice as many multiply-adds for the shipped group sizes + // (32/64/128/256), while the 16-wide stage keeps group_size == 16 + // on the fast path. int kk = 0; - for (; kk < chunk_end; kk += CHUNK) { - const ActVec av = *reinterpret_cast(act_row + k_base + kk); - const PackVec pv = *reinterpret_cast(w_row + (k_base + kk) / 2); -#pragma unroll - for (int b = 0; b < CHUNK / 2; ++b) { - int q0, q1; - decode_int4_pair(pv[b], q0, q1); - float w0, w1; - if constexpr (Asym) { - w0 = (static_cast(q0) - zero) * scale; - w1 = (static_cast(q1) - zero) * scale; - } else { - w0 = static_cast(q0) * scale; - w1 = static_cast(q1) * scale; - } - const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); - acc += static_cast(a0) * w0; - acc += static_cast(a1) * w1; - } + constexpr int CHUNK32 = 32; + const int end32 = (group_size / CHUNK32) * CHUNK32; + for (; kk < end32; kk += CHUNK32) { + int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, + zero, acc); + } + constexpr int CHUNK16 = 16; + const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; + for (; kk < end16; kk += CHUNK16) { + int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, + zero, acc); } - // Scalar tail for group_size not divisible by CHUNK. + // Scalar tail for group_size not divisible by 16. for (; kk < group_size; kk += 2) { const uint8_t packed = w_row[(k_base + kk) / 2]; int q0, q1; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index 35db93cac2..05bce99f76 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -193,6 +193,35 @@ class dpas_w8a16_policy_m_32 : public dpas_policy_base { using SGLayout = Layout, Stride<_4, _1, _0>>; }; +// --------------------------------------------------------------------------- +// 4-bit (S4 / w4a16) tile policies. +// +// The S4 mixed-input prefill mainloop (`sycl_tla_moe_prefill_s4_dpas.hpp`) +// reads a *halved* B-side byte stream (two nibbles per byte). With the same +// M/N footprint the packed weight fits in half the L2/GRF traffic of the INT8 +// path, so a larger N tile pays off: the default large-M policy uses a +// 128x256x32 WG tile (vs. the INT8 128x128x16), matching the reference +// `w4a16_policy` in vllm-xpu-kernels `csrc/xpu/grouped_gemm/xe_2/ +// gemm_xe2_policy.hpp`. The small-M buckets (m_8/m_16/m_32) reuse the same +// 64-wide N tiles as the INT8 path -- the reference uses identical shapes +// there. Only the default (large-M) tile and the new m_8 bucket differ, so we +// define those two here and alias m_16/m_32 to the shared shapes in the S4 +// header. +// --------------------------------------------------------------------------- +class dpas_w4a16_policy : public dpas_policy_base { + public: + using WGTile = Shape<_128, _256, _32>; + using SGLayout = Layout, Stride<_8, _1, _0>>; + + using GmemTiledCopyD = XE_STORE_2D<16, 8, 32>; +}; + +class dpas_w4a16_policy_m_8 : public dpas_policy_base { + public: + using WGTile = Shape<_8, _64, _32>; + using SGLayout = Layout, Stride<_4, _1, _0>>; +}; + // --------------------------------------------------------------------------- // `apply_scale` -- inline-asm per-lane multiply of a bf16/fp16 fragment by // an FP32 scalar. Copied verbatim from vllm-xpu-kernels `gemm_xe2.hpp`. diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index fad8d87eb7..693a8675e5 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -129,6 +129,15 @@ using ::ark::moe_dpas_fp8::dpas_w16a16_policy; using ::ark::moe_dpas_fp8::dpas_w8a16_policy; using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; using ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; +// Dedicated 4-bit tile policies. The default (large-M) tile is 128x256x32 +// (halved packed-B stream lets a wider N tile pay off) and the m_8 bucket is +// new; both mirror the reference `w4a16_policy*` in vllm-xpu-kernels. The +// m_16 / m_32 buckets share the INT8 64-wide N tiles, so they are aliased to +// the existing `dpas_w8a16_policy_m_16 / _m_32` shapes rather than duplicated. +using ::ark::moe_dpas_fp8::dpas_w4a16_policy; +using ::ark::moe_dpas_fp8::dpas_w4a16_policy_m_8; +using dpas_w4a16_policy_m_16 = ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_16; +using dpas_w4a16_policy_m_32 = ::ark::moe_dpas_fp8::dpas_w8a16_policy_m_32; using ::ark::moe_dpas_fp8::ScaleMode; using ::ark::moe_dpas_fp8::cute_scalar; using ::ark::moe_dpas_fp8::cute_scalar_t; @@ -649,12 +658,20 @@ void moe_prefill_s4_dpas_per_group_dispatch( static_cast(nullptr), outputs_ca, N, K, \ num_tokens_per_expert, E, group_size, atomic_buffer); - if (A_avg_M <= 8) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (A_avg_M <= 32) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); + // Four-tier bucketing on the average tokens-per-expert, matching the + // reference `w4a16` dispatch in vllm-xpu-kernels + // (`grouped_gemm_xe2_interface.hpp`): tiny M uses an 8-row tile, and the + // 32-row tile now covers M up to 128 (instead of jumping to the wide + // large-M tile at 33) so mid-size chunked-prefill batches avoid the + // padding waste of the 128-row tile. + if (A_avg_M <= 4) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_8); + } else if (A_avg_M <= 8) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_16); + } else if (A_avg_M <= 128) { + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy_m_32); } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy); } #undef ARK_DPAS_S4_PG_LAUNCH_SYM From 93cde8c02bb7b46dcb9ca35afc7406ecd7f35605 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 3 Aug 2026 02:47:05 +0000 Subject: [PATCH 14/47] fix: restore fp8 per-expert MoE prefill large-M tile policy threshold Merge ec61621 accidentally widened the dpas_w8a16_policy_m_32 bucket in the fp8 per-tensor (per-expert) prefill dispatch from A_avg_M <= 32 to <= 512, routing large-M prefill through the small 32x64 tile instead of the large-M 128x128 default tile and regressing performance. Restore the <= 32 threshold. --- .../wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index 05bce99f76..b67694e80b 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -973,7 +973,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 <= 512) { + } else if (A_avg_M <= 32) { ARK_DPAS_PT_LAUNCH(dpas_w8a16_policy_m_32); } else { ARK_DPAS_PT_LAUNCH(dpas_w8a16_policy); From f8877638ec441eb8a28027196af26ff9e534b74e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 3 Aug 2026 05:13:59 +0000 Subject: [PATCH 15/47] fix: restore fp8 per-group MoE prefill large-M tile policy threshold Merge ec61621 also widened the dpas_w8a16_policy_m_32 bucket in the fp8 per-group prefill dispatch from A_avg_M <= 32 to <= 512, regressing large-M group-size prefill for the same reason as the per-tensor path. Restore the <= 32 threshold so large-M prefill uses the 128x128 default tile. --- .../wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index b67694e80b..b9898f81f0 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -1031,7 +1031,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 <= 512) { + } else if (A_avg_M <= 32) { ARK_DPAS_PG_LAUNCH(dpas_w8a16_policy_m_32); } else { ARK_DPAS_PG_LAUNCH(dpas_w8a16_policy); From c8374a86a77781cdafcccb98838e58c671f5c71b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 3 Aug 2026 05:35:36 +0000 Subject: [PATCH 16/47] revert: undo fp8 per-group MoE prefill threshold change Revert commit f887763, restoring the fp8 per-group prefill dispatch threshold to A_avg_M <= 512. The per-expert (per-tensor) fix from 93cde8c is retained. --- .../wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index b9898f81f0..b67694e80b 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -1031,7 +1031,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); From 3b3d7148229fdb10b199a369bce67e30bd4137d4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 3 Aug 2026 07:08:38 +0000 Subject: [PATCH 17/47] fix: switch ark.moe auto dispatch to total-token threshold Migrate the auto-dispatch logic from branch copilot/update-phase-auto-dispatch-logic (commit 9605fe4): phase="auto" now dispatches to decode when activations.shape[0] <= threshold (total tokens) instead of inspecting num_tokens_per_expert.max(), avoiding a host-device sync. Adds ARK_MOE_AUTO_DECODE_MAX_TOKENS env override (default 256) and updates test_moe_unified.py accordingly. --- .../ark/auto_round_kernel/__init__.py | 69 +++++---- .../ark/test/test_moe_unified.py | 138 +++++++++++++++--- 2 files changed, 158 insertions(+), 49 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index b66e180d82..4b5a5726ba 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -2331,8 +2331,9 @@ def _native_fp8_prefill_enabled() -> bool: # # `moe_gemm_decode` and `moe_gemm_prefill` accept identical argument shapes # and dtypes -- the only difference is which underlying SYCL kernel is -# launched (a GEMV variant tuned for 1-2 tokens/expert vs. a Grouped GEMM -# variant tuned for many tokens/expert). Model code that runs through both +# launched (a GEMV variant tuned for smaller total-token workloads vs. a +# Grouped GEMM variant tuned for larger total-token workloads). Model code +# that runs through both # regimes (prefill of a prompt, then autoregressive decode) traditionally # has to keep two call sites and branch on phase. `moe(...)` collapses that # into a single API and auto-selects the right kernel from the token @@ -2340,21 +2341,40 @@ def _native_fp8_prefill_enabled() -> bool: # # Callers that already know the phase (e.g., a model's generation loop knows # whether it's in prefill or decode) should pass it via the `phase` argument -# to avoid the small host-device sync that `phase="auto"` needs to inspect -# `num_tokens_per_expert.max()`. +# to bypass the auto-dispatch heuristic entirely. # --------------------------------------------------------------------------- -# Default tokens-per-expert threshold used by `phase="auto"`. The decode -# GEMV kernel is faster when every expert sees only a handful of tokens -# (TopK >= 1 with batch size 1-4); above that the GEMM-tuned prefill kernel -# wins. The crossover is hardware-dependent but `4` is a conservative default -# that matches the regime `moe_gemm_decode`'s docstring describes -# ("typically only 1-2 tokens", up to top-k * small batch). -_MOE_AUTO_DECODE_MAX_TOKENS_PER_EXPERT = 4 +# Default total-token threshold used by `phase="auto"`: dispatch to decode +# when `activations.shape[0] <= threshold`, otherwise prefill. This threshold +# is hardware-dependent and can be overridden via +# `ARK_MOE_AUTO_DECODE_MAX_TOKENS`. +_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS = 256 _MOE_VALID_PHASES = ("auto", "decode", "prefill") +def _moe_auto_decode_max_total_tokens() -> int: + """Return auto decode threshold from env or the module default. + + ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` is accepted when it is a positive + integer. Unset/empty/invalid values fall back to + ``_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS``. + """ + env = os.environ.get("ARK_MOE_AUTO_DECODE_MAX_TOKENS") + if env is None: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + env = env.strip() + if not env: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + try: + value = int(env) + except ValueError: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + if value <= 0: + return _MOE_AUTO_DECODE_MAX_TOTAL_TOKENS + return value + + def moe( activations: torch.Tensor, weights: torch.Tensor, @@ -2366,7 +2386,7 @@ def moe( group_size: int = 128, asym: bool = False, phase: str = "auto", - decode_threshold: int = _MOE_AUTO_DECODE_MAX_TOKENS_PER_EXPERT, + decode_threshold: Optional[int] = None, ) -> torch.Tensor: """Unified MoE GEMM entry point that dispatches to decode or prefill. @@ -2387,18 +2407,18 @@ def moe( underlying kernel; see :func:`moe_gemm_decode`. phase: dispatch mode. - * ``"auto"`` (default): inspect ``num_tokens_per_expert.max()`` - and pick decode if every expert sees ``<= decode_threshold`` - tokens, otherwise prefill. This incurs one small host-device - sync per call. + * ``"auto"`` (default): dispatch to decode when + ``activations.shape[0] <= decode_threshold`` (total tokens), + otherwise prefill. * ``"decode"``: always dispatch to :func:`moe_gemm_decode`. Use when the model's generation loop already knows it is in the - decode phase; avoids the sync. + decode phase. * ``"prefill"``: always dispatch to :func:`moe_gemm_prefill`. Use when the model knows it is in the prefill phase. - decode_threshold: ``"auto"`` mode dispatches to decode when - ``num_tokens_per_expert.max() <= decode_threshold``. Defaults to - 4 (the regime the decode GEMV kernel is tuned for). + decode_threshold: Total-token threshold for ``"auto"`` mode. If not + provided, uses ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` when set to a + valid positive integer, otherwise defaults to 256. Explicit + argument values take precedence over the environment variable. Returns: ``[total_tokens, N]`` in the activations dtype. Bit-identical to the @@ -2408,14 +2428,11 @@ def moe( raise ValueError(f"phase must be one of {_MOE_VALID_PHASES}, got {phase!r}") if phase == "auto": - # `.max().item()` triggers a host-device sync; callers in tight - # decode loops should pass `phase="decode"` explicitly to skip this. - # We tolerate a non-int32 / non-contiguous tensor here because the - # downstream kernel wrappers will normalise it anyway. + threshold = _moe_auto_decode_max_total_tokens() if decode_threshold is None else int(decode_threshold) if num_tokens_per_expert.numel() == 0: raise ValueError("num_tokens_per_expert must be non-empty") - max_tpe = int(num_tokens_per_expert.max().item()) - phase = "decode" if max_tpe <= int(decode_threshold) else "prefill" + total_tokens = int(activations.shape[0]) + phase = "decode" if total_tokens <= threshold else "prefill" if phase == "decode": return moe_gemm_decode( diff --git a/auto_round_extension/ark/test/test_moe_unified.py b/auto_round_extension/ark/test/test_moe_unified.py index 4e769456fc..98bea8057b 100644 --- a/auto_round_extension/ark/test/test_moe_unified.py +++ b/auto_round_extension/ark/test/test_moe_unified.py @@ -25,8 +25,8 @@ This file checks: - * Dispatch correctness: ``phase="auto"`` picks decode when every expert - sees few tokens and prefill otherwise. + * Dispatch correctness: ``phase="auto"`` picks decode when total tokens are + below threshold and prefill otherwise. * Bit-parity: ``moe(phase="auto")`` matches the kernel it dispatched to. * Explicit-phase parity: ``moe(phase="decode")`` matches ``moe_gemm_decode``, ``moe(phase="prefill")`` matches @@ -88,11 +88,12 @@ def _unified_skip_reason() -> str: # --------------------------------------------------------------------------- -# Small shapes (one decode-shaped, one prefill-shaped) -- keep wall-clock low. +# Small shapes (one decode-sized by total tokens, one prefill-sized by total +# tokens) -- keep wall-clock low. # --------------------------------------------------------------------------- -_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[1, 2, 0, 2], N=128, K=256) -_PREFILL_SHAPE = dict(num_experts=4, tokens_per_expert=[16, 8, 0, 20], N=128, K=256) +_AUTO_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[64, 64, 64, 64], N=128, K=256) # total_tokens=256 +_AUTO_PREFILL_SHAPE = dict(num_experts=4, tokens_per_expert=[80, 80, 80, 80], N=128, K=256) # total_tokens=320 def _make_int4_sym(E, N, K, group_size, dtype, total_tokens): @@ -163,8 +164,8 @@ def _make_fp8(E, N, K, group_size, dtype, total_tokens, fp8_dtype): class TestMoeUnifiedDispatch: """Tests for the auto-dispatch logic itself.""" - def test_auto_picks_decode_for_small_tokens_per_expert(self): - shape = _DECODE_SHAPE + def test_auto_picks_decode_for_small_total_tokens(self): + shape = _AUTO_DECODE_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -192,12 +193,12 @@ def test_auto_picks_decode_for_small_tokens_per_expert(self): group_size=group_size, asym=False, ) - # max tokens/expert = 2 (<= default threshold 4) -> dispatched to decode + # total tokens = 256 (<= default threshold 256) -> dispatched to decode # -> output must be bit-identical to moe_gemm_decode. torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) - def test_auto_picks_prefill_for_large_tokens_per_expert(self): - shape = _PREFILL_SHAPE + def test_auto_picks_prefill_for_large_total_tokens(self): + shape = _AUTO_PREFILL_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -228,9 +229,66 @@ def test_auto_picks_prefill_for_large_tokens_per_expert(self): torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) def test_decode_threshold_override(self): - # Same prefill-shaped input but bump the threshold above the max - # tokens/expert -> auto must now pick decode. - shape = _PREFILL_SHAPE + decode_shape = _AUTO_DECODE_SHAPE + decode_total_tokens = sum(decode_shape["tokens_per_expert"]) + E, N, K = decode_shape["num_experts"], decode_shape["N"], decode_shape["K"] + group_size = 128 + dtype = torch.float16 + + activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, decode_total_tokens) + ntpe = torch.tensor(decode_shape["tokens_per_expert"], dtype=torch.int32, device="xpu") + + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=128, + ) + out_prefill = ark.moe_gemm_prefill( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) + + prefill_shape = _AUTO_PREFILL_SHAPE + prefill_total_tokens = sum(prefill_shape["tokens_per_expert"]) + activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, prefill_total_tokens) + ntpe = torch.tensor(prefill_shape["tokens_per_expert"], dtype=torch.int32, device="xpu") + + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=prefill_total_tokens, + ) + out_decode = ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) + + def test_decode_threshold_env_override(self, monkeypatch): + shape = _AUTO_PREFILL_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -239,7 +297,7 @@ def test_decode_threshold_override(self): activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, total_tokens) ntpe = torch.tensor(shape["tokens_per_expert"], dtype=torch.int32, device="xpu") - max_tpe = max(shape["tokens_per_expert"]) + monkeypatch.setenv("ARK_MOE_AUTO_DECODE_MAX_TOKENS", "512") out_auto = ark.moe( activations, packed, @@ -249,7 +307,6 @@ def test_decode_threshold_override(self): group_size=group_size, asym=False, phase="auto", - decode_threshold=max_tpe + 1, ) out_decode = ark.moe_gemm_decode( activations, @@ -262,8 +319,43 @@ def test_decode_threshold_override(self): ) torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=128, + ) + out_prefill = ark.moe_gemm_prefill( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) + + monkeypatch.setenv("ARK_MOE_AUTO_DECODE_MAX_TOKENS", "invalid") + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + ) + torch.testing.assert_close(out_auto, out_prefill, rtol=0, atol=0) + def test_invalid_phase_raises(self): - shape = _DECODE_SHAPE + shape = _AUTO_DECODE_SHAPE total_tokens = sum(shape["tokens_per_expert"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -298,8 +390,8 @@ class TestMoeUnifiedBitParity: @pytest.mark.parametrize( "shape_name,shape", [ - ("decode-shape", _DECODE_SHAPE), - ("prefill-shape", _PREFILL_SHAPE), + ("decode-shape", _AUTO_DECODE_SHAPE), + ("prefill-shape", _AUTO_PREFILL_SHAPE), ], ) def test_fp_unquantized(self, dtype, shape_name, shape): @@ -324,8 +416,8 @@ def test_fp_unquantized(self, dtype, shape_name, shape): @pytest.mark.parametrize( "shape_name,shape", [ - ("decode-shape", _DECODE_SHAPE), - ("prefill-shape", _PREFILL_SHAPE), + ("decode-shape", _AUTO_DECODE_SHAPE), + ("prefill-shape", _AUTO_PREFILL_SHAPE), ], ) def test_int4(self, dtype, asym, shape_name, shape): @@ -353,7 +445,7 @@ def test_int4(self, dtype, asym, shape_name, shape): def test_int8(self, dtype, asym): # Single shape -- the quant path is the same on both shapes, so # iterating both would just slow the test suite down. - shape = _PREFILL_SHAPE + shape = _AUTO_PREFILL_SHAPE E, N, K = shape["num_experts"], shape["N"], shape["K"] total_tokens = sum(shape["tokens_per_expert"]) group_size = 128 @@ -375,7 +467,7 @@ def test_int8(self, dtype, asym): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("asym", [False, True]) def test_int2(self, dtype, asym): - shape = _PREFILL_SHAPE + shape = _AUTO_PREFILL_SHAPE E, N, K = shape["num_experts"], shape["N"], shape["K"] total_tokens = sum(shape["tokens_per_expert"]) group_size = 128 @@ -397,7 +489,7 @@ def test_int2(self, dtype, asym): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) def test_fp8(self, dtype, fp8_dtype): - shape = _PREFILL_SHAPE + shape = _AUTO_PREFILL_SHAPE E, N, K = shape["num_experts"], shape["N"], shape["K"] total_tokens = sum(shape["tokens_per_expert"]) group_size = 128 From 6c3d26f28c07260b26fde9ad8edb403dec20a824 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 4 Aug 2026 07:38:06 +0000 Subject: [PATCH 18/47] feat: route int4-sym MoE decode to shared S4 DPAS grouped GEMM Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 60 +++++++++++++++++-- auto_round_extension/ark/test/test_moe.py | 49 +++++++++++++++ 2 files changed, 105 insertions(+), 4 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 0f03604ccf..9ec3e0a64a 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -35,10 +35,17 @@ #include #include +#include #include +#include #include "bestla/bestla.h" #include "sycl_tla_moe_dequant.hpp" +// S4-sym per-group DPAS grouped-GEMM (shared with the prefill path). The +// header self-guards on `ARK_XPU && ARK_SYCL_TLA`, so including it here is a +// no-op when the DPAS backend is disabled. Decode routes small-M int4-sym +// GEMV through this kernel; see `moe_gemm_decode` below. +#include "sycl_tla_moe_prefill_s4_dpas.hpp" #ifdef ARK_XPU #include @@ -740,7 +747,30 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh } // namespace moe_decode_detail // ---------------------------------------------------------------------------- -// Public API +// Env-flag helper -- `ARK_MOE_DECODE_DPAS_S4` (default ON). When ON, int4-sym +// (S4_CLIP, !asym) decode is routed to the shared per-group S4 DPAS grouped +// GEMM (`moe_dpas_s4::moe_prefill_s4_dpas_per_group_dispatch`) instead of the +// scalar FMA GEMV (`launch_int4`). The DPAS path already handles the tiny +// total-token counts typical of decode (its `A_avg_M <= 4` bucket selects the +// 8-row `dpas_w4a16_policy_m_8` tile) and reads the same `[E, N, K/2]` packed +// weights + `[E, N, K/group]` scales, so no repack is needed. +// +// Setting the var to "0" / "false" / "off" / "no" (case-insensitive) forces +// the legacy scalar GEMV, for A/B comparison and regression escape. Asym +// weights and shapes that fail the DPAS shape gate always fall back to the +// scalar path regardless of this flag. Re-read on every call so tests / +// benchmarks can toggle the path in-process. +// ---------------------------------------------------------------------------- +inline bool moe_decode_dpas_s4_enabled() { + const char* env = std::getenv("ARK_MOE_DECODE_DPAS_S4"); + if (env == nullptr) return true; // default ON + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "0" || s == "false" || s == "off" || s == "no") return false; + return true; +} + +// ---------------------------------------------------------------------------- // // weight_dtype: // BTLA_DTYPE::F16 / BF16 : weights stored as [E, N, K] in matching @@ -783,6 +813,30 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo } if (weight_dtype == BTLA_DTYPE::S4_CLIP) { + if (act_dtype != BTLA_DTYPE::F16 && act_dtype != BTLA_DTYPE::BF16) { + throw std::invalid_argument("moe_gemm_decode(int4): act_dtype must be FP16 or BF16"); + } + // Fast path: sym int4 through the shared per-group S4 DPAS grouped GEMM. + // Falls back to the scalar GEMV for asym weights (DPAS S4 is sym-only), + // when the env flag is off, or when the shape gate rejects the tile + // geometry (e.g. N%64!=0, K%32!=0, unsupported group_size). + if (!asym && moe_decode_dpas_s4_enabled() && + moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size)) { + if (act_dtype == BTLA_DTYPE::F16) { + moe_dpas_s4::moe_prefill_s4_dpas_per_group_dispatch( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(outputs), num_tokens_per_expert, + num_experts, N, K, group_size, total_tokens); + } else { + using BF = sycl::ext::oneapi::bfloat16; + moe_dpas_s4::moe_prefill_s4_dpas_per_group_dispatch( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(outputs), num_tokens_per_expert, num_experts, N, K, + group_size, total_tokens); + } + return; + } + // Scalar FMA GEMV fallback (asym, flag off, or shape gate miss). if (act_dtype == BTLA_DTYPE::F16) { if (asym) { moe_decode_detail::launch_int4( @@ -795,7 +849,7 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo static_cast(scales), static_cast(zeros), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); } - } else if (act_dtype == BTLA_DTYPE::BF16) { + } else { using BF = sycl::ext::oneapi::bfloat16; if (asym) { moe_decode_detail::launch_int4( @@ -808,8 +862,6 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo static_cast(scales), static_cast(zeros), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); } - } else { - throw std::invalid_argument("moe_gemm_decode(int4): act_dtype must be FP16 or BF16"); } return; } diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 8bad20cf98..ebeef5da52 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -480,6 +480,55 @@ def test_decode_int4_sym(self, dtype, group_size): assert out.shape == (total_tokens, N) torch.testing.assert_close(out, ref, rtol=5e-2, atol=5e-2) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("group_size", [32, 128]) + @pytest.mark.parametrize("tokens_per_expert", [[1, 0, 1, 1], [1, 1, 1, 1], [2, 0, 0, 1]]) + def test_decode_int4_sym_dpas_matches_scalar(self, monkeypatch, dtype, group_size, tokens_per_expert): + """int4-sym decode: the S4 DPAS path (ARK_MOE_DECODE_DPAS_S4=1, default) + must match both the scalar GEMV fallback (ARK_MOE_DECODE_DPAS_S4=0) and + the dequant->bmm reference within quantization tolerance. + + Shapes satisfy the DPAS per-group shape gate (N%64==0, K%32==0, + group_size in {32,64,128,256}) so the fast path is actually taken. + """ + num_experts = 4 + total_tokens = sum(tokens_per_expert) + N, K = 320, 256 # N%64==0, K%32==0 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_sym(w_float, scales, group_size) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + dequant = _dequant_int4_sym(packed, scales, group_size).to(dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + out_dpas = _run() + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + out_scalar = _run() + + assert out_dpas.shape == (total_tokens, N) + assert out_scalar.shape == (total_tokens, N) + # Both kernels approximate the same dequant reference. + torch.testing.assert_close(out_dpas, ref, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(out_scalar, ref, rtol=5e-2, atol=5e-2) + # And they must agree with each other within the same tolerance. + torch.testing.assert_close(out_dpas, out_scalar, rtol=5e-2, atol=5e-2) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_decode_int4_asym(self, dtype): num_experts = 4 From 0c946597e3012e855e011e71b3a9eefac49ef206 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 4 Aug 2026 07:39:11 +0000 Subject: [PATCH 19/47] test: add int4-sym decode DPAS vs scalar perf comparison; docs Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 7 +++- .../ark/test/test_moe_decode_perf.py | 41 ++++++++++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 9ec3e0a64a..003c924dff 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -780,7 +780,12 @@ inline bool moe_decode_dpas_s4_enabled() { // unsigned with zero-points when asym==true) // BTLA_DTYPE::S4_CLIP : packed int4 weights [E, N, K/2] (uint8), // scales [E, N, K/group_size] in act dtype, -// zeros optional (asym==true requires it) +// zeros optional (asym==true requires it). +// Sym weights are routed to the shared +// per-group S4 DPAS grouped GEMM by default +// (`ARK_MOE_DECODE_DPAS_S4`, default ON); +// asym, a disabled flag, or a shape-gate +// miss falls back to the scalar GEMV. // BTLA_DTYPE::S2_CLIP : packed int2 weights [E, N, K/4] (uint8), // 4 values per byte, sym/asym like int4 // BTLA_DTYPE::F8_E4M3 / F8_E5M2 : FP8 weights [E, N, K] (uint8 buffer), diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index cf9b9a74cb..94017e5c8b 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -370,8 +370,45 @@ def test_perf_int4(self, dtype, asym): _print_row(label, N, K, total_tokens, base_ms, ark_ms) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) - @pytest.mark.parametrize("asym", [False, True]) - def test_perf_int8(self, dtype, asym): + def test_perf_int4_sym_dpas_vs_scalar(self, monkeypatch, dtype): + """int4-sym decode: compare the S4 DPAS path (ARK_MOE_DECODE_DPAS_S4=1) + against the scalar GEMV fallback (ARK_MOE_DECODE_DPAS_S4=0). + + ``speedup`` here is ``scalar / dpas`` (the DPAS path is the "ark" + column), isolating the DPAS routing win from the dequant reference. + Only shapes that clear the DPAS shape gate are timed on both paths. + """ + group_size = 128 + _print_header( + f"INT4 sym DPAS vs scalar (group_size={group_size}, " + f"act={str(dtype).split('.')[-1]}) -- scalar GEMV (baseline) vs S4 DPAS (ark)" + ) + for label, E, tpe, N, K in DECODE_SHAPES: + if K % group_size != 0 or N % 64 != 0 or K % 32 != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_sym(w_float, scales, group_size) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + scalar_ms = _xpu_time_ms(_run) + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + dpas_ms = _xpu_time_ms(_run) + _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) group_size = 128 kind = "asym" if asym else "sym" _print_header( From ffd2926c6136a38bdb4961049784d1d516e605ea Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 4 Aug 2026 08:52:24 +0000 Subject: [PATCH 20/47] fix: split orphaned int8 block out of int4-sym dpas perf test Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- auto_round_extension/ark/test/test_moe_decode_perf.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index 94017e5c8b..ccbc08edc3 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -409,6 +409,10 @@ def _run(): monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") dpas_ms = _xpu_time_ms(_run) _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("asym", [False, True]) + def test_perf_int8(self, dtype, asym): group_size = 128 kind = "asym" if asym else "sym" _print_header( From 87b617ea4c90b51910d8afe810b737ea3fe1c467 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 4 Aug 2026 09:15:24 +0000 Subject: [PATCH 21/47] fix: lower default MoE auto-decode threshold 256 -> 32 The shared S4 DPAS grouped-GEMM (prefill path) already beats the scalar GEMV decode kernel by ~2x at 256 tokens (bs32) and only loses at the single-stream bs1 (8-token) extreme. Routing 256-token batches to decode was leaving ~2x on the table, so lower the auto-dispatch default and update coupled unified-dispatch tests and perf-test notes. Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/auto_round_kernel/__init__.py | 13 +++++++++++-- .../ark/test/test_moe_decode_perf.py | 6 ++++++ auto_round_extension/ark/test/test_moe_unified.py | 7 ++++--- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 4b5a5726ba..69147abc2f 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -2348,7 +2348,16 @@ def _native_fp8_prefill_enabled() -> bool: # when `activations.shape[0] <= threshold`, otherwise prefill. This threshold # is hardware-dependent and can be overridden via # `ARK_MOE_AUTO_DECODE_MAX_TOKENS`. -_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS = 256 +# +# Empirically (see `test_perf_int4_sym_dpas_vs_scalar`) the shared S4 DPAS +# grouped-GEMM that the prefill path uses already beats the scalar-GEMV decode +# kernel by ~2x once total tokens reach 256 (bs32), while it loses (~0.5x) only +# at the single-stream bs1 extreme (8 tokens). The crossover therefore sits far +# below 256, so the default is kept small: the scalar decode kernel is only +# preferred for the tiny single-/few-stream case where every expert sees well +# under one DPAS tile row. Mirrors vLLM-xpu-kernels' `w4a16` dispatch, which +# buckets on average tokens-per-expert rather than a large total-token cutoff. +_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS = 32 _MOE_VALID_PHASES = ("auto", "decode", "prefill") @@ -2417,7 +2426,7 @@ def moe( Use when the model knows it is in the prefill phase. decode_threshold: Total-token threshold for ``"auto"`` mode. If not provided, uses ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` when set to a - valid positive integer, otherwise defaults to 256. Explicit + valid positive integer, otherwise defaults to 32. Explicit argument values take precedence over the environment variable. Returns: diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index ccbc08edc3..59158a3a8e 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -377,6 +377,12 @@ def test_perf_int4_sym_dpas_vs_scalar(self, monkeypatch, dtype): ``speedup`` here is ``scalar / dpas`` (the DPAS path is the "ark" column), isolating the DPAS routing win from the dequant reference. Only shapes that clear the DPAS shape gate are timed on both paths. + + Observed crossover: at bs1 (8 tokens, ~1 token/expert) the DPAS tile is + nearly empty and the scalar GEMV wins (~0.5x), but by bs32 (256 tokens) + the shared S4 DPAS grouped-GEMM already wins ~2x. The auto-dispatch + default threshold (`_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS`) is therefore kept + small so 256-token batches are routed to prefill/DPAS, not decode. """ group_size = 128 _print_header( diff --git a/auto_round_extension/ark/test/test_moe_unified.py b/auto_round_extension/ark/test/test_moe_unified.py index 98bea8057b..2d5eb232fe 100644 --- a/auto_round_extension/ark/test/test_moe_unified.py +++ b/auto_round_extension/ark/test/test_moe_unified.py @@ -92,7 +92,7 @@ def _unified_skip_reason() -> str: # tokens) -- keep wall-clock low. # --------------------------------------------------------------------------- -_AUTO_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[64, 64, 64, 64], N=128, K=256) # total_tokens=256 +_AUTO_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[4, 4, 4, 4], N=128, K=256) # total_tokens=16 _AUTO_PREFILL_SHAPE = dict(num_experts=4, tokens_per_expert=[80, 80, 80, 80], N=128, K=256) # total_tokens=320 @@ -193,7 +193,7 @@ def test_auto_picks_decode_for_small_total_tokens(self): group_size=group_size, asym=False, ) - # total tokens = 256 (<= default threshold 256) -> dispatched to decode + # total tokens = 16 (<= default threshold 32) -> dispatched to decode # -> output must be bit-identical to moe_gemm_decode. torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) @@ -238,6 +238,7 @@ def test_decode_threshold_override(self): activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, decode_total_tokens) ntpe = torch.tensor(decode_shape["tokens_per_expert"], dtype=torch.int32, device="xpu") + # A threshold strictly below the shape's total tokens forces prefill. out_auto = ark.moe( activations, packed, @@ -247,7 +248,7 @@ def test_decode_threshold_override(self): group_size=group_size, asym=False, phase="auto", - decode_threshold=128, + decode_threshold=decode_total_tokens - 1, ) out_prefill = ark.moe_gemm_prefill( activations, From 23ba0dbfca17c23d59e68238ad2df2ebd9cf9c96 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 01:43:44 +0000 Subject: [PATCH 22/47] test: add int4-sym DPAS vs scalar token-count sweep for MoE decode threshold tuning Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/test/test_moe_decode_perf.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index 59158a3a8e..584287528d 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -214,6 +214,32 @@ def _default_moe_decode(activations, dequant_weights, num_tokens_per_expert): # Backwards-compatible alias (older code/tests referenced ``_MINIMAX_TPE``). _MINIMAX_TPE = _MINIMAX_TPE_BS1 + +def _spread_tokens(total_tokens: int, num_experts: int = 192) -> list: + """Distribute ``total_tokens`` across ``num_experts`` round-robin. + + Returns a ``[num_experts]`` histogram summing to ``total_tokens`` where the + load is striped across the expert range (expert ``i`` gets a token before + ``i+1`` gets its second), mirroring the spread a real top-k router produces + rather than clustering all tokens onto the first few experts. Used by the + threshold-sweep test to synthesise decode workloads of an exact size. + """ + tpe = [0] * num_experts + for i in range(total_tokens): + tpe[i % num_experts] += 1 + return tpe + + +# Total-token counts swept by ``test_perf_int4_sym_dpas_vs_scalar_threshold`` +# to locate the DPAS-vs-scalar crossover for the auto-dispatch threshold. +_INT4_THRESHOLD_TOKEN_COUNTS = [16, 32, 64, 128] + +# MiniMax-M2 up/down-proj (N, K) pairs reused by the threshold sweep. +_INT4_THRESHOLD_NK = [ + (1536, 3072), # gate/up-proj + (3072, 1536), # down-proj +] + DECODE_SHAPES = [ # (label, num_experts, tokens_per_expert, N, K) # batch=1 decode (single-stream). @@ -416,6 +442,55 @@ def _run(): dpas_ms = _xpu_time_ms(_run) _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_perf_int4_sym_dpas_vs_scalar_threshold(self, monkeypatch, dtype): + """int4-sym decode threshold sweep: DPAS vs scalar GEMV across a range + of total-token counts (16/32/64/128) at ``group_size=32``. + + Same comparison as ``test_perf_int4_sym_dpas_vs_scalar`` (``speedup`` is + ``scalar / dpas``, DPAS is the "ark" column) but instead of the fixed + MiniMax bs1/bs32 shapes it synthesises decode workloads of an exact + total-token size via :func:`_spread_tokens`. The ``speedup`` column + crosses 1.0x at the total-token count where the shared S4 DPAS + grouped-GEMM starts beating the scalar GEMV, which is the value to feed + into ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` / the ``moe(...)`` + ``decode_threshold`` auto-dispatch cutoff. + """ + group_size = 32 + E = 192 + _print_header( + f"INT4 sym DPAS vs scalar threshold sweep (group_size={group_size}, " + f"act={str(dtype).split('.')[-1]}) -- scalar GEMV (baseline) vs S4 DPAS (ark)" + ) + for N, K in _INT4_THRESHOLD_NK: + if K % group_size != 0 or N % 64 != 0 or K % 32 != 0: + continue + for total_tokens in _INT4_THRESHOLD_TOKEN_COUNTS: + tpe = _spread_tokens(total_tokens, E) + label = f"int4 {N}x{K} t{total_tokens}" + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_sym(w_float, scales, group_size) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + scalar_ms = _xpu_time_ms(_run) + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + dpas_ms = _xpu_time_ms(_run) + _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("asym", [False, True]) def test_perf_int8(self, dtype, asym): From a1480af3c661a706cec653488c45bb650177f275 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 05:33:38 +0000 Subject: [PATCH 23/47] perf: hoist per-group scale and split accumulators in int4 MoE decode GEMV Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 77 ++++++++++++------- 1 file changed, 51 insertions(+), 26 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 003c924dff..6a00ea4c75 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -218,10 +218,23 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight // stage for the remainder, which keeps the fast path active for group sizes // that are a multiple of 32 (32/64/128/256 -- the shipped quant configs) // without regressing group_size == 16 (which drops straight to the 16-wide -// stage). The math is identical to the scalar path. +// stage). +// +// The per-group scale (and asym zero-point) is NOT applied here: this +// accumulates the raw integer-weighted dot product ``Σ a·q`` into ``acc_q`` +// and, for the asym case, the plain activation sum ``Σ a`` into ``acc_a``. +// The caller folds the group's scale/zero in once (Σ a·((q−z)·s) == +// s·(Σ a·q − z·Σ a); sym collapses to s·Σ a·q). Hoisting the scale removes one +// float multiply per K element on the decode hot path, and because the scale +// fold is exact-once per group the result stays well within the kernel's +// existing quantization tolerance. +// +// Two independent partial accumulators (``acc_q0``/``acc_q1``) break the +// single fp32 dependency chain so the FMA pipeline is not latency-bound; the +// caller reduces the pair. ``acc_a`` (asym only) reuses the same split. template -static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float scale, float zero, - float& acc) { +static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc_q0, float& acc_q1, + float& acc_a) { static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); // sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so a single @@ -239,18 +252,15 @@ static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_pt for (int b = 0; b < SUB / 2; ++b) { int q0, q1; decode_int4_pair(pv[b], q0, q1); - float w0, w1; - if constexpr (Asym) { - w0 = (static_cast(q0) - zero) * scale; - w1 = (static_cast(q1) - zero) * scale; - } else { - w0 = static_cast(q0) * scale; - w1 = static_cast(q1) * scale; - } const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); - acc += static_cast(a0) * w0; - acc += static_cast(a1) * w1; + const float fa0 = static_cast(a0); + const float fa1 = static_cast(a1); + acc_q0 += fa0 * static_cast(q0); + acc_q1 += fa1 * static_cast(q1); + if constexpr (Asym) { + acc_a += fa0 + fa1; + } } } } @@ -311,35 +321,50 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // across twice as many multiply-adds for the shipped group sizes // (32/64/128/256), while the 16-wide stage keeps group_size == 16 // on the fast path. + // + // The scale (and asym zero) is constant across the group, so the + // wide stages accumulate the raw integer-weighted dot product + // ``Σ a·q`` (split across two partial accumulators to break the + // fp32 dependency chain) plus ``Σ a`` for asym, and the fold below + // applies the scale/zero exactly once per group: + // sym : acc += scale * (acc_q0 + acc_q1) + // asym: acc += scale * ((acc_q0 + acc_q1) - zero * acc_a) + float acc_q0 = 0.0f; + float acc_q1 = 0.0f; + float acc_a = 0.0f; int kk = 0; constexpr int CHUNK32 = 32; const int end32 = (group_size / CHUNK32) * CHUNK32; for (; kk < end32; kk += CHUNK32) { - int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, - zero, acc); + int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, acc_q0, + acc_q1, acc_a); } constexpr int CHUNK16 = 16; const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; for (; kk < end16; kk += CHUNK16) { - int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, scale, - zero, acc); + int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, acc_q0, + acc_q1, acc_a); } - // Scalar tail for group_size not divisible by 16. + // Scalar tail for group_size not divisible by 16. Uses the same + // raw-accumulation convention as the wide stages so the single + // scale/zero fold below stays valid. for (; kk < group_size; kk += 2) { const uint8_t packed = w_row[(k_base + kk) / 2]; int q0, q1; decode_int4_pair(packed, q0, q1); - float w0, w1; + const float fa0 = static_cast(act_row[k_base + kk]); + const float fa1 = static_cast(act_row[k_base + kk + 1]); + acc_q0 += fa0 * static_cast(q0); + acc_q1 += fa1 * static_cast(q1); if constexpr (Asym) { - w0 = (static_cast(q0) - zero) * scale; - w1 = (static_cast(q1) - zero) * scale; - } else { - w0 = static_cast(q0) * scale; - w1 = static_cast(q1) * scale; + acc_a += fa0 + fa1; } - acc += static_cast(act_row[k_base + kk]) * w0; - acc += static_cast(act_row[k_base + kk + 1]) * w1; } + float group_dot = acc_q0 + acc_q1; + if constexpr (Asym) { + group_dot -= zero * acc_a; + } + acc += scale * group_dot; } outputs[static_cast(token) * N + n_global] = static_cast(acc); From c20c0b2c7c3f6e7e206d0c08a033c31dbb33f293 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:59:22 +0000 Subject: [PATCH 24/47] perf: coalesce int4 MoE decode weight loads via on-device N-tiled repack Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 237 ++++++++++++++++-- auto_round_extension/ark/test/test_moe.py | 62 +++++ 2 files changed, 281 insertions(+), 18 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 6a00ea4c75..de9df4cdf8 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -11,7 +11,11 @@ // int8 per byte (sym: signed -128..127; // asym: unsigned 0..255 with zero-point) // - weights (int4 packed): [num_experts, N, K/2] row-major, two -// 4-bit values per byte (low nibble at lower K) +// 4-bit values per byte (low nibble at lower K). +// The scalar-GEMV fallback repacks this on-device +// into an N-tiled [E, N/16, K/2, 16] layout so that +// sub-group weight loads are coalesced; the external +// [E, N, K/2] contract is unchanged. // - weights (int2 packed): [num_experts, N, K/4] row-major, four // 2-bit values per byte (field j at K index // 4*i+j is bits [2j+1:2j]) @@ -87,6 +91,11 @@ class MoEDecodeKernelFP; template class MoEDecodeKernelInt4; +template +class MoEDecodeKernelInt4Coalesced; + +class MoEDecodeRepackInt4; + template class MoEDecodeKernelInt8; @@ -371,6 +380,146 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig }); } +// ---------------------------------------------------------------------------- +// INT4 (S4_CLIP) coalesced-load GEMV. +// +// The scalar `launch_int4` above is memory-bandwidth-bound: for a single decode +// token it just streams the whole packed weight matrix once with ~1 MAC per +// byte, so the arithmetic tweaks (split accumulators, hoisted scale) cannot +// help. Its real cost is that weight loads are *not coalesced across the +// sub-group*: with the `[E, N, K/2]` (K-contiguous) layout, lane `l` and lane +// `l+1` of a sub-group read packed bytes `K/2` apart at a fixed `k`, so each +// step issues 16 scattered transactions instead of one contiguous cache line. +// +// This path fixes that by first repacking the weights on-device into an +// N-tiled layout `[E, N/16, K/2, 16]`: the trailing dim of 16 holds one packed +// byte for each of the 16 columns owned by a sub-group tile, so at a fixed +// packed-byte index the 16 lanes read 16 contiguous bytes -> a single coalesced +// load. The dequant math is byte-for-byte identical to `launch_int4` (same +// `decode_int4_pair`, same per-group scale/zero fold), only the weight memory +// access pattern changes. The repack buffer is a transient USM device +// allocation freed after the queue drains; the caller's `[E, N, K/2]` weight +// contract is unchanged. +// +// The trailing lane stride means each lane's own K-bytes are 16 apart, so the +// vectorized `int4_decode_chunk` (contiguous per-lane load) does not apply +// here; the inner loop reads one packed byte per lane per step, which the +// hardware coalesces across the sub-group into one wide transaction. +// ---------------------------------------------------------------------------- +template +void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, + const ScalarT* scales, const ScalarT* zeros, ScalarT* outputs, + const int* expert_id_per_token, int total_tokens, int N, int K, int group_size, + int num_experts) { + if (N % N_TILE != 0) { + throw std::invalid_argument("moe_gemm_decode(int4): N must be a multiple of 16"); + } + if (K % group_size != 0 || (group_size & 1) != 0) { + throw std::invalid_argument("moe_gemm_decode(int4): K must be a multiple of group_size and group_size must be even"); + } + if (Asym && zeros == nullptr) { + throw std::invalid_argument("moe_gemm_decode(int4): zeros pointer required when asym=true"); + } + if (total_tokens == 0) return; + + const int n_tiles = N / N_TILE; + const int num_groups_k = K / group_size; + const int k_packed = K / 2; // bytes of packed weight per (expert, n) + + const size_t repacked_bytes = + static_cast(num_experts) * static_cast(n_tiles) * + static_cast(k_packed) * static_cast(N_TILE); + uint8_t* repacked = sycl::malloc_device(repacked_bytes, *q); + if (repacked == nullptr) { + throw std::runtime_error("moe_gemm_decode(int4): failed to allocate repack buffer"); + } + + // Repack kernel: one work-item per (expert, column, packed byte). The write + // index places the 16 columns of a tile contiguously in the trailing dim. + { + sycl::range<3> rp_global{static_cast(num_experts), static_cast(N), + static_cast(k_packed)}; + q->parallel_for(rp_global, [=](sycl::id<3> id) { + const int e = static_cast(id[0]); + const int n = static_cast(id[1]); + const int kb = static_cast(id[2]); + const int t = n / N_TILE; + const int l = n % N_TILE; + const size_t src = (static_cast(e) * N + static_cast(n)) * k_packed + kb; + const size_t dst = + ((static_cast(e) * n_tiles + t) * k_packed + kb) * N_TILE + l; + repacked[dst] = weights[src]; + }); + } + + sycl::range<2> global{static_cast(total_tokens), static_cast(n_tiles * SG_SIZE)}; + sycl::range<2> local{1, static_cast(SG_SIZE)}; + + q->parallel_for>( + sycl::nd_range<2>(global, local), + [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { + const int token = static_cast(it.get_global_id(0)); + const int n_tile = static_cast(it.get_group(1)); + const int lane = static_cast(it.get_local_id(1)); + const int n_global = n_tile * N_TILE + lane; + + const int expert = expert_id_per_token[token]; + const ScalarT* act_row = activations + static_cast(token) * K; + + // Base of this (expert, n_tile) weight tile in the repacked buffer. + // Layout [E, N/16, K/2, 16]; this lane reads byte kb at + // w_tile[kb*16 + lane], so adjacent lanes read adjacent bytes. + const uint8_t* w_tile = + repacked + ((static_cast(expert) * n_tiles + n_tile) * k_packed) * N_TILE; + const ScalarT* s_row = + scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* z_row = Asym + ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k + : nullptr; + + float acc = 0.0f; + for (int g = 0; g < num_groups_k; ++g) { + const float scale = static_cast(s_row[g]); + float zero = 0.0f; + if constexpr (Asym) { + zero = static_cast(z_row[g]); + } + const int k_base = g * group_size; + // Two partial accumulators break the fp32 dependency chain; the + // per-group scale/zero is folded once after the K-loop, exactly as + // in the scalar path. Each iteration processes two K-elements (one + // packed byte); the byte load is coalesced across the sub-group. + float acc_q0 = 0.0f; + float acc_q1 = 0.0f; + float acc_a = 0.0f; + const int kb_base = k_base / 2; + const int kb_count = group_size / 2; + for (int kb = 0; kb < kb_count; ++kb) { + const uint8_t packed = w_tile[(kb_base + kb) * N_TILE + lane]; + int q0, q1; + decode_int4_pair(packed, q0, q1); + const float fa0 = static_cast(act_row[k_base + 2 * kb]); + const float fa1 = static_cast(act_row[k_base + 2 * kb + 1]); + acc_q0 += fa0 * static_cast(q0); + acc_q1 += fa1 * static_cast(q1); + if constexpr (Asym) { + acc_a += fa0 + fa1; + } + } + float group_dot = acc_q0 + acc_q1; + if constexpr (Asym) { + group_dot -= zero * acc_a; + } + acc += scale * group_dot; + } + + outputs[static_cast(token) * N + n_global] = static_cast(acc); + }); + + q->wait(); + sycl::free(repacked, *q); +} + // ---------------------------------------------------------------------------- // INT8 (S8) GEMV with group-wise dequantization. // @@ -795,6 +944,24 @@ inline bool moe_decode_dpas_s4_enabled() { return true; } +// ---------------------------------------------------------------------------- +// Env-flag helper -- `ARK_MOE_DECODE_COALESCE_INT4` (default ON). When ON, the +// int4 scalar-GEMV fallback (asym, or sym with the DPAS path disabled / shape +// gate miss) uses `launch_int4_coalesced`, which repacks the weights on-device +// into an N-tiled layout so sub-group weight loads are coalesced. Setting the +// var to "0" / "false" / "off" / "no" (case-insensitive) forces the legacy +// per-lane-strided `launch_int4`, for A/B comparison and regression escape. +// Re-read on every call so tests / benchmarks can toggle it in-process. +// ---------------------------------------------------------------------------- +inline bool moe_decode_coalesce_int4_enabled() { + const char* env = std::getenv("ARK_MOE_DECODE_COALESCE_INT4"); + if (env == nullptr) return true; // default ON + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "0" || s == "false" || s == "off" || s == "no") return false; + return true; +} + // ---------------------------------------------------------------------------- // // weight_dtype: @@ -866,31 +1033,65 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo } return; } - // Scalar FMA GEMV fallback (asym, flag off, or shape gate miss). + // Scalar FMA GEMV fallback (asym, flag off, or shape gate miss). By + // default this uses the coalesced-load variant, which repacks the weights + // on-device so sub-group loads are contiguous; `ARK_MOE_DECODE_COALESCE_INT4=0` + // forces the legacy per-lane-strided kernel. + const bool coalesce = moe_decode_coalesce_int4_enabled(); if (act_dtype == BTLA_DTYPE::F16) { if (asym) { - moe_decode_detail::launch_int4( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(zeros), - static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + if (coalesce) { + moe_decode_detail::launch_int4_coalesced( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size, + num_experts); + } else { + moe_decode_detail::launch_int4( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + } } else { - moe_decode_detail::launch_int4( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(zeros), - static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + if (coalesce) { + moe_decode_detail::launch_int4_coalesced( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size, + num_experts); + } else { + moe_decode_detail::launch_int4( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), + static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, group_size); + } } } else { using BF = sycl::ext::oneapi::bfloat16; if (asym) { - moe_decode_detail::launch_int4( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(zeros), static_cast(outputs), - expert_id_per_token_buf, total_tokens, N, K, group_size); + if (coalesce) { + moe_decode_detail::launch_int4_coalesced( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size, num_experts); + } else { + moe_decode_detail::launch_int4( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size); + } } else { - moe_decode_detail::launch_int4( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(zeros), static_cast(outputs), - expert_id_per_token_buf, total_tokens, N, K, group_size); + if (coalesce) { + moe_decode_detail::launch_int4_coalesced( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size, num_experts); + } else { + moe_decode_detail::launch_int4( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(zeros), static_cast(outputs), + expert_id_per_token_buf, total_tokens, N, K, group_size); + } } } return; diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index ebeef5da52..3a53e0701e 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -560,6 +560,68 @@ def test_decode_int4_asym(self, dtype): assert out.shape == (total_tokens, N) torch.testing.assert_close(out, ref, rtol=5e-2, atol=5e-2) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("asym", [False, True]) + @pytest.mark.parametrize("group_size", [32, 128]) + def test_decode_int4_coalesced_matches_scalar(self, monkeypatch, dtype, asym, group_size): + """int4 scalar-GEMV fallback: the coalesced-load variant + (ARK_MOE_DECODE_COALESCE_INT4=1, default) must match both the legacy + per-lane-strided kernel (ARK_MOE_DECODE_COALESCE_INT4=0) and the + dequant->bmm reference within quantization tolerance. + + The S4 DPAS fast path is disabled so both runs exercise the scalar + fallback (this is the only path the coalesce flag affects). Shapes use + N%16==0 so the N-tiled repack is exact. + """ + num_experts = 4 + tokens_per_expert = [1, 0, 2, 1] + total_tokens = sum(tokens_per_expert) + N, K = 256, 256 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + if asym: + zeros = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_asym(w_float, scales, zeros, group_size) + dequant = _dequant_int4_asym(packed, scales, zeros, group_size).to(dtype) + else: + zeros = None + packed = _pack_int4_sym(w_float, scales, group_size) + dequant = _dequant_int4_sym(packed, scales, group_size).to(dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + zeros=zeros, + weight_bits=4, + group_size=group_size, + asym=asym, + ) + + # Force the scalar-GEMV fallback so the coalesce flag actually applies. + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") + out_coalesced = _run() + + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "0") + out_scalar = _run() + + assert out_coalesced.shape == (total_tokens, N) + assert out_scalar.shape == (total_tokens, N) + torch.testing.assert_close(out_coalesced, ref, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(out_scalar, ref, rtol=5e-2, atol=5e-2) + # The two kernels are numerically identical (same dequant math, only + # the weight memory layout differs), so require a tight match. + torch.testing.assert_close(out_coalesced, out_scalar, rtol=1e-3, atol=1e-3) + def test_decode_validation_errors(self): """Sanity-check that Python-side validation catches misuse.""" num_experts = 2 From 8c732ed8c16e3523572ed39116985ee547f4d876 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:49:19 +0000 Subject: [PATCH 25/47] perf: block tokens-per-expert in coalesced int4 MoE decode GEMV Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 164 +++++++++++++----- auto_round_extension/ark/test/test_moe.py | 63 +++++++ 2 files changed, 182 insertions(+), 45 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index de9df4cdf8..0f052f6c90 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -82,6 +82,16 @@ namespace moe_decode_detail { constexpr int SG_SIZE = 16; constexpr int N_TILE = SG_SIZE; // one output element per sub-group lane +// Token-blocking factor for the coalesced int4 decode GEMV. A work-item that +// owns one (n_tile, lane) output column processes up to TOKEN_BLOCK consecutive +// tokens, loading each packed weight byte from the (expert, n_tile) tile once +// and applying it to every token in the block that routes to the same expert. +// When decode routing is bursty (runs of tokens hitting the same expert), this +// amortizes the dominant weight traffic across the block instead of re-reading +// the tile once per token, moving the problem from pure GEMV toward GEMM. A +// value of 1 reproduces the one-token-per-work-item behaviour exactly. +constexpr int TOKEN_BLOCK = 4; + // ---------------------------------------------------------------------------- // Kernel name tags (one per specialization, required for SYCL kernel naming) // ---------------------------------------------------------------------------- @@ -405,6 +415,17 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // vectorized `int4_decode_chunk` (contiguous per-lane load) does not apply // here; the inner loop reads one packed byte per lane per step, which the // hardware coalesces across the sub-group into one wide transaction. +// +// On top of coalescing, this path blocks tokens: each work-item owns one +// output column but processes up to `TOKEN_BLOCK` consecutive tokens. For each +// distinct expert appearing in the block it makes a single weight-streaming +// pass and reuses every loaded (coalesced) byte across all tokens in the block +// routed to that expert. When decode routing is bursty -- runs of tokens +// hitting the same expert -- this amortizes the dominant weight traffic across +// the block (GEMV -> small GEMM). Fully-scattered routing degrades gracefully +// to one pass per token with the same per-pass weight reads as before, so the +// result is bit-identical to the one-token-per-work-item kernel regardless of +// routing. // ---------------------------------------------------------------------------- template void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, @@ -452,68 +473,121 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin }); } - sycl::range<2> global{static_cast(total_tokens), static_cast(n_tiles * SG_SIZE)}; + sycl::range<2> global{static_cast((total_tokens + TOKEN_BLOCK - 1) / TOKEN_BLOCK), + static_cast(n_tiles * SG_SIZE)}; sycl::range<2> local{1, static_cast(SG_SIZE)}; q->parallel_for>( sycl::nd_range<2>(global, local), [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { - const int token = static_cast(it.get_global_id(0)); + const int token_base = static_cast(it.get_global_id(0)) * TOKEN_BLOCK; const int n_tile = static_cast(it.get_group(1)); const int lane = static_cast(it.get_local_id(1)); const int n_global = n_tile * N_TILE + lane; - const int expert = expert_id_per_token[token]; - const ScalarT* act_row = activations + static_cast(token) * K; + // Number of tokens this work-item owns (last block may be short). + int block = TOKEN_BLOCK; + if (token_base + block > total_tokens) { + block = total_tokens - token_base; + } - // Base of this (expert, n_tile) weight tile in the repacked buffer. - // Layout [E, N/16, K/2, 16]; this lane reads byte kb at - // w_tile[kb*16 + lane], so adjacent lanes read adjacent bytes. - const uint8_t* w_tile = - repacked + ((static_cast(expert) * n_tiles + n_tile) * k_packed) * N_TILE; - const ScalarT* s_row = - scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; - const ScalarT* z_row = Asym - ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k - : nullptr; + // Experts routed by each token in the block. The tile weight byte is + // loaded once per k-step and reused only for tokens whose expert + // matches the byte's owning expert, so blocking tokens that share an + // expert amortizes the dominant weight traffic; tokens with a + // different expert contribute nothing from this pass and are handled + // by the pass whose leader expert matches theirs. + int experts[TOKEN_BLOCK]; + for (int b = 0; b < block; ++b) { + experts[b] = expert_id_per_token[token_base + b]; + } - float acc = 0.0f; - for (int g = 0; g < num_groups_k; ++g) { - const float scale = static_cast(s_row[g]); - float zero = 0.0f; - if constexpr (Asym) { - zero = static_cast(z_row[g]); + // Which distinct experts appear in this block. For each we make one + // weight-streaming pass, reusing every loaded byte across all tokens + // in the block routed to that expert. Bursty routing collapses to a + // single pass; fully-scattered routing degrades to one pass per token + // (i.e. the previous behaviour) with no extra weight reads per pass. + for (int lead = 0; lead < block; ++lead) { + const int expert = experts[lead]; + // Skip experts already streamed by an earlier token in this block. + bool seen = false; + for (int p = 0; p < lead; ++p) { + if (experts[p] == expert) { + seen = true; + break; + } } - const int k_base = g * group_size; - // Two partial accumulators break the fp32 dependency chain; the - // per-group scale/zero is folded once after the K-loop, exactly as - // in the scalar path. Each iteration processes two K-elements (one - // packed byte); the byte load is coalesced across the sub-group. - float acc_q0 = 0.0f; - float acc_q1 = 0.0f; - float acc_a = 0.0f; - const int kb_base = k_base / 2; - const int kb_count = group_size / 2; - for (int kb = 0; kb < kb_count; ++kb) { - const uint8_t packed = w_tile[(kb_base + kb) * N_TILE + lane]; - int q0, q1; - decode_int4_pair(packed, q0, q1); - const float fa0 = static_cast(act_row[k_base + 2 * kb]); - const float fa1 = static_cast(act_row[k_base + 2 * kb + 1]); - acc_q0 += fa0 * static_cast(q0); - acc_q1 += fa1 * static_cast(q1); + if (seen) continue; + + // Base of this (expert, n_tile) weight tile in the repacked buffer. + // Layout [E, N/16, K/2, 16]; this lane reads byte kb at + // w_tile[kb*16 + lane], so adjacent lanes read adjacent bytes. + const uint8_t* w_tile = + repacked + ((static_cast(expert) * n_tiles + n_tile) * k_packed) * N_TILE; + const ScalarT* s_row = + scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* z_row = Asym + ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k + : nullptr; + + float acc[TOKEN_BLOCK]; + for (int b = 0; b < block; ++b) acc[b] = 0.0f; + + for (int g = 0; g < num_groups_k; ++g) { + const float scale = static_cast(s_row[g]); + float zero = 0.0f; if constexpr (Asym) { - acc_a += fa0 + fa1; + zero = static_cast(z_row[g]); + } + const int k_base = g * group_size; + // Per-token split accumulators; the per-group scale/zero is folded + // once after the K-loop, exactly as in the scalar path. Each + // iteration processes two K-elements (one packed byte); the byte + // load is coalesced across the sub-group and reused across every + // token in the block routed to `expert`. + float acc_q0[TOKEN_BLOCK]; + float acc_q1[TOKEN_BLOCK]; + float acc_a[TOKEN_BLOCK]; + for (int b = 0; b < block; ++b) { + acc_q0[b] = 0.0f; + acc_q1[b] = 0.0f; + acc_a[b] = 0.0f; + } + const int kb_base = k_base / 2; + const int kb_count = group_size / 2; + for (int kb = 0; kb < kb_count; ++kb) { + const uint8_t packed = w_tile[(kb_base + kb) * N_TILE + lane]; + int q0, q1; + decode_int4_pair(packed, q0, q1); + const float fq0 = static_cast(q0); + const float fq1 = static_cast(q1); + for (int b = 0; b < block; ++b) { + if (experts[b] != expert) continue; + const ScalarT* act_row = activations + static_cast(token_base + b) * K; + const float fa0 = static_cast(act_row[k_base + 2 * kb]); + const float fa1 = static_cast(act_row[k_base + 2 * kb + 1]); + acc_q0[b] += fa0 * fq0; + acc_q1[b] += fa1 * fq1; + if constexpr (Asym) { + acc_a[b] += fa0 + fa1; + } + } + } + for (int b = 0; b < block; ++b) { + if (experts[b] != expert) continue; + float group_dot = acc_q0[b] + acc_q1[b]; + if constexpr (Asym) { + group_dot -= zero * acc_a[b]; + } + acc[b] += scale * group_dot; } } - float group_dot = acc_q0 + acc_q1; - if constexpr (Asym) { - group_dot -= zero * acc_a; + + for (int b = 0; b < block; ++b) { + if (experts[b] != expert) continue; + outputs[static_cast(token_base + b) * N + n_global] = static_cast(acc[b]); } - acc += scale * group_dot; } - - outputs[static_cast(token) * N + n_global] = static_cast(acc); }); q->wait(); diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 3a53e0701e..bcecc1446e 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -622,6 +622,69 @@ def _run(): # the weight memory layout differs), so require a tight match. torch.testing.assert_close(out_coalesced, out_scalar, rtol=1e-3, atol=1e-3) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("asym", [False, True]) + def test_decode_int4_coalesced_token_blocking(self, monkeypatch, dtype, asym): + """The coalesced int4 fallback blocks up to TOKEN_BLOCK consecutive + tokens per work-item, reusing each loaded weight byte across tokens + routed to the same expert. This must stay bit-identical to the legacy + per-lane-strided kernel regardless of routing, so exercise: + - a run of many tokens on one expert (full + short trailing block), + - blocks that straddle an expert boundary (mixed experts per block), + - an expert with zero tokens. + Shapes keep N%16==0 so the N-tiled repack is exact. + """ + num_experts = 4 + # 7 tokens on expert 0, none on expert 1, 5 on expert 2, 3 on expert 3. + # With TOKEN_BLOCK=4 this yields full blocks, short trailing blocks and + # at least one block straddling the 0->2 and 2->3 expert boundaries. + tokens_per_expert = [7, 0, 5, 3] + total_tokens = sum(tokens_per_expert) + N, K = 256, 128 + group_size = 32 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + if asym: + zeros = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_asym(w_float, scales, zeros, group_size) + dequant = _dequant_int4_asym(packed, scales, zeros, group_size).to(dtype) + else: + zeros = None + packed = _pack_int4_sym(w_float, scales, group_size) + dequant = _dequant_int4_sym(packed, scales, group_size).to(dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + zeros=zeros, + weight_bits=4, + group_size=group_size, + asym=asym, + ) + + # Force the scalar-GEMV fallback so the coalesce/token-blocking path runs. + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") + out_blocked = _run() + + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "0") + out_scalar = _run() + + assert out_blocked.shape == (total_tokens, N) + torch.testing.assert_close(out_blocked, ref, rtol=5e-2, atol=5e-2) + # Token blocking only changes weight reuse, not the dequant math, so the + # blocked and legacy kernels must match tightly for every routing shape. + torch.testing.assert_close(out_blocked, out_scalar, rtol=1e-3, atol=1e-3) + def test_decode_validation_errors(self): """Sanity-check that Python-side validation catches misuse.""" num_experts = 2 From cc939f8d6cb81d442368f62b092c219ddc75dcc6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:28:37 +0000 Subject: [PATCH 26/47] fix: parameterize MoEDecodeRepackInt4 kernel name to avoid duplicate mangled names Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 0f052f6c90..298df971f2 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -104,6 +104,7 @@ class MoEDecodeKernelInt4; template class MoEDecodeKernelInt4Coalesced; +template class MoEDecodeRepackInt4; template @@ -460,7 +461,7 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin { sycl::range<3> rp_global{static_cast(num_experts), static_cast(N), static_cast(k_packed)}; - q->parallel_for(rp_global, [=](sycl::id<3> id) { + q->parallel_for>(rp_global, [=](sycl::id<3> id) { const int e = static_cast(id[0]); const int n = static_cast(id[1]); const int kb = static_cast(id[2]); From 590ce59a0528db1b0463046c64046b84a9f46df5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 13:51:11 +0000 Subject: [PATCH 27/47] perf: gate coalesced int4 decode on repack amortization to fix tiny-batch regression Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 298df971f2..d66b5b4c2b 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1037,6 +1037,32 @@ inline bool moe_decode_coalesce_int4_enabled() { return true; } +// ---------------------------------------------------------------------------- +// Amortization gate for the coalesced int4 decode kernel. The coalesced path +// repacks the *entire* weight tensor for all `num_experts` on every call +// (cost proportional to num_experts * N * K/2) before running the GEMV. That +// one-time repack only pays off when a work-group reuses each repacked weight +// tile across many tokens -- i.e. when there are enough active tokens relative +// to the number of experts. For tiny decode batches (e.g. 8 tokens spread +// across 192 experts) the repack dominates and the coalesced kernel is far +// slower than the per-lane-strided `launch_int4`, which reads the weights in +// place with no repack. Require at least one full TOKEN_BLOCK worth of tokens +// per expert on average before coalescing; otherwise fall back to `launch_int4`. +// `ARK_MOE_DECODE_COALESCE_MIN_TOKENS` overrides the threshold (tokens per +// expert scaled by TOKEN_BLOCK); "0" disables the gate (always coalesce). +// ---------------------------------------------------------------------------- +inline bool moe_decode_coalesce_int4_amortized(int total_tokens, int num_experts) { + if (num_experts <= 0) return true; + long long min_tokens = static_cast(num_experts) * TOKEN_BLOCK; + const char* env = std::getenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS"); + if (env != nullptr) { + char* end = nullptr; + long long v = std::strtoll(env, &end, 10); + if (end != env && v >= 0) min_tokens = v; + } + return static_cast(total_tokens) >= min_tokens; +} + // ---------------------------------------------------------------------------- // // weight_dtype: @@ -1112,7 +1138,8 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo // default this uses the coalesced-load variant, which repacks the weights // on-device so sub-group loads are contiguous; `ARK_MOE_DECODE_COALESCE_INT4=0` // forces the legacy per-lane-strided kernel. - const bool coalesce = moe_decode_coalesce_int4_enabled(); + const bool coalesce = moe_decode_coalesce_int4_enabled() && + moe_decode_coalesce_int4_amortized(total_tokens, num_experts); if (act_dtype == BTLA_DTYPE::F16) { if (asym) { if (coalesce) { From d80a9a134e39854c78f66dac474d6a725f932a68 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:24:56 +0000 Subject: [PATCH 28/47] fix: qualify TOKEN_BLOCK with moe_decode_detail namespace Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index d66b5b4c2b..7c141b0eab 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1053,7 +1053,7 @@ inline bool moe_decode_coalesce_int4_enabled() { // ---------------------------------------------------------------------------- inline bool moe_decode_coalesce_int4_amortized(int total_tokens, int num_experts) { if (num_experts <= 0) return true; - long long min_tokens = static_cast(num_experts) * TOKEN_BLOCK; + long long min_tokens = static_cast(num_experts) * moe_decode_detail::TOKEN_BLOCK; const char* env = std::getenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS"); if (env != nullptr) { char* end = nullptr; From 7db3d5f24fb2a156717ba77275e2f38b49da70c7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:09:30 +0000 Subject: [PATCH 29/47] perf: hoist routing filter out of int4 coalesced decode hot loop Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 57 ++++++++++++------- 1 file changed, 36 insertions(+), 21 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 7c141b0eab..8a60fd6417 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -531,8 +531,24 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k : nullptr; + // Compact the tokens routed to `expert` into a dense member list + // once per pass. Hoisting the routing filter out of the hot k-loop + // removes a per-(kb, token) branch and lets the compiler keep the + // per-member activation base pointers in registers; the numerics + // are identical to the previous per-kb `experts[b] != expert` + // filter. + int members[TOKEN_BLOCK]; + const ScalarT* act_rows[TOKEN_BLOCK]; + int nmembers = 0; + for (int b = 0; b < block; ++b) { + if (experts[b] != expert) continue; + members[nmembers] = b; + act_rows[nmembers] = activations + static_cast(token_base + b) * K; + ++nmembers; + } + float acc[TOKEN_BLOCK]; - for (int b = 0; b < block; ++b) acc[b] = 0.0f; + for (int m = 0; m < nmembers; ++m) acc[m] = 0.0f; for (int g = 0; g < num_groups_k; ++g) { const float scale = static_cast(s_row[g]); @@ -549,10 +565,10 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin float acc_q0[TOKEN_BLOCK]; float acc_q1[TOKEN_BLOCK]; float acc_a[TOKEN_BLOCK]; - for (int b = 0; b < block; ++b) { - acc_q0[b] = 0.0f; - acc_q1[b] = 0.0f; - acc_a[b] = 0.0f; + for (int m = 0; m < nmembers; ++m) { + acc_q0[m] = 0.0f; + acc_q1[m] = 0.0f; + acc_a[m] = 0.0f; } const int kb_base = k_base / 2; const int kb_count = group_size / 2; @@ -562,31 +578,30 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin decode_int4_pair(packed, q0, q1); const float fq0 = static_cast(q0); const float fq1 = static_cast(q1); - for (int b = 0; b < block; ++b) { - if (experts[b] != expert) continue; - const ScalarT* act_row = activations + static_cast(token_base + b) * K; - const float fa0 = static_cast(act_row[k_base + 2 * kb]); - const float fa1 = static_cast(act_row[k_base + 2 * kb + 1]); - acc_q0[b] += fa0 * fq0; - acc_q1[b] += fa1 * fq1; + const int k0 = k_base + 2 * kb; + for (int m = 0; m < nmembers; ++m) { + const ScalarT* act_row = act_rows[m]; + const float fa0 = static_cast(act_row[k0]); + const float fa1 = static_cast(act_row[k0 + 1]); + acc_q0[m] += fa0 * fq0; + acc_q1[m] += fa1 * fq1; if constexpr (Asym) { - acc_a[b] += fa0 + fa1; + acc_a[m] += fa0 + fa1; } } } - for (int b = 0; b < block; ++b) { - if (experts[b] != expert) continue; - float group_dot = acc_q0[b] + acc_q1[b]; + for (int m = 0; m < nmembers; ++m) { + float group_dot = acc_q0[m] + acc_q1[m]; if constexpr (Asym) { - group_dot -= zero * acc_a[b]; + group_dot -= zero * acc_a[m]; } - acc[b] += scale * group_dot; + acc[m] += scale * group_dot; } } - for (int b = 0; b < block; ++b) { - if (experts[b] != expert) continue; - outputs[static_cast(token_base + b) * N + n_global] = static_cast(acc[b]); + for (int m = 0; m < nmembers; ++m) { + const int b = members[m]; + outputs[static_cast(token_base + b) * N + n_global] = static_cast(acc[m]); } } }); From c8ef1b8e59c2a0375b6fa8e8d243065091b6191b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:45:10 +0000 Subject: [PATCH 30/47] feat: add dedicated w4a16_policy_m_8 decode DPAS dispatch for S4 sym Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 20 +++-- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 81 +++++++++++++++++++ 2 files changed, 93 insertions(+), 8 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 8a60fd6417..b1d637f9ea 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1012,12 +1012,16 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // ---------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_DECODE_DPAS_S4` (default ON). When ON, int4-sym -// (S4_CLIP, !asym) decode is routed to the shared per-group S4 DPAS grouped -// GEMM (`moe_dpas_s4::moe_prefill_s4_dpas_per_group_dispatch`) instead of the -// scalar FMA GEMV (`launch_int4`). The DPAS path already handles the tiny -// total-token counts typical of decode (its `A_avg_M <= 4` bucket selects the -// 8-row `dpas_w4a16_policy_m_8` tile) and reads the same `[E, N, K/2]` packed -// weights + `[E, N, K/group]` scales, so no repack is needed. +// (S4_CLIP, !asym) decode is routed to the dedicated decode-phase S4 DPAS +// grouped GEMM (`moe_dpas_s4::moe_decode_s4_dpas_per_group_dispatch`) instead +// of the scalar FMA GEMV (`launch_int4`). Ported from vLLM-xpu-kernels' +// dedicated `w4a16` decode dispatch, this path hard-pins the 8-row +// `dpas_w4a16_policy_m_8` tile (the only tile the prefill `A_avg_M` ladder +// would ever pick for decode-sized batches) and reuses the shared per-group +// mainloop's 2D VNNI block load + register-resident per-N scale. It reads the +// same `[E, N, K/2]` packed weights + `[E, N, K/group]` scales, so no repack is +// needed. (`ARK_MOE_DECODE_S4_DPAS_M8=0` defers to the full prefill bucket +// ladder for A/B comparison; the two are numerically identical.) // // Setting the var to "0" / "false" / "off" / "no" (case-insensitive) forces // the legacy scalar GEMV, for A/B comparison and regression escape. Asym @@ -1136,13 +1140,13 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo if (!asym && moe_decode_dpas_s4_enabled() && moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size)) { if (act_dtype == BTLA_DTYPE::F16) { - moe_dpas_s4::moe_prefill_s4_dpas_per_group_dispatch( + moe_dpas_s4::moe_decode_s4_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), static_cast(scales), static_cast(outputs), num_tokens_per_expert, num_experts, N, K, group_size, total_tokens); } else { using BF = sycl::ext::oneapi::bfloat16; - moe_dpas_s4::moe_prefill_s4_dpas_per_group_dispatch( + moe_dpas_s4::moe_decode_s4_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), static_cast(scales), static_cast(outputs), num_tokens_per_expert, num_experts, N, K, group_size, total_tokens); diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index 693a8675e5..f59a863cc7 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -678,6 +678,87 @@ void moe_prefill_s4_dpas_per_group_dispatch( sycl::free(atomic_buffer, *q); } +// --------------------------------------------------------------------------- +// Decode-specialized S4 (sym) per-K-group DPAS dispatch. +// +// Ported from vllm-xpu-kernels' dedicated `w4a16` *decode* dispatch: for the +// generation (decode) phase every expert sees at most a handful of tokens, so +// the 8-row `dpas_w4a16_policy_m_8` tile is always the right choice. Rather +// than piggy-backing on `moe_prefill_s4_dpas_per_group_dispatch` and re-running +// its four-tier `A_avg_M` bucket ladder (which only selects `_m_8` at the tiny +// end anyway), this entry point hard-pins the 8-row tile. +// +// The underlying grouped-GEMM mainloop (`xe_gemm_*` per-group loop) is reused +// verbatim: it already performs the 2D VNNI block load via +// `get_block_2d_copy_A/B` + `make_block_2d_prefetch`, and caches the per-N +// scale in the `sg_scale[]` register array, folding it once per K-group. So no +// new mainloop math is written here -- only the tile selection differs from the +// prefill dispatch. +// +// `ARK_MOE_DECODE_S4_DPAS_M8` (default ON) can be set to "0"/"false"/"off"/"no" +// to fall back to the full prefill `A_avg_M` bucket ladder for A/B comparison +// (the two paths are numerically identical; only the tile shape differs). +// --------------------------------------------------------------------------- +inline bool moe_decode_s4_dpas_m8_enabled() { + const char* env = std::getenv("ARK_MOE_DECODE_S4_DPAS_M8"); + if (env == nullptr) return true; // default ON + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "0" || s == "false" || s == "off" || s == "no") return false; + return true; +} + +template +void moe_decode_s4_dpas_per_group_dispatch( + sycl::queue* q, const ScalarT* activations, const uint8_t* weights_NKp, + const ScalarT* scales, ScalarT* outputs, + const int* num_tokens_per_expert, int E, int N, int K, int group_size, + int total_tokens) { + // A/B escape: defer to the prefill bucket ladder when the m_8 pin is + // disabled. Identical math; only the DPAS tile shape differs. + if (!moe_decode_s4_dpas_m8_enabled()) { + moe_prefill_s4_dpas_per_group_dispatch( + q, activations, weights_NKp, scales, outputs, num_tokens_per_expert, E, + N, K, group_size, total_tokens); + return; + } + + if (E == 0 || N == 0 || K == 0 || total_tokens == 0) return; + if (K % group_size != 0) { + throw std::invalid_argument( + "moe_decode_s4_dpas(per-group): K must be a multiple of group_size"); + } + if ((K & 1) != 0) { + throw std::invalid_argument( + "moe_decode_s4_dpas(per-group): K must be even (packed nibbles)"); + } + + compat::set_default_queue(*q); + + using ElementA = cute_scalar_t; + const auto* activations_ca = + reinterpret_cast(activations); + const auto* scales_ca = reinterpret_cast(scales); + auto* outputs_ca = reinterpret_cast(outputs); + const auto* weights_i4 = + reinterpret_cast(weights_NKp); + + int32_t* atomic_buffer = sycl::malloc_device(1, *q); + if (atomic_buffer == nullptr) { + throw std::runtime_error( + "moe_decode_s4_dpas(per-group): failed to allocate atomic buffer"); + } + + // Hard-pin the 8-row tile -- the only tile the prefill ladder would ever + // select for decode-sized batches. + MoEGEMMLauncher_s4<'R', 'C', dpas_w4a16_policy_m_8>( + *q, activations_ca, weights_i4, scales_ca, + static_cast(nullptr), outputs_ca, N, K, + num_tokens_per_expert, E, group_size, atomic_buffer); + + sycl::free(atomic_buffer, *q); +} + // --------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default ON, semantics // identical to `moe_prefill_dpas_int_enabled` / `moe_prefill_dpas_fp8 From a3c82ced15b5b796e6834b4d7a2599c545e188a9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:46:48 +0000 Subject: [PATCH 31/47] docs+test: document decode S4 DPAS path and add A/B parity assertion Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/test/README_MOE_PREFILL_PERF.md | 15 +++++++++++++++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 14 ++++++++++++++ auto_round_extension/ark/test/test_moe.py | 10 ++++++++++ 3 files changed, 39 insertions(+) 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 c1ed20ef5b..bea7f499b1 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -311,6 +311,21 @@ The mid-size `32×64` tile now covers `A_avg_M` up to 128 (previously it jumped to the wide tile at 33), which avoids padding waste on the common chunked-prefill batch sizes. +**S4 DPAS decode path** — the *decode* phase (`sycl_tla_moe_decode.hpp`, +int4-sym / `S4_CLIP`, `!asym`, `ARK_MOE_DECODE_DPAS_S4` default ON) now +has its own dedicated dispatch, `moe_decode_s4_dpas_per_group_dispatch`, +ported from vLLM-xpu-kernels' dedicated `w4a16` *decode* dispatch. +Because every expert sees at most a handful of tokens during generation, +it hard-pins the 8-row `dpas_w4a16_policy_m_8` tile (the only tile the +prefill `A_avg_M` ladder would ever pick for decode-sized batches) rather +than re-running that ladder. It reuses the shared per-group mainloop's +2D VNNI block load (`get_block_2d_copy_A/B` + `make_block_2d_prefetch`) +and register-resident per-N scale (`sg_scale[]`, folded once per K-group), +reading the same `[E, N, K/2]` packed weights + `[E, N, K/group]` scales +with no repack. `ARK_MOE_DECODE_S4_DPAS_M8=0` defers to the full prefill +bucket ladder for A/B comparison (numerically identical; only the tile +shape differs). **Status: NEEDS-HARDWARE-VALIDATION** (untested port). + Accuracy parity is covered by `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group`, which forces `ARK_MOE_PREFILL_DPAS_S4=1` + 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 713474651d..96d7a53396 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 @@ -240,6 +240,20 @@ nibble 的 B 流字节量是 INT8 路径的一半,大 M tile 加宽到 `128×256 中等大小的 `32×64` tile 现在覆盖 `A_avg_M` 至 128(此前在 33 就跳到大 tile),避免了常见 chunked-prefill batch 大小下的 padding 浪费。 +**S4 DPAS decode 路径** — decode(生成)阶段(`sycl_tla_moe_decode.hpp`, +int4-sym / `S4_CLIP`,`!asym`,`ARK_MOE_DECODE_DPAS_S4` 默认开启)现在拥有 +独立的 dispatch `moe_decode_s4_dpas_per_group_dispatch`,移植自 +vLLM-xpu-kernels 专用的 `w4a16` *decode* dispatch。由于生成阶段每个专家 +最多只见到少量 token,它直接钉死 8 行的 `dpas_w4a16_policy_m_8` tile +(prefill 的 `A_avg_M` 阶梯对 decode 规模的 batch 也只会选这个 tile), +而不再运行整个阶梯。它复用共享的 per-group mainloop 的 2D VNNI 块加载 +(`get_block_2d_copy_A/B` + `make_block_2d_prefetch`)与寄存器驻留的 +per-N scale(`sg_scale[]`,每个 K-group 折叠一次),读取相同的 +`[E, N, K/2]` 打包权重 + `[E, N, K/group]` scale,无需重新打包。 +`ARK_MOE_DECODE_S4_DPAS_M8=0` 会回退到完整的 prefill 阶梯以便 A/B 对比 +(数值完全相同,仅 tile 形状不同)。**状态:NEEDS-HARDWARE-VALIDATION** +(未经测试的移植)。 + 精度对齐由 `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group` 覆盖,该用例强制 `ARK_MOE_PREFILL_DPAS_S4=1` + diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index bcecc1446e..a57f87b089 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -516,8 +516,15 @@ def _run(): ) monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + monkeypatch.setenv("ARK_MOE_DECODE_S4_DPAS_M8", "1") out_dpas = _run() + # A/B escape: deferring to the prefill A_avg_M bucket ladder must be + # numerically identical (only the DPAS tile shape differs). + monkeypatch.setenv("ARK_MOE_DECODE_S4_DPAS_M8", "0") + out_dpas_ladder = _run() + monkeypatch.setenv("ARK_MOE_DECODE_S4_DPAS_M8", "1") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") out_scalar = _run() @@ -528,6 +535,9 @@ def _run(): torch.testing.assert_close(out_scalar, ref, rtol=5e-2, atol=5e-2) # And they must agree with each other within the same tolerance. torch.testing.assert_close(out_dpas, out_scalar, rtol=5e-2, atol=5e-2) + # The m_8-pinned decode dispatch and the prefill bucket ladder are the + # same DPAS math on decode-sized batches. + torch.testing.assert_close(out_dpas, out_dpas_ladder, rtol=5e-2, atol=5e-2) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_decode_int4_asym(self, dtype): From 339608282dd6df73d0901a194bcfb7bf93194b17 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:02:45 +0000 Subject: [PATCH 32/47] perf: reuse persistent atomic buffer in S4 DPAS dispatch Replace per-call sycl::malloc_device/free of the 1-int32 work-group counter in the prefill and decode S4 DPAS grouped-GEMM dispatchers with a persistent per-queue buffer. Each malloc/free forces a queue sync, which is pure overhead on the decode hot path. The kernel self-zeroes the counter at launch and every launch is synchronous, so a single reused buffer per queue is safe. Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 52 ++++++++++++++----- 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index f59a863cc7..7b015ff314 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -91,9 +91,11 @@ #pragma once #include +#include #include #include #include +#include #ifdef ARK_XPU #include @@ -143,6 +145,38 @@ using ::ark::moe_dpas_fp8::cute_scalar; using ::ark::moe_dpas_fp8::cute_scalar_t; using ::ark::moe_dpas_fp8::make_moe_tensor; +// --------------------------------------------------------------------------- +// Persistent per-queue atomic work-group counter. +// +// The grouped-GEMM launcher below needs a single `int32_t` device slot as a +// global work-group counter (`atomicAdd`). The kernel self-initialises it to 0 +// at launch (group 0 / lane 0 does `atm.store(0)`), so the host never has to +// reset it between calls. Previously each dispatch call allocated this slot +// with `sycl::malloc_device` and released it with `sycl::free`; both operations +// force a queue synchronization, which is pure overhead on the decode hot path +// where the GEMM itself is only tens of microseconds. +// +// Instead, hand out one persistent buffer per queue and reuse it across calls. +// This is safe because every launcher call is synchronous (`event.wait()` in +// `MoEGEMMLauncher_s4`), so two launches can never share the buffer +// concurrently. Buffers live until process exit (one `int32_t` per queue), +// matching the singleton lifetime already used by `EventManager`. +// --------------------------------------------------------------------------- +inline int32_t* get_persistent_atomic_buffer(sycl::queue* q) { + static std::mutex mtx; + static std::unordered_map cache; + std::lock_guard lock(mtx); + auto it = cache.find(q); + if (it != cache.end()) return it->second; + int32_t* buf = sycl::malloc_device(1, *q); + if (buf == nullptr) { + throw std::runtime_error( + "moe_dpas_s4: failed to allocate persistent atomic buffer"); + } + cache.emplace(q, buf); + return buf; +} + // --------------------------------------------------------------------------- // Variant B -- per-K-group S4 (sym) mainloop. // @@ -646,11 +680,9 @@ void moe_prefill_s4_dpas_per_group_dispatch( int A_avg_M = total_tokens / E; - int32_t* atomic_buffer = sycl::malloc_device(1, *q); - if (atomic_buffer == nullptr) { - throw std::runtime_error( - "moe_prefill_s4_dpas(per-group): failed to allocate atomic buffer"); - } + // Reusable per-queue work-group counter (self-zeroed by the kernel); avoids + // a malloc_device/free (each a queue sync) on every dispatch call. + int32_t* atomic_buffer = get_persistent_atomic_buffer(q); #define ARK_DPAS_S4_PG_LAUNCH_SYM(policy) \ MoEGEMMLauncher_s4<'R', 'C', policy>( \ @@ -674,8 +706,6 @@ void moe_prefill_s4_dpas_per_group_dispatch( ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w4a16_policy); } #undef ARK_DPAS_S4_PG_LAUNCH_SYM - - sycl::free(atomic_buffer, *q); } // --------------------------------------------------------------------------- @@ -743,11 +773,7 @@ void moe_decode_s4_dpas_per_group_dispatch( const auto* weights_i4 = reinterpret_cast(weights_NKp); - int32_t* atomic_buffer = sycl::malloc_device(1, *q); - if (atomic_buffer == nullptr) { - throw std::runtime_error( - "moe_decode_s4_dpas(per-group): failed to allocate atomic buffer"); - } + int32_t* atomic_buffer = get_persistent_atomic_buffer(q); // Hard-pin the 8-row tile -- the only tile the prefill ladder would ever // select for decode-sized batches. @@ -755,8 +781,6 @@ void moe_decode_s4_dpas_per_group_dispatch( *q, activations_ca, weights_i4, scales_ca, static_cast(nullptr), outputs_ca, N, K, num_tokens_per_expert, E, group_size, atomic_buffer); - - sycl::free(atomic_buffer, *q); } // --------------------------------------------------------------------------- From ba2c18d78fb7b2a650d298a38b8cc35daff13fe7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:31:27 +0000 Subject: [PATCH 33/47] perf: trim device-timeline overhead on S4 DPAS decode hot path Two waste removals on the int4-sym decode path, motivated by the device-event perf measurement (host-side malloc/free reuse showed no change because it never lands on the device timeline): - Skip the fill_expert_id_per_token kernel when the S4-sym DPAS fast path is taken; that path uses num_tokens_per_expert directly and never reads expert_id_per_token_buf, so the extra device kernel launch was pure overhead. The predicate is shared with the S4 block so the two decisions cannot diverge. - Cache the per-launch query_device_multiprocessor_count host driver query in MoEGEMMLauncher_s4 (device SM count is fixed), removing a Level Zero round-trip from every decode launch. Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 20 ++++++++++++++----- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 6 ++++-- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index b1d637f9ea..d8cb44b32c 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1108,8 +1108,17 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo void* outputs, int* expert_id_per_token_buf, BTLA_DTYPE act_dtype, BTLA_DTYPE weight_dtype, int N, int K, int group_size, int* num_tokens_per_expert, int num_experts, int total_tokens, bool asym) { - moe_decode_detail::fill_expert_id_per_token(q, expert_id_per_token_buf, num_tokens_per_expert, num_experts, - total_tokens); + // The S4-sym DPAS fast path consumes `num_tokens_per_expert` directly and + // never reads `expert_id_per_token_buf`. Skipping the fill on that path + // removes an extra device-timeline kernel launch from the decode hot path; + // every other path (fp, int8, int2, fp8, and the scalar int4 fallback) still + // needs the per-token expert mapping. + const bool s4_dpas_fastpath = weight_dtype == BTLA_DTYPE::S4_CLIP && !asym && moe_decode_dpas_s4_enabled() && + moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size); + if (!s4_dpas_fastpath) { + moe_decode_detail::fill_expert_id_per_token(q, expert_id_per_token_buf, num_tokens_per_expert, num_experts, + total_tokens); + } if (weight_dtype == BTLA_DTYPE::F16 || weight_dtype == BTLA_DTYPE::BF16) { if (weight_dtype != act_dtype) { @@ -1136,9 +1145,10 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo // Fast path: sym int4 through the shared per-group S4 DPAS grouped GEMM. // Falls back to the scalar GEMV for asym weights (DPAS S4 is sym-only), // when the env flag is off, or when the shape gate rejects the tile - // geometry (e.g. N%64!=0, K%32!=0, unsupported group_size). - if (!asym && moe_decode_dpas_s4_enabled() && - moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size)) { + // geometry (e.g. N%64!=0, K%32!=0, unsupported group_size). Reuses the + // `s4_dpas_fastpath` predicate computed above (which also gated the + // `fill_expert_id_per_token` skip) so the two decisions cannot diverge. + if (s4_dpas_fastpath) { if (act_dtype == BTLA_DTYPE::F16) { moe_dpas_s4::moe_decode_s4_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index 7b015ff314..7c43e08e45 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -591,8 +591,10 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, SGLayout>::TiledMMA; auto mma = MMA{}; - int sm_count = - cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + // Device SM count is a host-side driver query (Level Zero round-trip). The + // value is fixed for a given device, so cache it once instead of paying the + // round-trip on every launch -- the decode hot path calls this per step. + static const int sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); auto MaxThreadsPerWorkgroup = size(mma); static constexpr int MaxThreadsPerSM = 512; From 5b88c2ec6d9cca59883647134547c2fed8ccfb43 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:50:11 +0000 Subject: [PATCH 34/47] perf(ark): use A_avg_M tile ladder for int4 moe sym decode (match vllm-xpu-kernels) Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 16 +++---- .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 46 +++++++++++-------- .../ark/test/README_MOE_PREFILL_PERF.md | 31 +++++++------ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 24 ++++++---- 4 files changed, 67 insertions(+), 50 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index d8cb44b32c..2aa6300ab6 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1014,14 +1014,14 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // Env-flag helper -- `ARK_MOE_DECODE_DPAS_S4` (default ON). When ON, int4-sym // (S4_CLIP, !asym) decode is routed to the dedicated decode-phase S4 DPAS // grouped GEMM (`moe_dpas_s4::moe_decode_s4_dpas_per_group_dispatch`) instead -// of the scalar FMA GEMV (`launch_int4`). Ported from vLLM-xpu-kernels' -// dedicated `w4a16` decode dispatch, this path hard-pins the 8-row -// `dpas_w4a16_policy_m_8` tile (the only tile the prefill `A_avg_M` ladder -// would ever pick for decode-sized batches) and reuses the shared per-group -// mainloop's 2D VNNI block load + register-resident per-N scale. It reads the -// same `[E, N, K/2]` packed weights + `[E, N, K/group]` scales, so no repack is -// needed. (`ARK_MOE_DECODE_S4_DPAS_M8=0` defers to the full prefill bucket -// ladder for A/B comparison; the two are numerically identical.) +// of the scalar FMA GEMV (`launch_int4`). Mirroring vLLM-xpu-kernels' +// `w4a16` decode dispatch, this path selects the DPAS tile from the average +// tokens-per-expert (`A_avg_M`) ladder (`_m_8` -> `_m_16` -> `_m_32` -> wide), +// reusing the shared per-group mainloop's 2D VNNI block load + register-resident +// per-N scale. It reads the same `[E, N, K/2]` packed weights + `[E, N, K/group]` +// scales, so no repack is needed. (`ARK_MOE_DECODE_S4_DPAS_M8=1` forces the +// legacy hard-pinned 8-row tile for A/B comparison; the two are numerically +// identical.) // // Setting the var to "0" / "false" / "off" / "no" (case-insensitive) forces // the legacy scalar GEMV, for A/B comparison and regression escape. Asym diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index 7c43e08e45..ba2bb9ca6e 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -713,31 +713,38 @@ void moe_prefill_s4_dpas_per_group_dispatch( // --------------------------------------------------------------------------- // Decode-specialized S4 (sym) per-K-group DPAS dispatch. // -// Ported from vllm-xpu-kernels' dedicated `w4a16` *decode* dispatch: for the -// generation (decode) phase every expert sees at most a handful of tokens, so -// the 8-row `dpas_w4a16_policy_m_8` tile is always the right choice. Rather -// than piggy-backing on `moe_prefill_s4_dpas_per_group_dispatch` and re-running -// its four-tier `A_avg_M` bucket ladder (which only selects `_m_8` at the tiny -// end anyway), this entry point hard-pins the 8-row tile. +// Mirrors vllm-xpu-kernels' `w4a16` dispatch (`grouped_gemm_xe2_interface.hpp`): +// the DPAS tile is selected from the average tokens-per-expert (`A_avg_M`) +// ladder -- `_m_8` for `A_avg_M <= 4`, `_m_16` for `<= 8`, `_m_32` for `<= 128`, +// then the wide tile. An earlier revision hard-pinned the 8-row +// `dpas_w4a16_policy_m_8` tile on the assumption that decode only ever sees a +// handful of tokens per expert, but that assumption breaks for larger decode +// batches (many sequences, high top-k, or few experts): whenever more than 4 +// tokens route to an expert on average the 8-row tile under-fills the M +// dimension and the (bandwidth-bound) packed weights are re-streamed 2-4x more +// than necessary, roughly halving throughput versus the reference. Selecting +// the same tile as the reference collapses to `_m_8` for tiny batches (so the +// small-decode case is unchanged) and grows the M tile exactly when reuse pays +// off. // // The underlying grouped-GEMM mainloop (`xe_gemm_*` per-group loop) is reused // verbatim: it already performs the 2D VNNI block load via // `get_block_2d_copy_A/B` + `make_block_2d_prefetch`, and caches the per-N // scale in the `sg_scale[]` register array, folding it once per K-group. So no // new mainloop math is written here -- only the tile selection differs from the -// prefill dispatch. +// prefill dispatch (which uses the identical ladder). // -// `ARK_MOE_DECODE_S4_DPAS_M8` (default ON) can be set to "0"/"false"/"off"/"no" -// to fall back to the full prefill `A_avg_M` bucket ladder for A/B comparison -// (the two paths are numerically identical; only the tile shape differs). +// `ARK_MOE_DECODE_S4_DPAS_M8` (default OFF) can be set to "1"/"true"/"on"/"yes" +// to force the legacy hard-pinned 8-row tile for A/B comparison and regression +// escape (the two paths are numerically identical; only the tile shape differs). // --------------------------------------------------------------------------- -inline bool moe_decode_s4_dpas_m8_enabled() { +inline bool moe_decode_s4_dpas_force_m8() { const char* env = std::getenv("ARK_MOE_DECODE_S4_DPAS_M8"); - if (env == nullptr) return true; // default ON + if (env == nullptr) return false; // default OFF -- use the A_avg_M ladder std::string s(env); for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "0" || s == "false" || s == "off" || s == "no") return false; - return true; + if (s == "1" || s == "true" || s == "on" || s == "yes") return true; + return false; } template @@ -746,9 +753,10 @@ void moe_decode_s4_dpas_per_group_dispatch( const ScalarT* scales, ScalarT* outputs, const int* num_tokens_per_expert, int E, int N, int K, int group_size, int total_tokens) { - // A/B escape: defer to the prefill bucket ladder when the m_8 pin is - // disabled. Identical math; only the DPAS tile shape differs. - if (!moe_decode_s4_dpas_m8_enabled()) { + // Default: select the DPAS tile from the shared `A_avg_M` ladder, matching + // the reference `w4a16` decode dispatch. Identical math; only the DPAS tile + // shape differs. `ARK_MOE_DECODE_S4_DPAS_M8=1` forces the legacy m_8 pin. + if (!moe_decode_s4_dpas_force_m8()) { moe_prefill_s4_dpas_per_group_dispatch( q, activations, weights_NKp, scales, outputs, num_tokens_per_expert, E, N, K, group_size, total_tokens); @@ -777,8 +785,8 @@ void moe_decode_s4_dpas_per_group_dispatch( int32_t* atomic_buffer = get_persistent_atomic_buffer(q); - // Hard-pin the 8-row tile -- the only tile the prefill ladder would ever - // select for decode-sized batches. + // Legacy opt-in path (`ARK_MOE_DECODE_S4_DPAS_M8=1`): hard-pin the 8-row + // tile. Kept for A/B comparison against the default `A_avg_M` ladder above. MoEGEMMLauncher_s4<'R', 'C', dpas_w4a16_policy_m_8>( *q, activations_ca, weights_i4, scales_ca, static_cast(nullptr), outputs_ca, N, K, 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 bea7f499b1..c7bb58ea76 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -312,19 +312,24 @@ jumped to the wide tile at 33), which avoids padding waste on the common chunked-prefill batch sizes. **S4 DPAS decode path** — the *decode* phase (`sycl_tla_moe_decode.hpp`, -int4-sym / `S4_CLIP`, `!asym`, `ARK_MOE_DECODE_DPAS_S4` default ON) now -has its own dedicated dispatch, `moe_decode_s4_dpas_per_group_dispatch`, -ported from vLLM-xpu-kernels' dedicated `w4a16` *decode* dispatch. -Because every expert sees at most a handful of tokens during generation, -it hard-pins the 8-row `dpas_w4a16_policy_m_8` tile (the only tile the -prefill `A_avg_M` ladder would ever pick for decode-sized batches) rather -than re-running that ladder. It reuses the shared per-group mainloop's -2D VNNI block load (`get_block_2d_copy_A/B` + `make_block_2d_prefetch`) -and register-resident per-N scale (`sg_scale[]`, folded once per K-group), -reading the same `[E, N, K/2]` packed weights + `[E, N, K/group]` scales -with no repack. `ARK_MOE_DECODE_S4_DPAS_M8=0` defers to the full prefill -bucket ladder for A/B comparison (numerically identical; only the tile -shape differs). **Status: NEEDS-HARDWARE-VALIDATION** (untested port). +int4-sym / `S4_CLIP`, `!asym`, `ARK_MOE_DECODE_DPAS_S4` default ON) has +its own dedicated dispatch, `moe_decode_s4_dpas_per_group_dispatch`, +mirroring vLLM-xpu-kernels' `w4a16` decode dispatch. It selects the DPAS +tile from the same `A_avg_M` ladder as prefill (`_m_8` → `_m_16` → `_m_32` +→ wide): the 8-row tile is used only for the tiny-batch tail (`A_avg_M ≤ +4`), and the M tile grows once more than four tokens route to an expert +on average. An earlier revision hard-pinned the 8-row tile on the +assumption that decode only ever sees a handful of tokens per expert, but +that under-fills the M dimension and re-streams the (bandwidth-bound) +packed weights 2–4× on larger decode batches (many sequences, high top-k, +or few experts), roughly halving throughput versus the reference. It +reuses the shared per-group mainloop's 2D VNNI block load +(`get_block_2d_copy_A/B` + `make_block_2d_prefetch`) and register-resident +per-N scale (`sg_scale[]`, folded once per K-group), reading the same +`[E, N, K/2]` packed weights + `[E, N, K/group]` scales with no repack. +`ARK_MOE_DECODE_S4_DPAS_M8=1` forces the legacy hard-pinned 8-row tile for +A/B comparison (numerically identical; only the tile shape differs). +**Status: NEEDS-HARDWARE-VALIDATION** (untested port). Accuracy parity is covered by `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group`, 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 96d7a53396..1c7d00d8a2 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 @@ -241,17 +241,21 @@ nibble 的 B 流字节量是 INT8 路径的一半,大 M tile 加宽到 `128×256 tile),避免了常见 chunked-prefill batch 大小下的 padding 浪费。 **S4 DPAS decode 路径** — decode(生成)阶段(`sycl_tla_moe_decode.hpp`, -int4-sym / `S4_CLIP`,`!asym`,`ARK_MOE_DECODE_DPAS_S4` 默认开启)现在拥有 -独立的 dispatch `moe_decode_s4_dpas_per_group_dispatch`,移植自 -vLLM-xpu-kernels 专用的 `w4a16` *decode* dispatch。由于生成阶段每个专家 -最多只见到少量 token,它直接钉死 8 行的 `dpas_w4a16_policy_m_8` tile -(prefill 的 `A_avg_M` 阶梯对 decode 规模的 batch 也只会选这个 tile), -而不再运行整个阶梯。它复用共享的 per-group mainloop 的 2D VNNI 块加载 -(`get_block_2d_copy_A/B` + `make_block_2d_prefetch`)与寄存器驻留的 -per-N scale(`sg_scale[]`,每个 K-group 折叠一次),读取相同的 +int4-sym / `S4_CLIP`,`!asym`,`ARK_MOE_DECODE_DPAS_S4` 默认开启)拥有 +独立的 dispatch `moe_decode_s4_dpas_per_group_dispatch`,对齐 +vLLM-xpu-kernels 的 `w4a16` decode dispatch。它与 prefill 使用相同的 +`A_avg_M` 阶梯选择 DPAS tile(`_m_8` → `_m_16` → `_m_32` → 大 tile): +仅在极小 batch 尾部(`A_avg_M ≤ 4`)使用 8 行 tile,一旦平均每个专家 +路由超过 4 个 token,M tile 就随之增大。早先的版本直接钉死 8 行的 +`dpas_w4a16_policy_m_8` tile,假设 decode 阶段每个专家只见到少量 token, +但在较大的 decode batch(序列多、top-k 高或专家少)下,这会导致 M +维度欠填充,并把(受带宽约束的)打包权重重复流式加载 2–4 次,使吞吐 +大约只有参考实现的一半。它复用共享的 per-group mainloop 的 2D VNNI 块 +加载(`get_block_2d_copy_A/B` + `make_block_2d_prefetch`)与寄存器驻留 +的 per-N scale(`sg_scale[]`,每个 K-group 折叠一次),读取相同的 `[E, N, K/2]` 打包权重 + `[E, N, K/group]` scale,无需重新打包。 -`ARK_MOE_DECODE_S4_DPAS_M8=0` 会回退到完整的 prefill 阶梯以便 A/B 对比 -(数值完全相同,仅 tile 形状不同)。**状态:NEEDS-HARDWARE-VALIDATION** +`ARK_MOE_DECODE_S4_DPAS_M8=1` 会强制使用旧的钉死 8 行 tile 以便 A/B +对比(数值完全相同,仅 tile 形状不同)。**状态:NEEDS-HARDWARE-VALIDATION** (未经测试的移植)。 精度对齐由 From 1f7b6313f05469a58c2a1ec04a54730225ea024e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:27:55 +0000 Subject: [PATCH 35/47] perf(ark): route decode-sized int4-sym moe through the asym scalar GEMV Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 59 ++++++++++++++++--- .../ark/test/README_MOE_PREFILL_PERF.md | 15 +++++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 13 ++++ auto_round_extension/ark/test/test_moe.py | 6 +- .../ark/test/test_moe_decode_perf.py | 26 +++++--- 5 files changed, 102 insertions(+), 17 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 2aa6300ab6..c0a9865d28 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1025,8 +1025,10 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // // Setting the var to "0" / "false" / "off" / "no" (case-insensitive) forces // the legacy scalar GEMV, for A/B comparison and regression escape. Asym -// weights and shapes that fail the DPAS shape gate always fall back to the -// scalar path regardless of this flag. Re-read on every call so tests / +// weights, shapes that fail the DPAS shape gate, and batches that fail the +// tokens-per-expert occupancy gate (`moe_decode_dpas_s4_occupancy_ok`, see +// below -- this is what keeps real decode batches on the fast scalar GEMV) +// always fall back to the scalar path regardless of this flag. Re-read on every call so tests / // benchmarks can toggle the path in-process. // ---------------------------------------------------------------------------- inline bool moe_decode_dpas_s4_enabled() { @@ -1038,10 +1040,43 @@ inline bool moe_decode_dpas_s4_enabled() { return true; } +// ---------------------------------------------------------------------------- +// Occupancy gate for the int4-sym S4 DPAS decode path. +// +// The DPAS grouped GEMM pays off only when its M tile is actually filled: the +// smallest tile (`dpas_w4a16_policy_m_8`) processes 8 token rows per expert, so +// with fewer than 8 tokens routed to an expert on average the tile is mostly +// padding and the (bandwidth-bound) packed weights are streamed for rows that +// contribute nothing. Real decode batches are exactly that regime -- e.g. +// MiniMax-M2 decode is 8 tokens (bs1) or 256 tokens (bs32) spread over 192 +// experts, i.e. 0.04-1.3 tokens per expert -- and there the shared scalar GEMV +// (`launch_int4`, the very kernel the *asym* path uses, where sym is just +// `Asym=false`) is up to ~3x faster because it reads each weight byte exactly +// once per active token with no tile padding. +// +// So route int4-sym decode through the same scalar GEMV as int4-asym unless the +// batch has at least one full 8-row tile of tokens per expert on average. +// `ARK_MOE_DECODE_DPAS_S4_MIN_TPE` overrides the tokens-per-expert threshold; +// "0" disables the gate (always take DPAS when the shape gate allows), which is +// what the accuracy tests use to exercise the DPAS kernel on tiny shapes. +// ---------------------------------------------------------------------------- +inline bool moe_decode_dpas_s4_occupancy_ok(int total_tokens, int num_experts) { + if (num_experts <= 0) return true; + long long min_tokens_per_expert = 8; // rows in `dpas_w4a16_policy_m_8` + const char* env = std::getenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE"); + if (env != nullptr) { + char* end = nullptr; + long long v = std::strtoll(env, &end, 10); + if (end != env && v >= 0) min_tokens_per_expert = v; + } + if (min_tokens_per_expert == 0) return true; + return static_cast(total_tokens) >= min_tokens_per_expert * static_cast(num_experts); +} + // ---------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_DECODE_COALESCE_INT4` (default ON). When ON, the // int4 scalar-GEMV fallback (asym, or sym with the DPAS path disabled / shape -// gate miss) uses `launch_int4_coalesced`, which repacks the weights on-device +// or occupancy gate miss) uses `launch_int4_coalesced`, which repacks the weights on-device // into an N-tiled layout so sub-group weight loads are coalesced. Setting the // var to "0" / "false" / "off" / "no" (case-insensitive) forces the legacy // per-lane-strided `launch_int4`, for A/B comparison and regression escape. @@ -1094,10 +1129,12 @@ inline bool moe_decode_coalesce_int4_amortized(int total_tokens, int num_experts // scales [E, N, K/group_size] in act dtype, // zeros optional (asym==true requires it). // Sym weights are routed to the shared -// per-group S4 DPAS grouped GEMM by default -// (`ARK_MOE_DECODE_DPAS_S4`, default ON); -// asym, a disabled flag, or a shape-gate -// miss falls back to the scalar GEMV. +// per-group S4 DPAS grouped GEMM only when +// the batch fills its M tile (>= 8 tokens per +// expert on average, `ARK_MOE_DECODE_DPAS_S4` +// default ON); asym, a disabled flag, a +// shape-gate miss, or a decode-sized batch +// uses the shared scalar GEMV. // BTLA_DTYPE::S2_CLIP : packed int2 weights [E, N, K/4] (uint8), // 4 values per byte, sym/asym like int4 // BTLA_DTYPE::F8_E4M3 / F8_E5M2 : FP8 weights [E, N, K] (uint8 buffer), @@ -1114,6 +1151,7 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo // every other path (fp, int8, int2, fp8, and the scalar int4 fallback) still // needs the per-token expert mapping. const bool s4_dpas_fastpath = weight_dtype == BTLA_DTYPE::S4_CLIP && !asym && moe_decode_dpas_s4_enabled() && + moe_decode_dpas_s4_occupancy_ok(total_tokens, num_experts) && moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size); if (!s4_dpas_fastpath) { moe_decode_detail::fill_expert_id_per_token(q, expert_id_per_token_buf, num_tokens_per_expert, num_experts, @@ -1144,8 +1182,11 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo } // Fast path: sym int4 through the shared per-group S4 DPAS grouped GEMM. // Falls back to the scalar GEMV for asym weights (DPAS S4 is sym-only), - // when the env flag is off, or when the shape gate rejects the tile - // geometry (e.g. N%64!=0, K%32!=0, unsupported group_size). Reuses the + // when the env flag is off, when the batch is too small to fill the DPAS M + // tile (the usual decode case -- sym then runs the exact same + // `launch_int4*` kernel as asym, with `Asym=false`), or when the shape gate + // rejects the tile geometry (e.g. N%64!=0, K%32!=0, unsupported + // group_size). Reuses the // `s4_dpas_fastpath` predicate computed above (which also gated the // `fill_expert_id_per_token` skip) so the two decisions cannot diverge. if (s4_dpas_fastpath) { 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 c7bb58ea76..592b3f0907 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -331,6 +331,21 @@ per-N scale (`sg_scale[]`, folded once per K-group), reading the same A/B comparison (numerically identical; only the tile shape differs). **Status: NEEDS-HARDWARE-VALIDATION** (untested port). +**Occupancy gate — decode-sized batches use the int4-asym kernel.** Even +the smallest DPAS tile processes 8 token rows per expert, so a batch with +fewer than 8 tokens per expert on average pays full weight-streaming cost +for mostly-padding rows. That is precisely the decode regime: on +MiniMax-M2 (192 experts) bs1 is 8 tokens and bs32 is 256 tokens, i.e. +0.04–1.3 tokens per expert, and measurements showed int4-sym (DPAS) at +0.31–0.34 ms/1.55 ms against int4-asym (scalar GEMV) at 0.12 ms/1.45 ms +for the same shapes. int4-sym decode is therefore routed to the *same* +scalar GEMV kernel that int4-asym uses (`launch_int4` / its coalesced +variant, with `Asym=false` — sym is cheaper there because it skips the +zero-point fold) unless the batch supplies at least 8 tokens per expert. +`ARK_MOE_DECODE_DPAS_S4_MIN_TPE` overrides the tokens-per-expert +threshold; `0` disables the gate (always DPAS when the shape gate allows), +which is what the accuracy and DPAS-vs-scalar perf tests set. + Accuracy parity is covered by `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group`, which forces `ARK_MOE_PREFILL_DPAS_S4=1` + 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 1c7d00d8a2..8070ff2230 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 @@ -258,6 +258,19 @@ vLLM-xpu-kernels 的 `w4a16` decode dispatch。它与 prefill 使用相同的 对比(数值完全相同,仅 tile 形状不同)。**状态:NEEDS-HARDWARE-VALIDATION** (未经测试的移植)。 +**占用率门控 — decode 规模的 batch 直接复用 int4-asym 的实现。** 即使是 +最小的 DPAS tile 也要处理每个专家 8 行 token,因此平均每个专家不足 8 个 +token 的 batch 会为几乎全是 padding 的行付出完整的权重流式加载代价。 +decode 正好处于这一区间:MiniMax-M2(192 个专家)bs1 只有 8 个 token, +bs32 只有 256 个 token,即平均每个专家 0.04–1.3 个 token;实测同样形状下 +int4-sym(DPAS)为 0.31–0.34 ms / 1.55 ms,而 int4-asym(标量 GEMV)为 +0.12 ms / 1.45 ms。因此除非 batch 平均每个专家至少有 8 个 token,int4-sym +的 decode 会被路由到与 int4-asym *完全相同* 的标量 GEMV kernel +(`launch_int4` 及其 coalesced 变体,`Asym=false` — sym 少了 zero-point +折叠,反而更快)。`ARK_MOE_DECODE_DPAS_S4_MIN_TPE` 可覆盖该 +"每专家 token 数" 阈值;设为 `0` 则关闭门控(只要形状门控通过就走 DPAS), +精度测试与 DPAS/标量 对比性能测试即使用该设置。 + 精度对齐由 `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group` 覆盖,该用例强制 `ARK_MOE_PREFILL_DPAS_S4=1` + diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index a57f87b089..77b76c0575 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -489,7 +489,10 @@ def test_decode_int4_sym_dpas_matches_scalar(self, monkeypatch, dtype, group_siz the dequant->bmm reference within quantization tolerance. Shapes satisfy the DPAS per-group shape gate (N%64==0, K%32==0, - group_size in {32,64,128,256}) so the fast path is actually taken. + group_size in {32,64,128,256}) and ``ARK_MOE_DECODE_DPAS_S4_MIN_TPE=0`` + disables the tokens-per-expert occupancy gate (these tiny token counts + would otherwise be routed to the scalar GEMV, which is faster there) so + the DPAS fast path is actually taken. """ num_experts = 4 total_tokens = sum(tokens_per_expert) @@ -516,6 +519,7 @@ def _run(): ) monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE", "0") monkeypatch.setenv("ARK_MOE_DECODE_S4_DPAS_M8", "1") out_dpas = _run() diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index 584287528d..26641cc95c 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -403,12 +403,17 @@ def test_perf_int4_sym_dpas_vs_scalar(self, monkeypatch, dtype): ``speedup`` here is ``scalar / dpas`` (the DPAS path is the "ark" column), isolating the DPAS routing win from the dequant reference. Only shapes that clear the DPAS shape gate are timed on both paths. - - Observed crossover: at bs1 (8 tokens, ~1 token/expert) the DPAS tile is - nearly empty and the scalar GEMV wins (~0.5x), but by bs32 (256 tokens) - the shared S4 DPAS grouped-GEMM already wins ~2x. The auto-dispatch - default threshold (`_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS`) is therefore kept - small so 256-token batches are routed to prefill/DPAS, not decode. + ``ARK_MOE_DECODE_DPAS_S4_MIN_TPE=0`` disables the tokens-per-expert + occupancy gate so the DPAS column really runs DPAS (by default these + decode-sized batches are routed to the scalar GEMV). + + Observed on MiniMax-M2 decode shapes (192 experts): the DPAS tile is + starved at every decode batch measured -- 8 tokens (bs1, 0.04 + tokens/expert) and 256 tokens (bs32, 1.3 tokens/expert) are both far + below the 8 rows of `dpas_w4a16_policy_m_8` -- and the scalar GEMV wins, + which is why the default occupancy gate keeps decode on the scalar path + (the same kernel int4-asym uses) until a batch supplies >= 8 tokens per + expert. """ group_size = 128 _print_header( @@ -439,7 +444,9 @@ def _run(): monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") scalar_ms = _xpu_time_ms(_run) monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE", "0") dpas_ms = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE", raising=False) _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -454,7 +461,10 @@ def test_perf_int4_sym_dpas_vs_scalar_threshold(self, monkeypatch, dtype): crosses 1.0x at the total-token count where the shared S4 DPAS grouped-GEMM starts beating the scalar GEMV, which is the value to feed into ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` / the ``moe(...)`` - ``decode_threshold`` auto-dispatch cutoff. + ``decode_threshold`` auto-dispatch cutoff, and is also the sweep behind + the default ``ARK_MOE_DECODE_DPAS_S4_MIN_TPE`` occupancy gate (set to + ``0`` here so the DPAS column is not itself re-routed to the scalar + GEMV). """ group_size = 32 E = 192 @@ -488,7 +498,9 @@ def _run(): monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") scalar_ms = _xpu_time_ms(_run) monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "1") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE", "0") dpas_ms = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE", raising=False) _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) From 211e95f8238b1116ef2f02f8b8432b423a74493f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:55:11 +0000 Subject: [PATCH 36/47] perf(ark): decode int4-sym nibbles through the asym unsigned path (^0x88 sign-flip) Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 105 +++++++++++------- .../ark/test/README_MOE_PREFILL_PERF.md | 21 +++- .../ark/test/README_MOE_PREFILL_PERF_CN.md | 17 ++- 3 files changed, 97 insertions(+), 46 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index c0a9865d28..72121cebbd 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -240,18 +240,36 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight // without regressing group_size == 16 (which drops straight to the 16-wide // stage). // -// The per-group scale (and asym zero-point) is NOT applied here: this -// accumulates the raw integer-weighted dot product ``Σ a·q`` into ``acc_q`` -// and, for the asym case, the plain activation sum ``Σ a`` into ``acc_a``. -// The caller folds the group's scale/zero in once (Σ a·((q−z)·s) == -// s·(Σ a·q − z·Σ a); sym collapses to s·Σ a·q). Hoisting the scale removes one -// float multiply per K element on the decode hot path, and because the scale -// fold is exact-once per group the result stays well within the kernel's -// existing quantization tolerance. +// Both sym and asym decode the nibbles through the *same* unsigned path. +// Measurements on MiniMax-M2 decode shapes showed int4-sym ~1.9x slower than +// int4-asym in this very kernel even though sym does strictly fewer floating +// point operations. The only difference was the per-nibble sign extension +// (`(int8_t)(byte << 4) >> 4`) -- a serial shift/narrow/shift chain per nibble +// that defeats the byte-wise vectorization the asym mask+shift form gets. The +// sym decode is therefore expressed with the standard sign-flip identity +// +// signed_nibble == (unsigned_nibble ^ 8) - 8 +// +// so XOR-ing the packed byte with `0x88` (flipping the sign bit of *both* +// nibbles at once, on the whole loaded vector register) turns sym into exactly +// the asym computation with a constant zero-point of 8. The decoded integers +// are bit-identical to the sign-extending decode for all 256 byte values, so +// the only change is that sym now accumulates the biased sum and subtracts +// `8 * sum a` at the end -- exactly the fp32 accumulation pattern asym has +// always used, and well inside the kernel's existing quantization tolerance. +// +// The per-group scale and zero-point are NOT applied here: this accumulates the +// raw integer-weighted dot product into `acc_q0`/`acc_q1` and the plain +// activation sum into `acc_a` (now needed by both modes, since sym carries the +// constant zero-point of 8). The caller folds the group's scale/zero in once +// (sum a*((q-z)*s) == s*(sum a*q - z*sum a)). Hoisting the scale removes one +// float multiply per K element on the decode hot path, and because the fold is +// exact-once per group the result stays well within the kernel's existing +// quantization tolerance. // // Two independent partial accumulators (``acc_q0``/``acc_q1``) break the // single fp32 dependency chain so the FMA pipeline is not latency-bound; the -// caller reduces the pair. ``acc_a`` (asym only) reuses the same split. +// caller reduces the pair. ``acc_a`` reuses the same split. template static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc_q0, float& acc_q1, float& acc_a) { @@ -267,20 +285,26 @@ static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_pt #pragma unroll for (int s = 0; s < CHUNK / SUB; ++s) { const ActVec av = *reinterpret_cast(act_ptr + s * SUB); - const PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 2)); + PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 2)); + if constexpr (!Asym) { + // Sign-flip the whole packed vector in one vector XOR so the sym nibbles + // can be decoded by the (vectorizable) unsigned path below; the constant + // zero-point of 8 is folded by the caller. + pv = pv ^ PackVec(static_cast(0x88)); + } #pragma unroll for (int b = 0; b < SUB / 2; ++b) { int q0, q1; - decode_int4_pair(pv[b], q0, q1); + // Always the unsigned decode: asym nibbles are unsigned by definition and + // sym nibbles were biased by the XOR above. + decode_int4_pair(pv[b], q0, q1); const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); const float fa0 = static_cast(a0); const float fa1 = static_cast(a1); acc_q0 += fa0 * static_cast(q0); acc_q1 += fa1 * static_cast(q1); - if constexpr (Asym) { - acc_a += fa0 + fa1; - } + acc_a += fa0 + fa1; } } } @@ -329,7 +353,10 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig float acc = 0.0f; for (int g = 0; g < num_groups_k; ++g) { const float scale = static_cast(s_row[g]); - float zero = 0.0f; + // Sym uses the constant zero-point of 8 that the `^0x88` sign-flip + // in the decode introduces, so both modes run the identical fold. + // `if constexpr` keeps the null `z_row` out of the sym instantiation. + float zero = 8.0f; if constexpr (Asym) { zero = static_cast(z_row[g]); } @@ -342,13 +369,14 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // (32/64/128/256), while the 16-wide stage keeps group_size == 16 // on the fast path. // - // The scale (and asym zero) is constant across the group, so the - // wide stages accumulate the raw integer-weighted dot product + // The scale and zero are constant across the group, so the wide + // stages accumulate the raw integer-weighted dot product // ``Σ a·q`` (split across two partial accumulators to break the - // fp32 dependency chain) plus ``Σ a`` for asym, and the fold below - // applies the scale/zero exactly once per group: - // sym : acc += scale * (acc_q0 + acc_q1) - // asym: acc += scale * ((acc_q0 + acc_q1) - zero * acc_a) + // fp32 dependency chain) plus ``Σ a``, and the fold below applies + // the scale/zero exactly once per group, identically for both modes: + // acc += scale * ((acc_q0 + acc_q1) - zero * acc_a) + // with `zero` == the per-group zero-point (asym) or the constant 8 + // that the `^0x88` sign-flip decode introduces (sym). float acc_q0 = 0.0f; float acc_q1 = 0.0f; float acc_a = 0.0f; @@ -369,22 +397,17 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // raw-accumulation convention as the wide stages so the single // scale/zero fold below stays valid. for (; kk < group_size; kk += 2) { - const uint8_t packed = w_row[(k_base + kk) / 2]; + uint8_t packed = w_row[(k_base + kk) / 2]; + if constexpr (!Asym) packed ^= static_cast(0x88); int q0, q1; - decode_int4_pair(packed, q0, q1); + decode_int4_pair(packed, q0, q1); const float fa0 = static_cast(act_row[k_base + kk]); const float fa1 = static_cast(act_row[k_base + kk + 1]); acc_q0 += fa0 * static_cast(q0); acc_q1 += fa1 * static_cast(q1); - if constexpr (Asym) { - acc_a += fa0 + fa1; - } + acc_a += fa0 + fa1; } - float group_dot = acc_q0 + acc_q1; - if constexpr (Asym) { - group_dot -= zero * acc_a; - } - acc += scale * group_dot; + acc += scale * ((acc_q0 + acc_q1) - zero * acc_a); } outputs[static_cast(token) * N + n_global] = static_cast(acc); @@ -552,7 +575,10 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin for (int g = 0; g < num_groups_k; ++g) { const float scale = static_cast(s_row[g]); - float zero = 0.0f; + // Constant zero-point of 8 for sym (see `int4_decode_chunk`): the + // `^0x88` sign-flip lets sym reuse the asym unsigned decode and + // fold, so both modes emit the identical instruction stream. + float zero = 8.0f; if constexpr (Asym) { zero = static_cast(z_row[g]); } @@ -573,9 +599,10 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin const int kb_base = k_base / 2; const int kb_count = group_size / 2; for (int kb = 0; kb < kb_count; ++kb) { - const uint8_t packed = w_tile[(kb_base + kb) * N_TILE + lane]; + uint8_t packed = w_tile[(kb_base + kb) * N_TILE + lane]; + if constexpr (!Asym) packed ^= static_cast(0x88); int q0, q1; - decode_int4_pair(packed, q0, q1); + decode_int4_pair(packed, q0, q1); const float fq0 = static_cast(q0); const float fq1 = static_cast(q1); const int k0 = k_base + 2 * kb; @@ -585,17 +612,11 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin const float fa1 = static_cast(act_row[k0 + 1]); acc_q0[m] += fa0 * fq0; acc_q1[m] += fa1 * fq1; - if constexpr (Asym) { - acc_a[m] += fa0 + fa1; - } + acc_a[m] += fa0 + fa1; } } for (int m = 0; m < nmembers; ++m) { - float group_dot = acc_q0[m] + acc_q1[m]; - if constexpr (Asym) { - group_dot -= zero * acc_a[m]; - } - acc[m] += scale * group_dot; + acc[m] += scale * ((acc_q0[m] + acc_q1[m]) - zero * acc_a[m]); } } 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 592b3f0907..435ae28392 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -340,12 +340,29 @@ MiniMax-M2 (192 experts) bs1 is 8 tokens and bs32 is 256 tokens, i.e. 0.31–0.34 ms/1.55 ms against int4-asym (scalar GEMV) at 0.12 ms/1.45 ms for the same shapes. int4-sym decode is therefore routed to the *same* scalar GEMV kernel that int4-asym uses (`launch_int4` / its coalesced -variant, with `Asym=false` — sym is cheaper there because it skips the -zero-point fold) unless the batch supplies at least 8 tokens per expert. +variant, with `Asym=false`) unless the batch supplies at least 8 tokens per +expert. `ARK_MOE_DECODE_DPAS_S4_MIN_TPE` overrides the tokens-per-expert threshold; `0` disables the gate (always DPAS when the shape gate allows), which is what the accuracy and DPAS-vs-scalar perf tests set. +**Sym decodes through the asym nibble path.** Once both modes shared the +scalar GEMV, int4-sym was still ~1.9x slower than int4-asym in the *same* +kernel (2.83 ms vs 1.49 ms at bs32) despite doing strictly fewer floating +point operations. The only asymmetry was the nibble decode: sym +sign-extended each nibble with `(int8_t)(byte << 4) >> 4`, a serial +shift/narrow/shift chain per nibble, while asym used a plain mask+shift +that vectorizes over the whole loaded byte vector. Sym now uses the +sign-flip identity `signed == (unsigned ^ 8) - 8`: one vector XOR of the +packed bytes with `0x88` flips both nibbles' sign bits, after which sym is +*literally* the asym computation with a constant zero-point of 8 (same +unsigned decode, same `sum a` accumulator, same single per-group +scale/zero fold). The decoded integers are bit-identical to the previous +sign-extending decode for all 256 byte values; the only change is that +sym now accumulates the biased sum and subtracts `8 * sum a` at the end, +which is exactly the fp32 accumulation pattern asym has always used. It +applies to both `launch_int4` and `launch_int4_coalesced`. + Accuracy parity is covered by `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group`, which forces `ARK_MOE_PREFILL_DPAS_S4=1` + 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 8070ff2230..8e43d7aae4 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 @@ -266,11 +266,24 @@ bs32 只有 256 个 token,即平均每个专家 0.04–1.3 个 token;实测同 int4-sym(DPAS)为 0.31–0.34 ms / 1.55 ms,而 int4-asym(标量 GEMV)为 0.12 ms / 1.45 ms。因此除非 batch 平均每个专家至少有 8 个 token,int4-sym 的 decode 会被路由到与 int4-asym *完全相同* 的标量 GEMV kernel -(`launch_int4` 及其 coalesced 变体,`Asym=false` — sym 少了 zero-point -折叠,反而更快)。`ARK_MOE_DECODE_DPAS_S4_MIN_TPE` 可覆盖该 +(`launch_int4` 及其 coalesced 变体,`Asym=false`)。`ARK_MOE_DECODE_DPAS_S4_MIN_TPE` 可覆盖该 "每专家 token 数" 阈值;设为 `0` 则关闭门控(只要形状门控通过就走 DPAS), 精度测试与 DPAS/标量 对比性能测试即使用该设置。 +**sym 直接复用 asym 的 nibble 解码路径。** 在两者都走标量 GEMV 之后， +int4-sym 在 *同一个* kernel 里仍比 int4-asym 慢约 1.9 倍（bs32：2.83 ms +vs 1.49 ms），尽管 sym 的浮点运算严格更少。唯一的差异在于 nibble 解码： +sym 使用 `(int8_t)(byte << 4) >> 4` 逐 nibble 做符号扩展，这是一条 +移位/截断/移位 的串行依赖链；而 asym 的 掩码+移位 形式可以在整个已加载 +的字节向量上向量化。现在 sym 改用符号翻转恒等式 +`signed == (unsigned ^ 8) - 8`：对打包字节做一次向量 `^ 0x88`，同时翻转 +两个 nibble 的符号位，之后 sym *就是* zero-point 恒为 8 的 asym 计算 +（相同的无符号解码、相同的 `sum a` 累加器、相同的每组一次 scale/zero +折叠）。对全部 256 种字节取值，解码出的整数与原来的符号扩展逐位相同， +唯一的变化是 sym 现在累加的是有偏置的和，最后再减去 `8 * sum a`， +这正是 asym 一直在用的 fp32 累加方式；`launch_int4` 与 `launch_int4_coalesced` +两个 kernel 均已应用。 + 精度对齐由 `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group` 覆盖,该用例强制 `ARK_MOE_PREFILL_DPAS_S4=1` + From e4d25c5d60679c983e4a10d9910c65d1ed30c8ac Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 06:28:51 +0000 Subject: [PATCH 37/47] perf(ark): vectorize int4 decode GEMV loads, hoist activation sums, pool the repack buffer Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/auto_round_kernel/__init__.py | 16 + .../ark/auto_round_kernel/ark.cpp | 1 + .../wrapper/include/sycl_tla_moe_decode.hpp | 624 +++++++++++++----- 3 files changed, 464 insertions(+), 177 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 69147abc2f..0a485bd7c8 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -1670,6 +1670,22 @@ def moe_gemm_decode( return outputs +def moe_decode_release_scratch() -> None: + """Release the device scratch buffers held by the int4 decode fallbacks. + + :func:`moe_gemm_decode` serves its int4 weight-repack and activation-sum + buffers from grow-on-demand per-queue slabs that are kept for the lifetime + of the process so the decode hot path never allocates. Call this to hand + that memory back, or to drop a repack cached via + ``ARK_MOE_DECODE_INT4_REPACK_CACHE=1`` before the underlying weight tensor + is freed. A no-op when the XPU extension is not loaded. + """ + lib = xpu_lib + if lib is None or not hasattr(lib, "moe_decode_release_scratch"): + return + lib.moe_decode_release_scratch() + + def _validate_moe_quant_args( activations: torch.Tensor, weights: torch.Tensor, diff --git a/auto_round_extension/ark/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index d9c9c7c03f..dc701e0dad 100755 --- a/auto_round_extension/ark/auto_round_kernel/ark.cpp +++ b/auto_round_extension/ark/auto_round_kernel/ark.cpp @@ -800,6 +800,7 @@ PYBIND11_MODULE(PY_NAME, m) { m.def("sage_dynamic_quant_v_layout", &ark::sage_dynamic_quant_v_layout); m.def("moe_gemm", &ark::moe_gemm_wrapper); m.def("moe_gemm_decode", &ark::moe_gemm_decode_wrapper); + m.def("moe_decode_release_scratch", &ark::moe_decode_release_scratch); 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); diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 72121cebbd..bc128bda1d 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -13,9 +13,11 @@ // - weights (int4 packed): [num_experts, N, K/2] row-major, two // 4-bit values per byte (low nibble at lower K). // The scalar-GEMV fallback repacks this on-device -// into an N-tiled [E, N/16, K/2, 16] layout so that -// sub-group weight loads are coalesced; the external -// [E, N, K/2] contract is unchanged. +// into an N-tiled [E, N/16, ceil(K/8), 16, 4] +// layout so that sub-group weight loads are +// coalesced *and* each lane loads 4 packed bytes at +// a time; the external [E, N, K/2] contract is +// unchanged. // - weights (int2 packed): [num_experts, N, K/4] row-major, four // 2-bit values per byte (field j at K index // 4*i+j is bits [2j+1:2j]) @@ -40,6 +42,8 @@ #include #include #include +#include +#include #include #include @@ -107,6 +111,9 @@ class MoEDecodeKernelInt4Coalesced; template class MoEDecodeRepackInt4; +template +class MoEDecodeActGroupSum; + template class MoEDecodeKernelInt8; @@ -221,6 +228,168 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight }); } +// ---------------------------------------------------------------------------- +// Persistent per-queue device scratch pool. +// +// The int4 decode fallbacks need two device-side scratch buffers: the N-tiled +// weight repack and the per-(token, K-group) activation-sum table. Allocating +// them with `sycl::malloc_device` on every call is not viable on the decode hot +// path -- decode issues one call per generated token, and a USM allocation +// (plus the `queue::wait()` that has to precede the matching `sycl::free`) +// costs on the order of the GEMV itself. Instead each buffer is served from a +// slab that is allocated once per queue and grown on demand, so steady-state +// decode performs no allocation and needs no host-side synchronization: the +// in-order queue already serializes the producer kernel before the consumer, +// which is the same ordering guarantee `fill_expert_id_per_token` relies on. +// +// A slab additionally carries an optional *tag* -- the address of the source +// buffer it was derived from plus a caller-supplied key that must fold in +// everything else the derived contents depend on (shape, layout parameters). +// `acquire` reports whether the slab already holds the result for that exact +// tag, which lets the caller skip regenerating it. This is only consulted when +// the caller opts in (see `moe_decode_int4_repack_cache_enabled`), because the +// address half of a tag is a pointer identity and a freed-then-reallocated +// buffer can land on the same address. +// +// Slabs are intentionally never freed from a static destructor: the SYCL +// context may already be torn down at that point. `release_all` provides +// explicit teardown for callers that need it (exposed to Python as +// `moe_decode_release_scratch`). +// ---------------------------------------------------------------------------- +class DeviceScratchPool { + public: + uint8_t* acquire(sycl::queue* q, size_t bytes, const void* tag_ptr, size_t tag_key, bool use_tag, + bool* tag_hit) { + std::lock_guard lock(mu_); + Slab& slab = slabs_[q]; + if (slab.ptr == nullptr || slab.bytes < bytes) { + if (slab.ptr != nullptr) { + // The old slab may still be referenced by in-flight kernels. + q->wait(); + sycl::free(slab.ptr, *q); + slab = Slab{}; + } + uint8_t* p = sycl::malloc_device(bytes, *q); + if (p == nullptr) { + throw std::runtime_error("moe_gemm_decode: failed to allocate device scratch buffer"); + } + slab.ptr = p; + slab.bytes = bytes; + } + const bool hit = use_tag && slab.tagged && slab.tag_ptr == tag_ptr && slab.tag_key == tag_key; + if (tag_hit != nullptr) *tag_hit = hit; + if (!hit) { + slab.tagged = use_tag; + slab.tag_ptr = tag_ptr; + slab.tag_key = tag_key; + } + return slab.ptr; + } + + uint8_t* acquire(sycl::queue* q, size_t bytes) { + return acquire(q, bytes, nullptr, 0, false, nullptr); + } + + void release_all() { + std::lock_guard lock(mu_); + for (auto& kv : slabs_) { + if (kv.second.ptr != nullptr) { + kv.first->wait(); + sycl::free(kv.second.ptr, *kv.first); + } + } + slabs_.clear(); + } + + private: + struct Slab { + uint8_t* ptr = nullptr; + size_t bytes = 0; + bool tagged = false; + const void* tag_ptr = nullptr; + size_t tag_key = 0; + }; + std::mutex mu_; + std::map slabs_; +}; + +inline DeviceScratchPool& int4_repack_pool() { + static DeviceScratchPool pool; + return pool; +} + +inline DeviceScratchPool& act_group_sum_pool() { + static DeviceScratchPool pool; + return pool; +} + +// ---------------------------------------------------------------------------- +// Per-(token, K-group) activation sums. +// +// Both int4 GEMV kernels fold their per-group scale/zero as +// `scale * (Σ a·q - zero · Σ a)`, where `Σ a` runs over the group's K range. +// `Σ a` depends only on the activation row and the group, *not* on the output +// column, yet the GEMVs used to recompute it inside the inner loop -- once per +// sub-group lane (16x redundant) and again for every N-tile work-group (N/16x +// redundant). That cost one extra float add per K element on the hot path, +// which for sym is pure overhead introduced by the constant zero-point of 8 +// that the `^0x88` decode relies on. +// +// This pass computes the `[total_tokens, K/group_size]` table once, so the +// GEMVs only accumulate `Σ a·q` and read one float per group. The table is +// tiny (tokens x groups floats) and comes from the scratch pool, so no +// allocation happens in steady state. +// +// The summation order differs from the previous in-loop accumulation, so +// results move by a few float ULPs -- far inside the kernel's quantization +// tolerance. +// ---------------------------------------------------------------------------- +template +void launch_act_group_sums(sycl::queue* q, const ScalarT* activations, float* a_sums, int total_tokens, int K, + int group_size) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + const int num_groups_k = K / group_size; + q->parallel_for>( + sycl::range<2>(static_cast(total_tokens), static_cast(num_groups_k)), + [=](sycl::id<2> id) { + const int token = static_cast(id[0]); + const int g = static_cast(id[1]); + const ScalarT* row = activations + static_cast(token) * K + static_cast(g) * group_size; + // Split accumulators + a 16-wide vector load, mirroring the GEMV's own + // activation access pattern. + float s0 = 0.0f; + float s1 = 0.0f; + constexpr int SUB = 16; + using ActVec = sycl::vec; + int k = 0; + const int end = (group_size / SUB) * SUB; + for (; k < end; k += SUB) { + const ActVec av = *reinterpret_cast(row + k); +#pragma unroll + for (int u = 0; u < SUB; u += 2) { + s0 += static_cast(sycl::bit_cast(static_cast(av[u]))); + s1 += static_cast(sycl::bit_cast(static_cast(av[u + 1]))); + } + } + for (; k < group_size; ++k) { + s0 += static_cast(row[k]); + } + a_sums[static_cast(token) * num_groups_k + g] = s0 + s1; + }); +} + +// Convenience wrapper: fetch the activation-sum table from the scratch pool and +// (re)compute it for this call's activations. +template +float* compute_act_group_sums(sycl::queue* q, const ScalarT* activations, int total_tokens, int K, + int group_size) { + const int num_groups_k = K / group_size; + const size_t bytes = static_cast(total_tokens) * static_cast(num_groups_k) * sizeof(float); + float* a_sums = reinterpret_cast(act_group_sum_pool().acquire(q, bytes)); + launch_act_group_sums(q, activations, a_sums, total_tokens, K, group_size); + return a_sums; +} + // ---------------------------------------------------------------------------- // INT4 (S4_CLIP) GEMV with group-wise dequantization. // @@ -259,20 +428,22 @@ void launch_fp(sycl::queue* q, const ScalarT* activations, const ScalarT* weight // always used, and well inside the kernel's existing quantization tolerance. // // The per-group scale and zero-point are NOT applied here: this accumulates the -// raw integer-weighted dot product into `acc_q0`/`acc_q1` and the plain -// activation sum into `acc_a` (now needed by both modes, since sym carries the -// constant zero-point of 8). The caller folds the group's scale/zero in once +// raw integer-weighted dot product into `acc_q0`/`acc_q1`. The plain activation +// sum `Σ a` that the fold also needs (for both modes -- sym carries the +// constant zero-point of 8) is *not* accumulated here either: it is independent +// of the output column and is precomputed once per (token, K-group) by +// `launch_act_group_sums`. The caller folds the group's scale/zero in once // (sum a*((q-z)*s) == s*(sum a*q - z*sum a)). Hoisting the scale removes one -// float multiply per K element on the decode hot path, and because the fold is -// exact-once per group the result stays well within the kernel's existing -// quantization tolerance. +// float multiply per K element on the decode hot path, hoisting `Σ a` removes +// one float add per K element, and because the fold is exact-once per group the +// result stays well within the kernel's existing quantization tolerance. // // Two independent partial accumulators (``acc_q0``/``acc_q1``) break the // single fp32 dependency chain so the FMA pipeline is not latency-bound; the -// caller reduces the pair. ``acc_a`` reuses the same split. +// caller reduces the pair. template -static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc_q0, float& acc_q1, - float& acc_a) { +static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc_q0, + float& acc_q1) { static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); // sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so a single @@ -304,7 +475,6 @@ static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_pt const float fa1 = static_cast(a1); acc_q0 += fa0 * static_cast(q0); acc_q1 += fa1 * static_cast(q1); - acc_a += fa0 + fa1; } } } @@ -328,6 +498,10 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig const int num_groups_k = K / group_size; const int k_packed = K / 2; // bytes of packed weight per (expert, n) + // Per-(token, K-group) activation sums, shared by every lane and every + // N-tile instead of being recomputed inside the inner loop. + const float* a_sums = compute_act_group_sums(q, activations, total_tokens, K, group_size); + sycl::range<2> global{static_cast(total_tokens), static_cast(n_tiles * SG_SIZE)}; sycl::range<2> local{1, static_cast(SG_SIZE)}; @@ -341,6 +515,7 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig const int expert = expert_id_per_token[token]; const ScalarT* act_row = activations + static_cast(token) * K; + const float* a_sum_row = a_sums + static_cast(token) * num_groups_k; const uint8_t* w_row = weights + (static_cast(expert) * N + static_cast(n_global)) * k_packed; @@ -370,28 +545,28 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // on the fast path. // // The scale and zero are constant across the group, so the wide - // stages accumulate the raw integer-weighted dot product + // stages accumulate only the raw integer-weighted dot product // ``Σ a·q`` (split across two partial accumulators to break the - // fp32 dependency chain) plus ``Σ a``, and the fold below applies - // the scale/zero exactly once per group, identically for both modes: - // acc += scale * ((acc_q0 + acc_q1) - zero * acc_a) + // fp32 dependency chain); ``Σ a`` comes from the precomputed + // per-(token, group) table. The fold below applies the scale/zero + // exactly once per group, identically for both modes: + // acc += scale * ((acc_q0 + acc_q1) - zero * a_sum) // with `zero` == the per-group zero-point (asym) or the constant 8 // that the `^0x88` sign-flip decode introduces (sym). float acc_q0 = 0.0f; float acc_q1 = 0.0f; - float acc_a = 0.0f; int kk = 0; constexpr int CHUNK32 = 32; const int end32 = (group_size / CHUNK32) * CHUNK32; for (; kk < end32; kk += CHUNK32) { int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, acc_q0, - acc_q1, acc_a); + acc_q1); } constexpr int CHUNK16 = 16; const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; for (; kk < end16; kk += CHUNK16) { int4_decode_chunk(act_row + k_base + kk, w_row + (k_base + kk) / 2, acc_q0, - acc_q1, acc_a); + acc_q1); } // Scalar tail for group_size not divisible by 16. Uses the same // raw-accumulation convention as the wide stages so the single @@ -405,15 +580,39 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig const float fa1 = static_cast(act_row[k_base + kk + 1]); acc_q0 += fa0 * static_cast(q0); acc_q1 += fa1 * static_cast(q1); - acc_a += fa0 + fa1; } - acc += scale * ((acc_q0 + acc_q1) - zero * acc_a); + acc += scale * ((acc_q0 + acc_q1) - zero * a_sum_row[g]); } outputs[static_cast(token) * N + n_global] = static_cast(acc); }); } +// ---------------------------------------------------------------------------- +// Opt-in reuse of the int4 weight repack across calls. +// +// The repack output depends only on the weight buffer, which does not change +// between decode steps of a real inference loop, so in principle it can be +// built once and reused. The pool tag is a *pointer identity*, though, and a +// freed-then-reallocated weight tensor can land on the address of the previous +// one (torch's caching allocator makes this common in test loops that build a +// fresh packed tensor of the same shape per iteration). Reusing a stale repack +// would then silently produce wrong results, so this is off by default and must +// be enabled explicitly by a caller that owns the weight lifetime: +// +// ARK_MOE_DECODE_INT4_REPACK_CACHE=1 +// +// `ark::moe_decode_release_scratch()` (exposed to Python as +// `moe_decode_release_scratch`) drops the cached buffers. +// ---------------------------------------------------------------------------- +inline bool moe_decode_int4_repack_cache_enabled() { + const char* env = std::getenv("ARK_MOE_DECODE_INT4_REPACK_CACHE"); + if (env == nullptr) return false; // default OFF -- see comment above + std::string s(env); + for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); + return !(s == "0" || s == "false" || s == "off" || s == "no"); +} + // ---------------------------------------------------------------------------- // INT4 (S4_CLIP) coalesced-load GEMV. // @@ -426,19 +625,32 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // step issues 16 scattered transactions instead of one contiguous cache line. // // This path fixes that by first repacking the weights on-device into an -// N-tiled layout `[E, N/16, K/2, 16]`: the trailing dim of 16 holds one packed -// byte for each of the 16 columns owned by a sub-group tile, so at a fixed -// packed-byte index the 16 lanes read 16 contiguous bytes -> a single coalesced -// load. The dequant math is byte-for-byte identical to `launch_int4` (same -// `decode_int4_pair`, same per-group scale/zero fold), only the weight memory -// access pattern changes. The repack buffer is a transient USM device -// allocation freed after the queue drains; the caller's `[E, N, K/2]` weight +// N-tiled, 4-byte-blocked layout `[E, N/16, ceil(K/8), 16, 4]`: a *chunk* holds +// four consecutive packed bytes for each of the 16 columns owned by a sub-group +// tile, lane-major. Lane `l` therefore reads its four bytes at chunk offset +// `l*4`, and the 16 lanes of a sub-group together cover 64 contiguous bytes -> +// still a single coalesced transaction, but now each lane issues one +// `vec` load instead of four separate byte loads. For sym the +// `^0x88` sign flip that lets the unsigned (vectorizable) nibble decode stand in +// for the sign-extending one also becomes a single vector XOR per chunk instead +// of one scalar XOR per byte -- the last remaining per-byte instruction sym paid +// over asym in this kernel. The dequant math is otherwise identical to +// `launch_int4` (same `decode_int4_pair`, same per-group scale/zero fold); only +// the weight memory layout changes, and the caller's `[E, N, K/2]` weight // contract is unchanged. // -// The trailing lane stride means each lane's own K-bytes are 16 apart, so the -// vectorized `int4_decode_chunk` (contiguous per-lane load) does not apply -// here; the inner loop reads one packed byte per lane per step, which the -// hardware coalesces across the sub-group into one wide transaction. +// Group sizes that are a multiple of 8 (16/32/64/128/256 -- the shipped quant +// configs) start every K-group on a chunk boundary, so the vectorized stage +// covers the whole group. Other even group sizes are handled by a scalar +// prologue/epilogue around the vector stage, which reads the same layout one +// byte at a time. +// +// The repack buffer comes from the persistent per-queue scratch pool +// (`DeviceScratchPool`), so decode steady state performs no USM allocation and +// -- unlike the previous transient allocation, which had to be freed behind a +// blocking `queue::wait()` on every call -- introduces no host-side +// synchronization. The repack kernel itself still runs per call unless the +// caller opts into `ARK_MOE_DECODE_INT4_REPACK_CACHE`. // // On top of coalescing, this path blocks tokens: each work-item owns one // output column but processes up to `TOKEN_BLOCK` consecutive tokens. For each @@ -448,8 +660,7 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // hitting the same expert -- this amortizes the dominant weight traffic across // the block (GEMV -> small GEMM). Fully-scattered routing degrades gracefully // to one pass per token with the same per-pass weight reads as before, so the -// result is bit-identical to the one-token-per-work-item kernel regardless of -// routing. +// result is independent of routing. // ---------------------------------------------------------------------------- template void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, @@ -470,165 +681,211 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin const int n_tiles = N / N_TILE; const int num_groups_k = K / group_size; const int k_packed = K / 2; // bytes of packed weight per (expert, n) - - const size_t repacked_bytes = - static_cast(num_experts) * static_cast(n_tiles) * - static_cast(k_packed) * static_cast(N_TILE); - uint8_t* repacked = sycl::malloc_device(repacked_bytes, *q); - if (repacked == nullptr) { - throw std::runtime_error("moe_gemm_decode(int4): failed to allocate repack buffer"); - } - - // Repack kernel: one work-item per (expert, column, packed byte). The write - // index places the 16 columns of a tile contiguously in the trailing dim. - { + // Packed bytes are blocked by 4 along K so each lane can issue one 4-byte + // load. The last chunk is zero-padded when k_packed is not a multiple of 4. + constexpr int PACK_VEC = 4; + const int k_chunks = (k_packed + PACK_VEC - 1) / PACK_VEC; + const int chunk_stride = N_TILE * PACK_VEC; // bytes per (chunk) across the tile + + const size_t repacked_bytes = static_cast(num_experts) * static_cast(n_tiles) * + static_cast(k_chunks) * static_cast(chunk_stride); + // Reuse the repack across calls only when the caller opted in; the tag key + // folds in the full shape so a tensor of different dimensions cannot alias a + // cached repack that happens to sit at the same address. + const size_t repack_key = (static_cast(num_experts) * 1000003u + static_cast(N)) * 1000003u + + static_cast(k_packed); + bool repack_cached = false; + uint8_t* repacked = int4_repack_pool().acquire(q, repacked_bytes, weights, repack_key, + moe_decode_int4_repack_cache_enabled(), &repack_cached); + + // Repack kernel: one work-item per (expert, column, packed-byte slot). The + // write index places the 16 columns of a tile contiguously in chunks of 4 + // bytes, lane-major. Slots past `k_packed` are zero-filled so the padded tail + // of the last chunk is always initialized. + if (!repack_cached) { sycl::range<3> rp_global{static_cast(num_experts), static_cast(N), - static_cast(k_packed)}; + static_cast(k_chunks * PACK_VEC)}; q->parallel_for>(rp_global, [=](sycl::id<3> id) { const int e = static_cast(id[0]); const int n = static_cast(id[1]); const int kb = static_cast(id[2]); const int t = n / N_TILE; const int l = n % N_TILE; - const size_t src = (static_cast(e) * N + static_cast(n)) * k_packed + kb; - const size_t dst = - ((static_cast(e) * n_tiles + t) * k_packed + kb) * N_TILE + l; - repacked[dst] = weights[src]; + const int c = kb / PACK_VEC; + const int r = kb % PACK_VEC; + const size_t dst = ((static_cast(e) * n_tiles + t) * k_chunks + c) * chunk_stride + + static_cast(l) * PACK_VEC + r; + if (kb < k_packed) { + repacked[dst] = weights[(static_cast(e) * N + static_cast(n)) * k_packed + kb]; + } else { + repacked[dst] = 0; + } }); } + // Per-(token, K-group) activation sums, hoisted out of the inner loop. + const float* a_sums = compute_act_group_sums(q, activations, total_tokens, K, group_size); + sycl::range<2> global{static_cast((total_tokens + TOKEN_BLOCK - 1) / TOKEN_BLOCK), static_cast(n_tiles * SG_SIZE)}; sycl::range<2> local{1, static_cast(SG_SIZE)}; q->parallel_for>( - sycl::nd_range<2>(global, local), - [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { - const int token_base = static_cast(it.get_global_id(0)) * TOKEN_BLOCK; - const int n_tile = static_cast(it.get_group(1)); - const int lane = static_cast(it.get_local_id(1)); - const int n_global = n_tile * N_TILE + lane; - - // Number of tokens this work-item owns (last block may be short). - int block = TOKEN_BLOCK; - if (token_base + block > total_tokens) { - block = total_tokens - token_base; - } - - // Experts routed by each token in the block. The tile weight byte is - // loaded once per k-step and reused only for tokens whose expert - // matches the byte's owning expert, so blocking tokens that share an - // expert amortizes the dominant weight traffic; tokens with a - // different expert contribute nothing from this pass and are handled - // by the pass whose leader expert matches theirs. - int experts[TOKEN_BLOCK]; - for (int b = 0; b < block; ++b) { - experts[b] = expert_id_per_token[token_base + b]; - } - - // Which distinct experts appear in this block. For each we make one - // weight-streaming pass, reusing every loaded byte across all tokens - // in the block routed to that expert. Bursty routing collapses to a - // single pass; fully-scattered routing degrades to one pass per token - // (i.e. the previous behaviour) with no extra weight reads per pass. - for (int lead = 0; lead < block; ++lead) { - const int expert = experts[lead]; - // Skip experts already streamed by an earlier token in this block. - bool seen = false; - for (int p = 0; p < lead; ++p) { - if (experts[p] == expert) { - seen = true; - break; - } - } - if (seen) continue; - - // Base of this (expert, n_tile) weight tile in the repacked buffer. - // Layout [E, N/16, K/2, 16]; this lane reads byte kb at - // w_tile[kb*16 + lane], so adjacent lanes read adjacent bytes. - const uint8_t* w_tile = - repacked + ((static_cast(expert) * n_tiles + n_tile) * k_packed) * N_TILE; - const ScalarT* s_row = - scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; - const ScalarT* z_row = Asym - ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k - : nullptr; - - // Compact the tokens routed to `expert` into a dense member list - // once per pass. Hoisting the routing filter out of the hot k-loop - // removes a per-(kb, token) branch and lets the compiler keep the - // per-member activation base pointers in registers; the numerics - // are identical to the previous per-kb `experts[b] != expert` - // filter. - int members[TOKEN_BLOCK]; - const ScalarT* act_rows[TOKEN_BLOCK]; - int nmembers = 0; - for (int b = 0; b < block; ++b) { - if (experts[b] != expert) continue; - members[nmembers] = b; - act_rows[nmembers] = activations + static_cast(token_base + b) * K; - ++nmembers; - } - - float acc[TOKEN_BLOCK]; - for (int m = 0; m < nmembers; ++m) acc[m] = 0.0f; - - for (int g = 0; g < num_groups_k; ++g) { - const float scale = static_cast(s_row[g]); - // Constant zero-point of 8 for sym (see `int4_decode_chunk`): the - // `^0x88` sign-flip lets sym reuse the asym unsigned decode and - // fold, so both modes emit the identical instruction stream. - float zero = 8.0f; - if constexpr (Asym) { - zero = static_cast(z_row[g]); - } - const int k_base = g * group_size; - // Per-token split accumulators; the per-group scale/zero is folded - // once after the K-loop, exactly as in the scalar path. Each - // iteration processes two K-elements (one packed byte); the byte - // load is coalesced across the sub-group and reused across every - // token in the block routed to `expert`. - float acc_q0[TOKEN_BLOCK]; - float acc_q1[TOKEN_BLOCK]; - float acc_a[TOKEN_BLOCK]; - for (int m = 0; m < nmembers; ++m) { - acc_q0[m] = 0.0f; - acc_q1[m] = 0.0f; - acc_a[m] = 0.0f; - } - const int kb_base = k_base / 2; - const int kb_count = group_size / 2; - for (int kb = 0; kb < kb_count; ++kb) { - uint8_t packed = w_tile[(kb_base + kb) * N_TILE + lane]; - if constexpr (!Asym) packed ^= static_cast(0x88); - int q0, q1; - decode_int4_pair(packed, q0, q1); - const float fq0 = static_cast(q0); - const float fq1 = static_cast(q1); - const int k0 = k_base + 2 * kb; - for (int m = 0; m < nmembers; ++m) { - const ScalarT* act_row = act_rows[m]; - const float fa0 = static_cast(act_row[k0]); - const float fa1 = static_cast(act_row[k0 + 1]); - acc_q0[m] += fa0 * fq0; - acc_q1[m] += fa1 * fq1; - acc_a[m] += fa0 + fa1; - } - } - for (int m = 0; m < nmembers; ++m) { - acc[m] += scale * ((acc_q0[m] + acc_q1[m]) - zero * acc_a[m]); - } - } + sycl::nd_range<2>(global, local), + [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { + const int token_base = static_cast(it.get_global_id(0)) * TOKEN_BLOCK; + const int n_tile = static_cast(it.get_group(1)); + const int lane = static_cast(it.get_local_id(1)); + const int n_global = n_tile * N_TILE + lane; + + // Number of tokens this work-item owns (last block may be short). + int block = TOKEN_BLOCK; + if (token_base + block > total_tokens) { + block = total_tokens - token_base; + } - for (int m = 0; m < nmembers; ++m) { - const int b = members[m]; - outputs[static_cast(token_base + b) * N + n_global] = static_cast(acc[m]); - } - } - }); + // Experts routed by each token in the block. The tile weight byte is + // loaded once per k-step and reused only for tokens whose expert + // matches the byte's owning expert, so blocking tokens that share an + // expert amortizes the dominant weight traffic; tokens with a + // different expert contribute nothing from this pass and are handled + // by the pass whose leader expert matches theirs. + int experts[TOKEN_BLOCK]; + for (int b = 0; b < block; ++b) { + experts[b] = expert_id_per_token[token_base + b]; + } - q->wait(); - sycl::free(repacked, *q); + // Which distinct experts appear in this block. For each we make one + // weight-streaming pass, reusing every loaded byte across all tokens + // in the block routed to that expert. Bursty routing collapses to a + // single pass; fully-scattered routing degrades to one pass per token + // (i.e. the previous behaviour) with no extra weight reads per pass. + for (int lead = 0; lead < block; ++lead) { + const int expert = experts[lead]; + // Skip experts already streamed by an earlier token in this block. + bool seen = false; + for (int p = 0; p < lead; ++p) { + if (experts[p] == expert) { + seen = true; + break; + } + } + if (seen) continue; + + // Base of this (expert, n_tile) weight tile in the repacked buffer. + // Layout [E, N/16, ceil(K/8), 16, 4]; this lane reads packed byte + // `b_abs` at w_tile[(b_abs/4)*64 + lane*4 + b_abs%4], so the 16 lanes + // of the sub-group span 64 contiguous bytes per chunk. + const uint8_t* w_tile = + repacked + (static_cast(expert) * n_tiles + n_tile) * k_chunks * chunk_stride; + const ScalarT* s_row = + scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + const ScalarT* z_row = Asym + ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k + : nullptr; + + // Compact the tokens routed to `expert` into a dense member list + // once per pass. Hoisting the routing filter out of the hot k-loop + // removes a per-(kb, token) branch and lets the compiler keep the + // per-member activation base pointers in registers; the numerics + // are identical to the previous per-kb `experts[b] != expert` + // filter. + int members[TOKEN_BLOCK]; + const ScalarT* act_rows[TOKEN_BLOCK]; + const float* a_sum_rows[TOKEN_BLOCK]; + int nmembers = 0; + for (int b = 0; b < block; ++b) { + if (experts[b] != expert) continue; + members[nmembers] = b; + act_rows[nmembers] = activations + static_cast(token_base + b) * K; + a_sum_rows[nmembers] = a_sums + static_cast(token_base + b) * num_groups_k; + ++nmembers; + } + + float acc[TOKEN_BLOCK]; + for (int m = 0; m < nmembers; ++m) acc[m] = 0.0f; + + for (int g = 0; g < num_groups_k; ++g) { + const float scale = static_cast(s_row[g]); + // Constant zero-point of 8 for sym (see `int4_decode_chunk`): the + // `^0x88` sign-flip lets sym reuse the asym unsigned decode and + // fold, so both modes emit the identical instruction stream. + float zero = 8.0f; + if constexpr (Asym) { + zero = static_cast(z_row[g]); + } + const int k_base = g * group_size; + // Per-token split accumulators for the raw integer-weighted dot + // product; the per-group scale/zero (and the precomputed activation + // sum) are folded once after the K-loop, exactly as in the scalar + // path. + float acc_q0[TOKEN_BLOCK]; + float acc_q1[TOKEN_BLOCK]; + for (int m = 0; m < nmembers; ++m) { + acc_q0[m] = 0.0f; + acc_q1[m] = 0.0f; + } + + // Decode one already-sign-biased packed byte (two K elements) and + // accumulate it into every token of this pass. + auto accumulate_byte = [&](uint8_t biased, int k0) { + int q0, q1; + decode_int4_pair(biased, q0, q1); + const float fq0 = static_cast(q0); + const float fq1 = static_cast(q1); + for (int m = 0; m < nmembers; ++m) { + const ScalarT* act_row = act_rows[m]; + acc_q0[m] += static_cast(act_row[k0]) * fq0; + acc_q1[m] += static_cast(act_row[k0 + 1]) * fq1; + } + }; + // Load a single packed byte through the chunked layout. + auto load_byte = [&](int b_abs) { + uint8_t packed = w_tile[static_cast(b_abs / PACK_VEC) * chunk_stride + + static_cast(lane) * PACK_VEC + (b_abs % PACK_VEC)]; + if constexpr (!Asym) packed ^= static_cast(0x88); + return packed; + }; + + const int kb_base = k_base / 2; + const int kb_count = group_size / 2; + int kb = 0; + // Prologue to the next 4-byte chunk boundary. Empty whenever + // group_size % 8 == 0, i.e. for every shipped quant config. + for (; kb < kb_count && ((kb_base + kb) % PACK_VEC) != 0; ++kb) { + accumulate_byte(load_byte(kb_base + kb), k_base + 2 * kb); + } + using PackVec = sycl::vec; + for (; kb + PACK_VEC <= kb_count; kb += PACK_VEC) { + const int b_abs = kb_base + kb; // 4-byte aligned here + PackVec pv = *reinterpret_cast( + w_tile + static_cast(b_abs / PACK_VEC) * chunk_stride + + static_cast(lane) * PACK_VEC); + if constexpr (!Asym) { + // One vector XOR flips the sign bit of all 8 nibbles at once. + pv = pv ^ PackVec(static_cast(0x88)); + } +#pragma unroll + for (int u = 0; u < PACK_VEC; ++u) { + accumulate_byte(pv[u], k_base + 2 * (kb + u)); + } + } + // Scalar tail for group sizes that are not a multiple of 8. + for (; kb < kb_count; ++kb) { + accumulate_byte(load_byte(kb_base + kb), k_base + 2 * kb); + } + + for (int m = 0; m < nmembers; ++m) { + acc[m] += scale * ((acc_q0[m] + acc_q1[m]) - zero * a_sum_rows[m][g]); + } + } + + for (int m = 0; m < nmembers; ++m) { + const int b = members[m]; + outputs[static_cast(token_base + b) * N + n_global] = static_cast(acc[m]); + } + } + }); } // ---------------------------------------------------------------------------- @@ -1031,6 +1288,19 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh } // namespace moe_decode_detail +// ---------------------------------------------------------------------------- +// Release every device scratch buffer the int4 decode fallbacks hold (the +// N-tiled weight repack and the activation-sum table). Both are served from +// grow-on-demand per-queue slabs that are normally kept for the lifetime of the +// process; call this to hand the memory back, or to drop a repack cached under +// `ARK_MOE_DECODE_INT4_REPACK_CACHE` before the underlying weight buffer is +// freed. Safe to call at any time -- the next decode simply reallocates. +// ---------------------------------------------------------------------------- +inline void moe_decode_release_scratch() { + moe_decode_detail::int4_repack_pool().release_all(); + moe_decode_detail::act_group_sum_pool().release_all(); +} + // ---------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_DECODE_DPAS_S4` (default ON). When ON, int4-sym // (S4_CLIP, !asym) decode is routed to the dedicated decode-phase S4 DPAS From 4ec9d7fe3b6f775dd300147431a45af333b03877 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 06:32:40 +0000 Subject: [PATCH 38/47] test,docs: cover the new int4 decode paths and document the tuning knobs Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/test/README_MOE_PREFILL_PERF.md | 79 ++++++++++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 62 ++++++++ auto_round_extension/ark/test/test_moe.py | 135 +++++++++++++++++- .../ark/test/test_moe_decode_perf.py | 80 ++++++++++- 4 files changed, 350 insertions(+), 6 deletions(-) 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 435ae28392..2d624c2c0c 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -363,6 +363,85 @@ sym now accumulates the biased sum and subtracts `8 * sum a` at the end, which is exactly the fp32 accumulation pattern asym has always used. It applies to both `launch_int4` and `launch_int4_coalesced`. +**4-byte-blocked coalesced repack.** The coalesced fallback +(`launch_int4_coalesced`, `ARK_MOE_DECODE_COALESCE_INT4` default ON) +repacks the `[E, N, K/2]` weights on-device so sub-group loads are +contiguous. The original repack layout `[E, N/16, K/2, 16]` put one byte +per lane per step, so although the 16 lanes together covered one cache +line, each lane still issued a *byte* load — and for sym, one scalar +`^ 0x88` per byte. The layout is now `[E, N/16, ceil(K/8), 16, 4]`: a +chunk holds four consecutive packed bytes for each of the 16 columns of a +tile, lane-major, so lane `l` reads its four bytes at chunk offset `l*4` +and the sub-group still spans 64 contiguous bytes. Each lane therefore +issues one `vec` load instead of four byte loads (4× fewer +weight-load instructions), and the sym sign flip becomes one vector XOR +per chunk instead of four scalar XORs — removing the last per-byte +instruction sym paid over asym in this kernel. Group sizes that are a +multiple of 8 (16/32/64/128/256 — every shipped quant config) start each +K-group on a chunk boundary so the vector stage covers the whole group; +other even group sizes fall back to a scalar prologue/epilogue over the +same layout. The external `[E, N, K/2]` weight contract is unchanged. + +**Hoisted activation sums.** Both int4 GEMVs fold the per-group +scale/zero as `scale * (Σ a·q − zero · Σ a)`. `Σ a` depends only on the +activation row and the K-group, not on the output column, yet it used to +be recomputed inside the inner loop — once per sub-group lane (16× +redundant) and again for every N-tile work-group — costing one extra float +add per K element. It is now precomputed once into a +`[total_tokens, K/group_size]` fp32 table (`launch_act_group_sums`), so +the GEMV inner loop only accumulates `Σ a·q` and reads one float per +group. Sym benefits most: for sym `Σ a` exists purely to carry the +constant zero-point of 8 that the `^ 0x88` decode introduces. The +summation order changes by a few fp32 ULPs, far inside the kernel's +quantization tolerance. + +**Pooled scratch instead of per-call `malloc_device`.** The repack buffer +used to be a transient USM allocation that had to be freed behind a +blocking `queue::wait()` on *every* decode call — and decode issues one +call per generated token, so that allocation plus sync was on the order of +the GEMV itself. Both the repack buffer and the activation-sum table now +come from a persistent per-queue, grow-on-demand slab +(`DeviceScratchPool`), so steady-state decode performs no allocation and +introduces no host-side synchronization; ordering between the producer +kernels and the GEMV is already guaranteed by the in-order queue. +`ark.moe_decode_release_scratch()` (pybind `moe_decode_release_scratch`) +hands the memory back. + +The repack *kernel* still runs on every call by default. Setting +`ARK_MOE_DECODE_INT4_REPACK_CACHE=1` reuses the previous repack when the +weight buffer address and shape are unchanged, which is valid for a real +inference loop where the weights are fixed. It is **off by default** +because the tag is a pointer identity: a freed-then-reallocated weight +tensor can land on the same address (torch's caching allocator makes this +common in test loops), and a stale repack would silently produce wrong +results. Callers that enable it must call +`ark.moe_decode_release_scratch()` before dropping the weight tensor. + +| Env var | Default | Effect | +| ------- | ------- | ------ | +| `ARK_MOE_DECODE_COALESCE_INT4` | ON | Use the coalesced, 4-byte-blocked repack GEMV for the int4 scalar fallback; `0` forces the legacy per-lane-strided `launch_int4`. | +| `ARK_MOE_DECODE_COALESCE_MIN_TOKENS` | `num_experts * TOKEN_BLOCK` | Minimum total tokens before the coalesced kernel is worth its repack pass; `0` disables the gate (what the parity/A-B tests set). | +| `ARK_MOE_DECODE_INT4_REPACK_CACHE` | OFF | Reuse the repack across calls on the same weight buffer. Only safe when the caller owns the weight lifetime. | + +Perf A/B for the coalesced path is +`test_moe_decode_perf.py::test_perf_int4_coalesced_vs_strided` (toggles +`ARK_MOE_DECODE_COALESCE_INT4` 0/1 on the same shapes). Correctness is +covered by `test_moe.py::test_decode_int4_coalesced_matches_scalar`, +`::test_decode_int4_coalesced_token_blocking`, +`::test_decode_int4_coalesced_unaligned_group_size` (group sizes that are +not a multiple of 8, exercising the scalar prologue/epilogue) and +`::test_decode_int4_repack_cache`. + +**Occupancy-gate threshold sweep.** The default +`ARK_MOE_DECODE_DPAS_S4_MIN_TPE` of 8 was derived from the row count of +`dpas_w4a16_policy_m_8` rather than measured. The sweep that locates the +real crossing point is +`test_moe_decode_perf.py::test_perf_int4_sym_dpas_vs_scalar_threshold`; +its default token counts (16–128) all sit far below the gate (8 × 192 +experts == 1536 tokens), so pass `--all-shapes` to extend the sweep to +256/512/1024/1536/3072 tokens and bracket the gate from both sides. The +default stays at 8 until hardware numbers say otherwise. + Accuracy parity is covered by `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group`, which forces `ARK_MOE_PREFILL_DPAS_S4=1` + 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 8e43d7aae4..fabb2948eb 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 @@ -284,6 +284,68 @@ sym 使用 `(int8_t)(byte << 4) >> 4` 逐 nibble 做符号扩展，这是一条 这正是 asym 一直在用的 fp32 累加方式；`launch_int4` 与 `launch_int4_coalesced` 两个 kernel 均已应用。 +**按 4 字节分块的 coalesced repack。** coalesced 回退路径 +(`launch_int4_coalesced`，`ARK_MOE_DECODE_COALESCE_INT4` 默认开启)会在设备端把 +`[E, N, K/2]` 权重重排，使 sub-group 的加载连续。原先的重排布局 +`[E, N/16, K/2, 16]` 每个 lane 每步只放一个字节，因此虽然 16 个 lane 合起来覆盖 +一条 cache line，每个 lane 仍然发出的是*字节*加载 —— 对 sym 而言还额外附带每字节 +一次标量 `^ 0x88`。现在布局改为 `[E, N/16, ceil(K/8), 16, 4]`：一个 chunk 为 +tile 内 16 列中的每一列存放 4 个连续的打包字节，按 lane 主序排列，因此 lane `l` +在 chunk 偏移 `l*4` 处读取自己的 4 个字节，sub-group 整体仍然覆盖 64 个连续字节。 +于是每个 lane 只需一次 `vec` 加载而不是四次字节加载(权重加载指令数 +降为 1/4)，sym 的符号翻转也变成每个 chunk 一次向量 XOR 而不是四次标量 XOR —— +这消除了该 kernel 中 sym 相对 asym 仅存的逐字节额外指令。group_size 为 8 的倍数时 +(16/32/64/128/256，即全部已发布的量化配置)每个 K 组都从 chunk 边界开始，向量 +阶段覆盖整个组；其他偶数 group_size 则通过标量前导/收尾循环在同一布局上处理。 +对外的 `[E, N, K/2]` 权重约定保持不变。 + +**提取激活求和。** 两个 int4 GEMV 都按 +`scale * (Σ a·q − zero · Σ a)` 折叠每组的 scale/zero。`Σ a` 只依赖激活行与 K 组， +与输出列无关，但此前它是在内层循环里重复计算的 —— 每个 sub-group lane 算一遍 +(16 倍冗余)，每个 N-tile work-group 再算一遍 —— 每个 K 元素多付出一次浮点加法。 +现在它被预先计算成一张 `[total_tokens, K/group_size]` 的 fp32 表 +(`launch_act_group_sums`)，GEMV 内层循环只累加 `Σ a·q`，每组读取一个 float。 +sym 获益最大：对 sym 来说 `Σ a` 的存在纯粹是为了承载 `^ 0x88` 解码引入的常数 +zero-point 8。求和顺序的变化仅带来几个 fp32 ULP 的差异，远在 kernel 现有的 +量化容差之内。 + +**用 scratch 池替代每次调用的 `malloc_device`。** repack 缓冲区原本是临时的 USM +分配，每次 decode 调用都必须在一次阻塞的 `queue::wait()` 之后释放 —— 而 decode +每生成一个 token 就调用一次，因此这次分配加同步的开销已经与 GEMV 本身同量级。 +现在 repack 缓冲区与激活求和表都取自按 queue 持有、按需增长的常驻 slab +(`DeviceScratchPool`)，稳态 decode 不再有任何分配，也不引入主机侧同步；生产者 +kernel 与 GEMV 之间的顺序由 in-order queue 保证。 +`ark.moe_decode_release_scratch()`(pybind `moe_decode_release_scratch`)可将内存 +归还。 + +repack *kernel* 默认仍每次调用都执行。设置 +`ARK_MOE_DECODE_INT4_REPACK_CACHE=1` 可在权重缓冲区地址与形状不变时复用上一次的 +repack 结果 —— 这对权重固定的真实推理循环是成立的。它**默认关闭**，因为其 tag +是指针身份：被释放后重新分配的权重张量可能落在同一地址(torch 的缓存分配器在 +测试循环中很容易出现这种情况)，此时陈旧的 repack 会静默产生错误结果。启用它的 +调用方必须在丢弃权重张量之前调用 `ark.moe_decode_release_scratch()`。 + +| 环境变量 | 默认值 | 作用 | +| -------- | ------ | ---- | +| `ARK_MOE_DECODE_COALESCE_INT4` | 开启 | int4 标量回退使用按 4 字节分块的 coalesced repack GEMV；设为 `0` 则强制使用按 lane 跨步的旧版 `launch_int4`。 | +| `ARK_MOE_DECODE_COALESCE_MIN_TOKENS` | `num_experts * TOKEN_BLOCK` | coalesced kernel 值回其 repack 开销所需的最小总 token 数；设为 `0` 关闭该门控(一致性/A-B 测试即如此设置)。 | +| `ARK_MOE_DECODE_INT4_REPACK_CACHE` | 关闭 | 在同一权重缓冲区上跨调用复用 repack 结果。仅当调用方掌握权重生命周期时才安全。 | + +coalesced 路径的性能 A/B 见 +`test_moe_decode_perf.py::test_perf_int4_coalesced_vs_strided`(在相同形状上切换 +`ARK_MOE_DECODE_COALESCE_INT4` 0/1)。正确性由 +`test_moe.py::test_decode_int4_coalesced_matches_scalar`、 +`::test_decode_int4_coalesced_token_blocking`、 +`::test_decode_int4_coalesced_unaligned_group_size`(非 8 的倍数的 group_size， +覆盖标量前导/收尾路径)以及 `::test_decode_int4_repack_cache` 覆盖。 + +**占用率门控阈值扫描。** `ARK_MOE_DECODE_DPAS_S4_MIN_TPE` 的默认值 8 来自 +`dpas_w4a16_policy_m_8` 的 tile 行数，而非实测结果。定位真实交叉点的扫描用例是 +`test_moe_decode_perf.py::test_perf_int4_sym_dpas_vs_scalar_threshold`；它默认的 +token 数(16–128)都远低于该门控(8 × 192 个专家 == 1536 个 token)，因此需要传入 +`--all-shapes` 把扫描扩展到 256/512/1024/1536/3072 个 token，从两侧夹住门控。 +在拿到硬件数据之前，默认值仍保持为 8。 + 精度对齐由 `test_moe_prefill_accuracy.py::test_accuracy_int4_dpas_per_group` 覆盖,该用例强制 `ARK_MOE_PREFILL_DPAS_S4=1` + diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 77b76c0575..e1eb623531 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -584,8 +584,10 @@ def test_decode_int4_coalesced_matches_scalar(self, monkeypatch, dtype, asym, gr dequant->bmm reference within quantization tolerance. The S4 DPAS fast path is disabled so both runs exercise the scalar - fallback (this is the only path the coalesce flag affects). Shapes use - N%16==0 so the N-tiled repack is exact. + fallback (this is the only path the coalesce flag affects), and the + coalesce amortization gate is disabled so the coalesced kernel really + runs at this (tiny) token count instead of silently falling back to the + per-lane-strided one. Shapes use N%16==0 so the N-tiled repack is exact. """ num_experts = 4 tokens_per_expert = [1, 0, 2, 1] @@ -619,8 +621,11 @@ def _run(): asym=asym, ) - # Force the scalar-GEMV fallback so the coalesce flag actually applies. + # Force the scalar-GEMV fallback so the coalesce flag actually applies, + # and disable the tokens-per-expert amortization gate so the coalesced + # kernel is reached at this token count. monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS", "0") monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") out_coalesced = _run() @@ -684,8 +689,11 @@ def _run(): asym=asym, ) - # Force the scalar-GEMV fallback so the coalesce/token-blocking path runs. + # Force the scalar-GEMV fallback so the coalesce/token-blocking path + # runs, and disable the amortization gate so the coalesced kernel is + # reached at this token count. monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS", "0") monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") out_blocked = _run() @@ -699,6 +707,125 @@ def _run(): # blocked and legacy kernels must match tightly for every routing shape. torch.testing.assert_close(out_blocked, out_scalar, rtol=1e-3, atol=1e-3) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("asym", [False, True]) + @pytest.mark.parametrize("group_size", [4, 12]) + def test_decode_int4_coalesced_unaligned_group_size(self, monkeypatch, dtype, asym, group_size): + """The coalesced int4 kernel loads four packed bytes per lane out of a + ``[E, N/16, ceil(K/8), 16, 4]`` layout. Group sizes that are a multiple + of 8 start every K-group on a 4-byte chunk boundary, so the vectorized + stage covers the whole group; other even group sizes need the scalar + prologue/epilogue around it. + + ``group_size=4`` leaves every odd group misaligned with no room for a + vector step at all, and ``group_size=12`` mixes a misaligned prologue + with a vector step, so between them both non-vector paths are covered. + """ + num_experts = 3 + tokens_per_expert = [2, 0, 3] + total_tokens = sum(tokens_per_expert) + N, K = 32, 48 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + if asym: + zeros = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_asym(w_float, scales, zeros, group_size) + dequant = _dequant_int4_asym(packed, scales, zeros, group_size).to(dtype) + else: + zeros = None + packed = _pack_int4_sym(w_float, scales, group_size) + dequant = _dequant_int4_sym(packed, scales, group_size).to(dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + zeros=zeros, + weight_bits=4, + group_size=group_size, + asym=asym, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS", "0") + + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") + out_coalesced = _run() + + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "0") + out_scalar = _run() + + assert out_coalesced.shape == (total_tokens, N) + torch.testing.assert_close(out_coalesced, ref, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(out_coalesced, out_scalar, rtol=1e-3, atol=1e-3) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_decode_int4_repack_cache(self, monkeypatch, dtype): + """``ARK_MOE_DECODE_INT4_REPACK_CACHE=1`` lets the coalesced int4 kernel + reuse the N-tiled weight repack across calls instead of rebuilding it. + + Reuse is keyed on the weight buffer address plus its shape, so it is + only valid while the caller keeps that buffer alive. Verify that (a) + repeated calls on a live weight tensor keep matching the reference, and + (b) ``moe_decode_release_scratch()`` drops the cache so a *different* + weight tensor -- which torch's caching allocator may well hand back at + the same address -- is repacked again rather than answered from the + stale entry. + """ + num_experts = 3 + tokens_per_expert = [4, 0, 5] + total_tokens = sum(tokens_per_expert) + N, K = 64, 128 + group_size = 32 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + def _build(): + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_sym(w_float, scales, group_size) + dequant = _dequant_int4_sym(packed, scales, group_size).to(dtype) + return packed, scales, _moe_decode_reference(activations, dequant, num_tokens_per_expert) + + def _run(packed, scales): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS", "0") + monkeypatch.setenv("ARK_MOE_DECODE_INT4_REPACK_CACHE", "1") + + try: + packed_a, scales_a, ref_a = _build() + # First call builds the repack, second must hit the cache. + torch.testing.assert_close(_run(packed_a, scales_a), ref_a, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(_run(packed_a, scales_a), ref_a, rtol=5e-2, atol=5e-2) + + # Drop the cached repack before the buffer it was derived from goes + # away, then verify a freshly built weight tensor is honoured. + del packed_a, scales_a + ark.moe_decode_release_scratch() + packed_b, scales_b, ref_b = _build() + torch.testing.assert_close(_run(packed_b, scales_b), ref_b, rtol=5e-2, atol=5e-2) + finally: + ark.moe_decode_release_scratch() + def test_decode_validation_errors(self): """Sanity-check that Python-side validation catches misuse.""" num_experts = 2 diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index 26641cc95c..60861b407e 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -234,6 +234,15 @@ def _spread_tokens(total_tokens: int, num_experts: int = 192) -> list: # to locate the DPAS-vs-scalar crossover for the auto-dispatch threshold. _INT4_THRESHOLD_TOKEN_COUNTS = [16, 32, 64, 128] +# Extra (much larger) token counts appended when --all-shapes is passed. The +# default ARK_MOE_DECODE_DPAS_S4_MIN_TPE gate is 8 tokens per expert, i.e. +# 8 * 192 == 1536 total tokens for the sweep's expert count, so the counts above +# alone can never show where DPAS actually overtakes the scalar GEMV -- they all +# sit far below the gate. These bracket the gate from both sides so the measured +# crossing point can replace the tile-row-count heuristic the default was +# derived from. +_INT4_THRESHOLD_TOKEN_COUNTS_EXTENDED = [256, 512, 1024, 1536, 3072] + # MiniMax-M2 up/down-proj (N, K) pairs reused by the threshold sweep. _INT4_THRESHOLD_NK = [ (1536, 3072), # gate/up-proj @@ -395,6 +404,65 @@ def test_perf_int4(self, dtype, asym): ) _print_row(label, N, K, total_tokens, base_ms, ark_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("asym", [False, True]) + def test_perf_int4_coalesced_vs_strided(self, monkeypatch, dtype, asym): + """int4 scalar-GEMV fallback A/B: coalesced N-tiled weight loads + (``ARK_MOE_DECODE_COALESCE_INT4=1``) vs the legacy per-lane-strided + kernel (``=0``). + + ``speedup`` is ``strided / coalesced`` (the coalesced kernel is the + "ark" column). The S4 DPAS fast path is disabled so both columns run the + scalar fallback -- the only path the coalesce flag affects -- and the + amortization gate is disabled so the coalesced kernel is actually + reached at decode-sized token counts (it repacks the whole weight + tensor, so at very low tokens-per-expert the repack dominates, which is + exactly what the default gate exists to avoid; this row shows how much). + """ + group_size = 128 + kind = "asym" if asym else "sym" + _print_header( + f"INT4 {kind} coalesced vs strided (group_size={group_size}, " + f"act={str(dtype).split('.')[-1]}) -- strided GEMV (baseline) vs coalesced GEMV (ark)" + ) + for label, E, tpe, N, K in DECODE_SHAPES: + if K % group_size != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + if asym: + zeros = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_int4_asym(w_float, scales, zeros, group_size) + else: + zeros = None + packed = _pack_int4_sym(w_float, scales, group_size) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + zeros=zeros, + weight_bits=4, + group_size=group_size, + asym=asym, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4", "0") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS", "0") + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "0") + strided_ms = _xpu_time_ms(_run) + monkeypatch.setenv("ARK_MOE_DECODE_COALESCE_INT4", "1") + coalesced_ms = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_MOE_DECODE_COALESCE_MIN_TOKENS", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_COALESCE_INT4", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_S4", raising=False) + _print_row(label, N, K, total_tokens, strided_ms, coalesced_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_perf_int4_sym_dpas_vs_scalar(self, monkeypatch, dtype): """int4-sym decode: compare the S4 DPAS path (ARK_MOE_DECODE_DPAS_S4=1) @@ -450,7 +518,7 @@ def _run(): _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) - def test_perf_int4_sym_dpas_vs_scalar_threshold(self, monkeypatch, dtype): + def test_perf_int4_sym_dpas_vs_scalar_threshold(self, request, monkeypatch, dtype): """int4-sym decode threshold sweep: DPAS vs scalar GEMV across a range of total-token counts (16/32/64/128) at ``group_size=32``. @@ -465,9 +533,17 @@ def test_perf_int4_sym_dpas_vs_scalar_threshold(self, monkeypatch, dtype): the default ``ARK_MOE_DECODE_DPAS_S4_MIN_TPE`` occupancy gate (set to ``0`` here so the DPAS column is not itself re-routed to the scalar GEMV). + + By default only the small token counts are swept so a CI pass stays + short. Those all sit below the default occupancy gate (8 tokens per + expert == 1536 tokens for E=192), so pass ``--all-shapes`` to extend the + sweep across the gate and measure where the crossing actually is. """ group_size = 32 E = 192 + token_counts = list(_INT4_THRESHOLD_TOKEN_COUNTS) + if request.config.getoption("--all-shapes", default=False): + token_counts += _INT4_THRESHOLD_TOKEN_COUNTS_EXTENDED _print_header( f"INT4 sym DPAS vs scalar threshold sweep (group_size={group_size}, " f"act={str(dtype).split('.')[-1]}) -- scalar GEMV (baseline) vs S4 DPAS (ark)" @@ -475,7 +551,7 @@ def test_perf_int4_sym_dpas_vs_scalar_threshold(self, monkeypatch, dtype): for N, K in _INT4_THRESHOLD_NK: if K % group_size != 0 or N % 64 != 0 or K % 32 != 0: continue - for total_tokens in _INT4_THRESHOLD_TOKEN_COUNTS: + for total_tokens in token_counts: tpe = _spread_tokens(total_tokens, E) label = f"int4 {N}x{K} t{total_tokens}" activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") From 610a005d669e6e0ef469df782ba6bafcbbc782e0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 07:25:46 +0000 Subject: [PATCH 39/47] perf: speed up int4-sym MoE decode with word-native nibble decode Decode packed int4 nibbles from 32-bit words instead of 8-bit-typed sycl::vec ops, and restore true signed-nibble decode for sym so the zero-point/activation-sum fold (and its extra per-call device kernel launch) is skipped entirely on the sym path. Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 303 ++++++++++-------- .../wrapper/include/sycl_tla_moe_dequant.hpp | 28 +- 2 files changed, 186 insertions(+), 145 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index bc128bda1d..9ae797dfb6 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -39,6 +39,7 @@ #pragma once +#include #include #include #include @@ -86,6 +87,31 @@ namespace moe_decode_detail { constexpr int SG_SIZE = 16; constexpr int N_TILE = SG_SIZE; // one output element per sub-group lane +// ---------------------------------------------------------------------------- +// Allocation-free boolean env-var lookup. +// +// The int4 decode dispatch consults up to three of these on *every* call -- they +// are deliberately re-read rather than cached so tests and benchmarks can toggle +// a path in-process -- and decode issues one call per generated token. Building +// a `std::string` per lookup put a heap allocation on that hot path for nothing, +// so the comparison is done in place instead. The accepted spellings are +// unchanged: "0" / "false" / "off" / "no" (case-insensitive) mean off, any other +// value means on, and an unset variable falls back to `default_value`. +// ---------------------------------------------------------------------------- +inline bool env_flag_enabled(const char* name, bool default_value) { + const char* env = std::getenv(name); + if (env == nullptr) return default_value; + auto iequals = [](const char* value, const char* lowercase_literal) { + const char* a = value; + const char* b = lowercase_literal; + for (; *a != '\0' && *b != '\0'; ++a, ++b) { + if (static_cast(std::tolower(static_cast(*a))) != *b) return false; + } + return *a == '\0' && *b == '\0'; + }; + return !(iequals(env, "0") || iequals(env, "false") || iequals(env, "off") || iequals(env, "no")); +} + // Token-blocking factor for the coalesced int4 decode GEMV. A work-item that // owns one (n_tile, lane) output column processes up to TOKEN_BLOCK consecutive // tokens, loading each packed weight byte from the (expert, n_tile) tile once @@ -136,6 +162,7 @@ using moe_dequant::decode_fp8_e4m3_lut; using moe_dequant::decode_fp8_e5m2_bits; using moe_dequant::decode_fp8_e5m2_lut; using moe_dequant::decode_int2_quad; +using moe_dequant::decode_int4_octet; using moe_dequant::decode_int4_pair; using moe_dequant::decode_int8; using moe_dequant::fp8_decode_use_lut; @@ -324,22 +351,25 @@ inline DeviceScratchPool& act_group_sum_pool() { } // ---------------------------------------------------------------------------- -// Per-(token, K-group) activation sums. +// Per-(token, K-group) activation sums (asym int4 only). // -// Both int4 GEMV kernels fold their per-group scale/zero as +// The asym int4 GEMVs fold their per-group scale/zero as // `scale * (Σ a·q - zero · Σ a)`, where `Σ a` runs over the group's K range. // `Σ a` depends only on the activation row and the group, *not* on the output // column, yet the GEMVs used to recompute it inside the inner loop -- once per // sub-group lane (16x redundant) and again for every N-tile work-group (N/16x -// redundant). That cost one extra float add per K element on the hot path, -// which for sym is pure overhead introduced by the constant zero-point of 8 -// that the `^0x88` decode relies on. +// redundant). That cost one extra float add per K element on the hot path. // // This pass computes the `[total_tokens, K/group_size]` table once, so the // GEMVs only accumulate `Σ a·q` and read one float per group. The table is // tiny (tokens x groups floats) and comes from the scratch pool, so no // allocation happens in steady state. // +// Sym does *not* use this at all: it decodes true signed nibbles, so its fold +// carries no zero-point term. That keeps this extra kernel launch -- a +// first-order cost when the GEMV itself is only tens of microseconds -- off the +// sym decode timeline entirely. +// // The summation order differs from the previous in-loop accumulation, so // results move by a few float ULPs -- far inside the kernel's quantization // tolerance. @@ -409,34 +439,36 @@ float* compute_act_group_sums(sycl::queue* q, const ScalarT* activations, int to // without regressing group_size == 16 (which drops straight to the 16-wide // stage). // -// Both sym and asym decode the nibbles through the *same* unsigned path. -// Measurements on MiniMax-M2 decode shapes showed int4-sym ~1.9x slower than -// int4-asym in this very kernel even though sym does strictly fewer floating -// point operations. The only difference was the per-nibble sign extension -// (`(int8_t)(byte << 4) >> 4`) -- a serial shift/narrow/shift chain per nibble -// that defeats the byte-wise vectorization the asym mask+shift form gets. The -// sym decode is therefore expressed with the standard sign-flip identity +// The packed weights are consumed as 32-bit *words* (four packed bytes, eight +// K elements) through the shared `decode_int4_octet` primitive rather than as +// a `sycl::vec` byte vector. On Xe the ALU is 32-bit-lane based +// and byte-typed vector operations lower to restricted byte regioning that IGC +// often has to expand, so every per-byte step in the hot loop -- the element +// extraction *and*, for sym, the sign handling -- paid that expansion. In +// word form both modes issue exactly two native DWORD operations per nibble: // -// signed_nibble == (unsigned_nibble ^ 8) - 8 +// asym: (word >> 4j) & 0xF +// sym : (int)(word << (28 - 4j)) >> 28 // -// so XOR-ing the packed byte with `0x88` (flipping the sign bit of *both* -// nibbles at once, on the whole loaded vector register) turns sym into exactly -// the asym computation with a constant zero-point of 8. The decoded integers -// are bit-identical to the sign-extending decode for all 256 byte values, so -// the only change is that sym now accumulates the biased sum and subtracts -// `8 * sum a` at the end -- exactly the fp32 accumulation pattern asym has -// always used, and well inside the kernel's existing quantization tolerance. +// so sym's sign extension is no longer a serial byte-typed shift/narrow/shift +// chain and costs the same as asym's mask+shift. That removes the reason the +// previous revision biased sym with a `^0x88` vector XOR and folded a constant +// zero-point of 8: sym now accumulates *true signed* nibbles, which means it +// no longer needs the `Σ a` term at all (see `launch_int4`) -- one fewer fp32 +// add per K element, one fewer table read per K-group, and one fewer device +// kernel launch per decode call than asym. The decoded integers are +// bit-identical to `decode_int4_pair` for every input word, so decode/prefill +// parity is unchanged. // -// The per-group scale and zero-point are NOT applied here: this accumulates the -// raw integer-weighted dot product into `acc_q0`/`acc_q1`. The plain activation -// sum `Σ a` that the fold also needs (for both modes -- sym carries the -// constant zero-point of 8) is *not* accumulated here either: it is independent -// of the output column and is precomputed once per (token, K-group) by -// `launch_act_group_sums`. The caller folds the group's scale/zero in once -// (sum a*((q-z)*s) == s*(sum a*q - z*sum a)). Hoisting the scale removes one -// float multiply per K element on the decode hot path, hoisting `Σ a` removes -// one float add per K element, and because the fold is exact-once per group the -// result stays well within the kernel's existing quantization tolerance. +// The per-group scale and zero-point are NOT applied here: this accumulates +// the raw integer-weighted dot product into `acc_q0`/`acc_q1`. The caller +// folds the group's scale (and, for asym, its zero-point against the +// precomputed `Σ a`) in exactly once per group: +// sym : acc += scale * (acc_q0 + acc_q1) +// asym: acc += scale * ((acc_q0 + acc_q1) - zero * Σ a) +// Hoisting the scale removes one float multiply per K element on the decode +// hot path, and because the fold is exact-once per group the result stays well +// within the kernel's existing quantization tolerance. // // Two independent partial accumulators (``acc_q0``/``acc_q1``) break the // single fp32 dependency chain so the FMA pipeline is not latency-bound; the @@ -449,32 +481,29 @@ static inline void int4_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_pt // sycl::vec only supports widths of 1, 2, 3, 4, 8 or 16, so a single // vec load is illegal. Process the chunk in 16-wide sub-blocks // (16 activations + 8 packed weight bytes each), which keeps CHUNK == 32 - // valid while reusing the same code path for CHUNK == 16. + // valid while reusing the same code path for CHUNK == 16. The 8 packed bytes + // are loaded as two 32-bit words in one 8-byte transaction -- the same + // access width (and the same 8-byte alignment requirement) as the byte + // vector it replaces. constexpr int SUB = 16; + constexpr int WORDS = SUB / 8; // one 32-bit word per 8 K elements using ActVec = sycl::vec; - using PackVec = sycl::vec; + using WordVec = sycl::vec; #pragma unroll for (int s = 0; s < CHUNK / SUB; ++s) { const ActVec av = *reinterpret_cast(act_ptr + s * SUB); - PackVec pv = *reinterpret_cast(w_ptr + s * (SUB / 2)); - if constexpr (!Asym) { - // Sign-flip the whole packed vector in one vector XOR so the sym nibbles - // can be decoded by the (vectorizable) unsigned path below; the constant - // zero-point of 8 is folded by the caller. - pv = pv ^ PackVec(static_cast(0x88)); - } + const WordVec wv = *reinterpret_cast(w_ptr + s * (SUB / 2)); #pragma unroll - for (int b = 0; b < SUB / 2; ++b) { - int q0, q1; - // Always the unsigned decode: asym nibbles are unsigned by definition and - // sym nibbles were biased by the XOR above. - decode_int4_pair(pv[b], q0, q1); - const ScalarT a0 = sycl::bit_cast(static_cast(av[2 * b])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[2 * b + 1])); - const float fa0 = static_cast(a0); - const float fa1 = static_cast(a1); - acc_q0 += fa0 * static_cast(q0); - acc_q1 += fa1 * static_cast(q1); + for (int w = 0; w < WORDS; ++w) { + int q[8]; + decode_int4_octet(wv[w], q); +#pragma unroll + for (int u = 0; u < 8; u += 2) { + const ScalarT a0 = sycl::bit_cast(static_cast(av[8 * w + u])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[8 * w + u + 1])); + acc_q0 += static_cast(a0) * static_cast(q[u]); + acc_q1 += static_cast(a1) * static_cast(q[u + 1]); + } } } } @@ -499,8 +528,16 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig const int k_packed = K / 2; // bytes of packed weight per (expert, n) // Per-(token, K-group) activation sums, shared by every lane and every - // N-tile instead of being recomputed inside the inner loop. - const float* a_sums = compute_act_group_sums(q, activations, total_tokens, K, group_size); + // N-tile instead of being recomputed inside the inner loop. Only the *asym* + // fold needs them (`Σ a·(q - z) == Σ a·q - z·Σ a`): sym decodes true signed + // nibbles, so its fold is a plain per-group scale multiply with no + // zero-point term. Skipping the pre-pass keeps a whole extra kernel launch + // off the sym decode timeline -- on decode-sized batches the GEMV itself is + // only tens of microseconds, so an extra dispatch is a first-order cost. + [[maybe_unused]] const float* a_sums = nullptr; + if constexpr (Asym) { + a_sums = compute_act_group_sums(q, activations, total_tokens, K, group_size); + } sycl::range<2> global{static_cast(total_tokens), static_cast(n_tiles * SG_SIZE)}; sycl::range<2> local{1, static_cast(SG_SIZE)}; @@ -515,26 +552,21 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig const int expert = expert_id_per_token[token]; const ScalarT* act_row = activations + static_cast(token) * K; - const float* a_sum_row = a_sums + static_cast(token) * num_groups_k; const uint8_t* w_row = weights + (static_cast(expert) * N + static_cast(n_global)) * k_packed; const ScalarT* s_row = scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; - const ScalarT* z_row = Asym - ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k - : nullptr; + [[maybe_unused]] const ScalarT* z_row = nullptr; + [[maybe_unused]] const float* a_sum_row = nullptr; + if constexpr (Asym) { + z_row = zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + a_sum_row = a_sums + static_cast(token) * num_groups_k; + } float acc = 0.0f; for (int g = 0; g < num_groups_k; ++g) { const float scale = static_cast(s_row[g]); - // Sym uses the constant zero-point of 8 that the `^0x88` sign-flip - // in the decode introduces, so both modes run the identical fold. - // `if constexpr` keeps the null `z_row` out of the sym instantiation. - float zero = 8.0f; - if constexpr (Asym) { - zero = static_cast(z_row[g]); - } const int k_base = g * group_size; // Vectorized ladder: process 32 K-elements at a time (16 packed // weight bytes + vec activation block), then a 16-wide @@ -547,12 +579,11 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // The scale and zero are constant across the group, so the wide // stages accumulate only the raw integer-weighted dot product // ``Σ a·q`` (split across two partial accumulators to break the - // fp32 dependency chain); ``Σ a`` comes from the precomputed - // per-(token, group) table. The fold below applies the scale/zero - // exactly once per group, identically for both modes: - // acc += scale * ((acc_q0 + acc_q1) - zero * a_sum) - // with `zero` == the per-group zero-point (asym) or the constant 8 - // that the `^0x88` sign-flip decode introduces (sym). + // fp32 dependency chain). The fold below applies the scale exactly + // once per group; asym additionally subtracts its per-group + // zero-point against the precomputed ``Σ a``: + // sym : acc += scale * (acc_q0 + acc_q1) + // asym: acc += scale * ((acc_q0 + acc_q1) - zero * a_sum) float acc_q0 = 0.0f; float acc_q1 = 0.0f; int kk = 0; @@ -572,16 +603,19 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // raw-accumulation convention as the wide stages so the single // scale/zero fold below stays valid. for (; kk < group_size; kk += 2) { - uint8_t packed = w_row[(k_base + kk) / 2]; - if constexpr (!Asym) packed ^= static_cast(0x88); + const uint8_t packed = w_row[(k_base + kk) / 2]; int q0, q1; - decode_int4_pair(packed, q0, q1); + decode_int4_pair(packed, q0, q1); const float fa0 = static_cast(act_row[k_base + kk]); const float fa1 = static_cast(act_row[k_base + kk + 1]); acc_q0 += fa0 * static_cast(q0); acc_q1 += fa1 * static_cast(q1); } - acc += scale * ((acc_q0 + acc_q1) - zero * a_sum_row[g]); + if constexpr (Asym) { + acc += scale * ((acc_q0 + acc_q1) - static_cast(z_row[g]) * a_sum_row[g]); + } else { + acc += scale * (acc_q0 + acc_q1); + } } outputs[static_cast(token) * N + n_global] = static_cast(acc); @@ -606,11 +640,7 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // `moe_decode_release_scratch`) drops the cached buffers. // ---------------------------------------------------------------------------- inline bool moe_decode_int4_repack_cache_enabled() { - const char* env = std::getenv("ARK_MOE_DECODE_INT4_REPACK_CACHE"); - if (env == nullptr) return false; // default OFF -- see comment above - std::string s(env); - for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - return !(s == "0" || s == "false" || s == "off" || s == "no"); + return env_flag_enabled("ARK_MOE_DECODE_INT4_REPACK_CACHE", false); // default OFF -- see comment above } // ---------------------------------------------------------------------------- @@ -629,15 +659,14 @@ inline bool moe_decode_int4_repack_cache_enabled() { // four consecutive packed bytes for each of the 16 columns owned by a sub-group // tile, lane-major. Lane `l` therefore reads its four bytes at chunk offset // `l*4`, and the 16 lanes of a sub-group together cover 64 contiguous bytes -> -// still a single coalesced transaction, but now each lane issues one -// `vec` load instead of four separate byte loads. For sym the -// `^0x88` sign flip that lets the unsigned (vectorizable) nibble decode stand in -// for the sign-extending one also becomes a single vector XOR per chunk instead -// of one scalar XOR per byte -- the last remaining per-byte instruction sym paid -// over asym in this kernel. The dequant math is otherwise identical to -// `launch_int4` (same `decode_int4_pair`, same per-group scale/zero fold); only -// the weight memory layout changes, and the caller's `[E, N, K/2]` weight -// contract is unchanged. +// still a single coalesced transaction, but now each lane issues one 32-bit +// word load instead of four separate byte loads. That word is decoded with the +// shared `decode_int4_octet` primitive, so all eight nibbles are extracted with +// native DWORD shift/mask pairs and neither mode touches the 8-bit ALU (see +// `int4_decode_chunk`). The dequant math is otherwise identical to +// `launch_int4` (bit-identical nibbles, same per-group scale/zero fold, sym +// accumulating true signed nibbles with no `Σ a` term); only the weight memory +// layout changes, and the caller's `[E, N, K/2]` weight contract is unchanged. // // Group sizes that are a multiple of 8 (16/32/64/128/256 -- the shipped quant // configs) start every K-group on a chunk boundary, so the vectorized stage @@ -723,8 +752,13 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin }); } - // Per-(token, K-group) activation sums, hoisted out of the inner loop. - const float* a_sums = compute_act_group_sums(q, activations, total_tokens, K, group_size); + // Per-(token, K-group) activation sums, hoisted out of the inner loop. Only + // asym needs them -- sym decodes true signed nibbles and folds a plain scale + // -- so the sym path skips this kernel launch entirely (see `launch_int4`). + [[maybe_unused]] const float* a_sums = nullptr; + if constexpr (Asym) { + a_sums = compute_act_group_sums(q, activations, total_tokens, K, group_size); + } sycl::range<2> global{static_cast((total_tokens + TOKEN_BLOCK - 1) / TOKEN_BLOCK), static_cast(n_tiles * SG_SIZE)}; @@ -780,9 +814,10 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin repacked + (static_cast(expert) * n_tiles + n_tile) * k_chunks * chunk_stride; const ScalarT* s_row = scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; - const ScalarT* z_row = Asym - ? zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k - : nullptr; + [[maybe_unused]] const ScalarT* z_row = nullptr; + if constexpr (Asym) { + z_row = zeros + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + } // Compact the tokens routed to `expert` into a dense member list // once per pass. Hoisting the routing filter out of the hot k-loop @@ -792,13 +827,15 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin // filter. int members[TOKEN_BLOCK]; const ScalarT* act_rows[TOKEN_BLOCK]; - const float* a_sum_rows[TOKEN_BLOCK]; + [[maybe_unused]] const float* a_sum_rows[TOKEN_BLOCK]; int nmembers = 0; for (int b = 0; b < block; ++b) { if (experts[b] != expert) continue; members[nmembers] = b; act_rows[nmembers] = activations + static_cast(token_base + b) * K; - a_sum_rows[nmembers] = a_sums + static_cast(token_base + b) * num_groups_k; + if constexpr (Asym) { + a_sum_rows[nmembers] = a_sums + static_cast(token_base + b) * num_groups_k; + } ++nmembers; } @@ -807,18 +844,11 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin for (int g = 0; g < num_groups_k; ++g) { const float scale = static_cast(s_row[g]); - // Constant zero-point of 8 for sym (see `int4_decode_chunk`): the - // `^0x88` sign-flip lets sym reuse the asym unsigned decode and - // fold, so both modes emit the identical instruction stream. - float zero = 8.0f; - if constexpr (Asym) { - zero = static_cast(z_row[g]); - } const int k_base = g * group_size; // Per-token split accumulators for the raw integer-weighted dot - // product; the per-group scale/zero (and the precomputed activation - // sum) are folded once after the K-loop, exactly as in the scalar - // path. + // product; the per-group scale (and, for asym, the zero-point + // against the precomputed activation sum) are folded once after + // the K-loop, exactly as in the scalar path. float acc_q0[TOKEN_BLOCK]; float acc_q1[TOKEN_BLOCK]; for (int m = 0; m < nmembers; ++m) { @@ -826,11 +856,9 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin acc_q1[m] = 0.0f; } - // Decode one already-sign-biased packed byte (two K elements) and - // accumulate it into every token of this pass. - auto accumulate_byte = [&](uint8_t biased, int k0) { - int q0, q1; - decode_int4_pair(biased, q0, q1); + // Accumulate one decoded nibble pair (two K elements) into every + // token of this pass. + auto accumulate_pair = [&](int q0, int q1, int k0) { const float fq0 = static_cast(q0); const float fq1 = static_cast(q1); for (int m = 0; m < nmembers; ++m) { @@ -839,12 +867,13 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin acc_q1[m] += static_cast(act_row[k0 + 1]) * fq1; } }; - // Load a single packed byte through the chunked layout. - auto load_byte = [&](int b_abs) { - uint8_t packed = w_tile[static_cast(b_abs / PACK_VEC) * chunk_stride + - static_cast(lane) * PACK_VEC + (b_abs % PACK_VEC)]; - if constexpr (!Asym) packed ^= static_cast(0x88); - return packed; + // Load and decode a single packed byte through the chunked layout. + auto accumulate_byte = [&](int b_abs, int k0) { + const uint8_t packed = w_tile[static_cast(b_abs / PACK_VEC) * chunk_stride + + static_cast(lane) * PACK_VEC + (b_abs % PACK_VEC)]; + int q0, q1; + decode_int4_pair(packed, q0, q1); + accumulate_pair(q0, q1, k0); }; const int kb_base = k_base / 2; @@ -853,30 +882,38 @@ void launch_int4_coalesced(sycl::queue* q, const ScalarT* activations, const uin // Prologue to the next 4-byte chunk boundary. Empty whenever // group_size % 8 == 0, i.e. for every shipped quant config. for (; kb < kb_count && ((kb_base + kb) % PACK_VEC) != 0; ++kb) { - accumulate_byte(load_byte(kb_base + kb), k_base + 2 * kb); + accumulate_byte(kb_base + kb, k_base + 2 * kb); } - using PackVec = sycl::vec; + // A lane's PACK_VEC == 4 bytes inside a chunk are contiguous, so + // they are exactly one little-endian 32-bit word: load it as such + // and decode all 8 nibbles with native DWORD ops (no 8-bit ALU, + // no sign-bias XOR for sym) via the shared octet primitive. for (; kb + PACK_VEC <= kb_count; kb += PACK_VEC) { const int b_abs = kb_base + kb; // 4-byte aligned here - PackVec pv = *reinterpret_cast( + const uint32_t word = *reinterpret_cast( w_tile + static_cast(b_abs / PACK_VEC) * chunk_stride + static_cast(lane) * PACK_VEC); - if constexpr (!Asym) { - // One vector XOR flips the sign bit of all 8 nibbles at once. - pv = pv ^ PackVec(static_cast(0x88)); - } + int qv[8]; + decode_int4_octet(word, qv); #pragma unroll for (int u = 0; u < PACK_VEC; ++u) { - accumulate_byte(pv[u], k_base + 2 * (kb + u)); + accumulate_pair(qv[2 * u], qv[2 * u + 1], k_base + 2 * (kb + u)); } } // Scalar tail for group sizes that are not a multiple of 8. for (; kb < kb_count; ++kb) { - accumulate_byte(load_byte(kb_base + kb), k_base + 2 * kb); + accumulate_byte(kb_base + kb, k_base + 2 * kb); } - for (int m = 0; m < nmembers; ++m) { - acc[m] += scale * ((acc_q0[m] + acc_q1[m]) - zero * a_sum_rows[m][g]); + if constexpr (Asym) { + const float zero = static_cast(z_row[g]); + for (int m = 0; m < nmembers; ++m) { + acc[m] += scale * ((acc_q0[m] + acc_q1[m]) - zero * a_sum_rows[m][g]); + } + } else { + for (int m = 0; m < nmembers; ++m) { + acc[m] += scale * (acc_q0[m] + acc_q1[m]); + } } } @@ -1323,12 +1360,7 @@ inline void moe_decode_release_scratch() { // benchmarks can toggle the path in-process. // ---------------------------------------------------------------------------- inline bool moe_decode_dpas_s4_enabled() { - const char* env = std::getenv("ARK_MOE_DECODE_DPAS_S4"); - if (env == nullptr) return true; // default ON - std::string s(env); - for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "0" || s == "false" || s == "off" || s == "no") return false; - return true; + return moe_decode_detail::env_flag_enabled("ARK_MOE_DECODE_DPAS_S4", true); // default ON } // ---------------------------------------------------------------------------- @@ -1374,12 +1406,7 @@ inline bool moe_decode_dpas_s4_occupancy_ok(int total_tokens, int num_experts) { // Re-read on every call so tests / benchmarks can toggle it in-process. // ---------------------------------------------------------------------------- inline bool moe_decode_coalesce_int4_enabled() { - const char* env = std::getenv("ARK_MOE_DECODE_COALESCE_INT4"); - if (env == nullptr) return true; // default ON - std::string s(env); - for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "0" || s == "false" || s == "off" || s == "no") return false; - return true; + return moe_decode_detail::env_flag_enabled("ARK_MOE_DECODE_COALESCE_INT4", true); // default ON } // ---------------------------------------------------------------------------- diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp index d142621fdc..ecc95f52c3 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp @@ -211,17 +211,31 @@ inline void decode_int2_quad(uint8_t packed, int q[4]) { // ... // q[7] = byte3 high nibble (k_base + 7) // -// The decoder is expressed as a `#pragma unroll` loop over `decode_int4_pair`, -// so it is bit-identical by construction to four scalar decodes of the same -// four bytes. This keeps the parity contract with the decode/GEMV path (which -// only ever calls `decode_int4_pair`) trivially satisfied. +// Collapsing that mapping, field `j` (K offset `j`) is simply bits +// `[4j+3 : 4j]` of the word, so every field can be extracted with a pair of +// *32-bit* ALU ops and the 8-bit datapath is never touched: +// asym: `(word >> 4j) & 0xF` +// sym : `(int)(word << (28 - 4j)) >> 28` -- park the nibble in the sign +// position, then arithmetic-shift it back down. This is exactly the +// 32-bit form of `int8_t(byte << 4) >> 4`, so the decoded integers are +// bit-identical to `decode_int4_pair` for all inputs and the +// decode/prefill parity contract is preserved. +// +// The 32-bit form matters on Xe: `sycl::vec` arithmetic and +// per-byte extraction lower to byte-typed regioning that IGC frequently has to +// expand, and that expansion is what made the sym sign-extension look +// inherently more expensive than the asym mask+shift. Both modes now issue the +// same two native DWORD operations per nibble. // ---------------------------------------------------------------------------- template inline void decode_int4_octet(uint32_t packed, int q[8]) { #pragma unroll - for (int i = 0; i < 4; ++i) { - const uint8_t byte = static_cast((packed >> (i * 8)) & 0xFFu); - decode_int4_pair(byte, q[2 * i], q[2 * i + 1]); + for (int j = 0; j < 8; ++j) { + if constexpr (Asym) { + q[j] = static_cast((packed >> (4 * j)) & 0xFu); + } else { + q[j] = static_cast(packed << (28 - 4 * j)) >> 28; + } } } From a2827941df6c5a9d8fa440d6d11848ed0411e7b7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 07:29:01 +0000 Subject: [PATCH 40/47] docs: update int4 decode perf notes for word-native nibble decode Rewrite the sym-decode, coalesced-repack and activation-sum sections of README_MOE_PREFILL_PERF.md and its CN counterpart to describe the 32-bit word decode path and the asym-only activation-sum pre-pass. Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/test/README_MOE_PREFILL_PERF.md | 92 +++++++++++-------- .../ark/test/README_MOE_PREFILL_PERF_CN.md | 58 +++++++----- 2 files changed, 90 insertions(+), 60 deletions(-) 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 2d624c2c0c..090bb3e1cf 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -346,54 +346,70 @@ expert. threshold; `0` disables the gate (always DPAS when the shape gate allows), which is what the accuracy and DPAS-vs-scalar perf tests set. -**Sym decodes through the asym nibble path.** Once both modes shared the -scalar GEMV, int4-sym was still ~1.9x slower than int4-asym in the *same* -kernel (2.83 ms vs 1.49 ms at bs32) despite doing strictly fewer floating -point operations. The only asymmetry was the nibble decode: sym -sign-extended each nibble with `(int8_t)(byte << 4) >> 4`, a serial -shift/narrow/shift chain per nibble, while asym used a plain mask+shift -that vectorizes over the whole loaded byte vector. Sym now uses the -sign-flip identity `signed == (unsigned ^ 8) - 8`: one vector XOR of the -packed bytes with `0x88` flips both nibbles' sign bits, after which sym is -*literally* the asym computation with a constant zero-point of 8 (same -unsigned decode, same `sum a` accumulator, same single per-group -scale/zero fold). The decoded integers are bit-identical to the previous -sign-extending decode for all 256 byte values; the only change is that -sym now accumulates the biased sum and subtracts `8 * sum a` at the end, -which is exactly the fp32 accumulation pattern asym has always used. It -applies to both `launch_int4` and `launch_int4_coalesced`. +**Word-native nibble decode; sym keeps its signed nibbles.** Once both +modes shared the scalar GEMV, int4-sym was still slower than int4-asym in +the *same* kernel despite doing strictly fewer floating point operations. +The asymmetry was the nibble decode, and the first attempt at fixing it +(the `^ 0x88` sign-flip identity `signed == (unsigned ^ 8) - 8`) did not +close the gap: it kept sym on 8-bit-typed operations — a `sycl::vec` +XOR plus per-byte mask/shift — which Xe expands into narrow-type ALU work +rather than executing on the native 32-bit datapath, and it forced sym to +carry a constant zero-point of 8 (see *activation sums* below). + +Both modes now decode through the shared `decode_int4_octet` primitive, +which takes the 8 nibbles of a packed *32-bit word* and extracts each one +with a single DWORD shift/mask pair (asym) or a DWORD shift-left + +arithmetic shift-right pair (sym). No 8-bit-typed vector, no XOR, no +narrowing casts, and one 32-bit load per 8 K elements instead of a byte +vector. The per-nibble results are bit-identical to `decode_int4_pair` for +every one of the 2^32 input words in both modes (verified exhaustively), so +this is a pure instruction-selection change. It applies to `launch_int4`, +`launch_int4_coalesced`, and — since the primitive is shared — the prefill +mixed-dtype path. + +Because sym once again recovers *true signed* nibbles, its per-group fold +collapses to `acc += scale * Σ a·q` with no zero-point term at all, whereas +asym keeps `acc += scale * (Σ a·q − zero · Σ a)`. **4-byte-blocked coalesced repack.** The coalesced fallback (`launch_int4_coalesced`, `ARK_MOE_DECODE_COALESCE_INT4` default ON) repacks the `[E, N, K/2]` weights on-device so sub-group loads are contiguous. The original repack layout `[E, N/16, K/2, 16]` put one byte per lane per step, so although the 16 lanes together covered one cache -line, each lane still issued a *byte* load — and for sym, one scalar -`^ 0x88` per byte. The layout is now `[E, N/16, ceil(K/8), 16, 4]`: a -chunk holds four consecutive packed bytes for each of the 16 columns of a -tile, lane-major, so lane `l` reads its four bytes at chunk offset `l*4` -and the sub-group still spans 64 contiguous bytes. Each lane therefore -issues one `vec` load instead of four byte loads (4× fewer -weight-load instructions), and the sym sign flip becomes one vector XOR -per chunk instead of four scalar XORs — removing the last per-byte -instruction sym paid over asym in this kernel. Group sizes that are a -multiple of 8 (16/32/64/128/256 — every shipped quant config) start each -K-group on a chunk boundary so the vector stage covers the whole group; -other even group sizes fall back to a scalar prologue/epilogue over the -same layout. The external `[E, N, K/2]` weight contract is unchanged. - -**Hoisted activation sums.** Both int4 GEMVs fold the per-group -scale/zero as `scale * (Σ a·q − zero · Σ a)`. `Σ a` depends only on the -activation row and the K-group, not on the output column, yet it used to -be recomputed inside the inner loop — once per sub-group lane (16× +line, each lane still issued a *byte* load. The layout is now +`[E, N/16, ceil(K/8), 16, 4]`: a chunk holds four consecutive packed bytes +for each of the 16 columns of a tile, lane-major, so lane `l` reads its +four bytes at chunk offset `l*4` and the sub-group still spans 64 +contiguous bytes. A lane's four bytes are contiguous, hence exactly one +little-endian 32-bit word: the lane issues a single DWORD load (4× fewer +weight-load instructions) and feeds it straight to `decode_int4_octet`, so +all eight nibbles come out with native 32-bit ops in both modes. Group +sizes that are a multiple of 8 (16/32/64/128/256 — every shipped quant +config) start each K-group on a chunk boundary so the vector stage covers +the whole group; other even group sizes fall back to a scalar +prologue/epilogue over the same layout. The external `[E, N, K/2]` weight +contract is unchanged. + +**Hoisted activation sums (asym only).** The asym int4 GEMVs fold the +per-group scale/zero as `scale * (Σ a·q − zero · Σ a)`. `Σ a` depends only +on the activation row and the K-group, not on the output column, yet it +used to be recomputed inside the inner loop — once per sub-group lane (16× redundant) and again for every N-tile work-group — costing one extra float add per K element. It is now precomputed once into a `[total_tokens, K/group_size]` fp32 table (`launch_act_group_sums`), so the GEMV inner loop only accumulates `Σ a·q` and reads one float per -group. Sym benefits most: for sym `Σ a` exists purely to carry the -constant zero-point of 8 that the `^ 0x88` decode introduces. The -summation order changes by a few fp32 ULPs, far inside the kernel's -quantization tolerance. +group. The summation order changes by a few fp32 ULPs, far inside the +kernel's quantization tolerance. + +**Sym skips the pre-pass entirely.** `launch_act_group_sums` is a separate +`parallel_for`, and on an in-order queue it fully serializes ahead of the +GEMV. That is a poor trade at decode sizes: it saves one float add per K +element in a loop that is already memory-bound, but adds a whole kernel +dispatch to a call whose GEMV is only tens of microseconds at bs1 — which +is why routing sym through the biased-unsigned decode made sym *slower*, +not faster. Now that sym decodes true signed nibbles it has no zero-point +term, so the table is computed (and the kernel launched) only when +`Asym` is true. **Pooled scratch instead of per-call `malloc_device`.** The repack buffer used to be a transient USM allocation that had to be freed behind a 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 fabb2948eb..201a2bb49b 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 @@ -270,44 +270,58 @@ int4-sym(DPAS)为 0.31–0.34 ms / 1.55 ms,而 int4-asym(标量 GEMV)为 "每专家 token 数" 阈值;设为 `0` 则关闭门控(只要形状门控通过就走 DPAS), 精度测试与 DPAS/标量 对比性能测试即使用该设置。 -**sym 直接复用 asym 的 nibble 解码路径。** 在两者都走标量 GEMV 之后， -int4-sym 在 *同一个* kernel 里仍比 int4-asym 慢约 1.9 倍（bs32：2.83 ms -vs 1.49 ms），尽管 sym 的浮点运算严格更少。唯一的差异在于 nibble 解码： -sym 使用 `(int8_t)(byte << 4) >> 4` 逐 nibble 做符号扩展，这是一条 -移位/截断/移位 的串行依赖链；而 asym 的 掩码+移位 形式可以在整个已加载 -的字节向量上向量化。现在 sym 改用符号翻转恒等式 -`signed == (unsigned ^ 8) - 8`：对打包字节做一次向量 `^ 0x88`，同时翻转 -两个 nibble 的符号位，之后 sym *就是* zero-point 恒为 8 的 asym 计算 -（相同的无符号解码、相同的 `sum a` 累加器、相同的每组一次 scale/zero -折叠）。对全部 256 种字节取值，解码出的整数与原来的符号扩展逐位相同， -唯一的变化是 sym 现在累加的是有偏置的和，最后再减去 `8 * sum a`， -这正是 asym 一直在用的 fp32 累加方式；`launch_int4` 与 `launch_int4_coalesced` -两个 kernel 均已应用。 +**基于 32 位字的 nibble 解码；sym 恢复真正的有符号 nibble。** 在两者都走 +标量 GEMV 之后，int4-sym 在 *同一个* kernel 里仍比 int4-asym 慢，尽管 sym +的浮点运算严格更少。差异在于 nibble 解码，而第一次尝试的修复（符号翻转 +恒等式 `signed == (unsigned ^ 8) - 8`，即 `^ 0x88`）并没有弥合差距：它让 +sym 仍然停留在 8 位类型的运算上 —— 一次 `sycl::vec` 的 XOR 加上 +逐字节的 掩码/移位 —— 而 Xe 会把这类窄类型运算展开，无法直接跑在原生 +32 位数据通路上；同时它还迫使 sym 携带一个恒为 8 的 zero-point（见下面的 +"激活求和"）。 + +现在两种模式都通过共享的 `decode_int4_octet` 原语解码：它接收一个打包的 +*32 位字* 中的 8 个 nibble，每个 nibble 只用一对 DWORD 移位/掩码（asym） +或一对 DWORD 左移 + 算术右移（sym）即可取出。没有 8 位类型的向量，没有 +XOR，没有窄化转换，并且每 8 个 K 元素只需一次 32 位加载而不是一次字节 +向量加载。在两种模式下，对全部 2^32 种输入字，逐 nibble 的结果都与 +`decode_int4_pair` 逐位相同（已穷举验证），因此这纯粹是指令选择层面的改动。 +它应用于 `launch_int4`、`launch_int4_coalesced`，并且由于该原语是共享的， +prefill 的混合精度路径同样受益。 + +由于 sym 重新恢复了 *真正的有符号* nibble，它的每组折叠退化为 +`acc += scale * Σ a·q`，完全没有 zero-point 项；asym 则仍是 +`acc += scale * (Σ a·q − zero · Σ a)`。 **按 4 字节分块的 coalesced repack。** coalesced 回退路径 (`launch_int4_coalesced`，`ARK_MOE_DECODE_COALESCE_INT4` 默认开启)会在设备端把 `[E, N, K/2]` 权重重排，使 sub-group 的加载连续。原先的重排布局 `[E, N/16, K/2, 16]` 每个 lane 每步只放一个字节，因此虽然 16 个 lane 合起来覆盖 -一条 cache line，每个 lane 仍然发出的是*字节*加载 —— 对 sym 而言还额外附带每字节 -一次标量 `^ 0x88`。现在布局改为 `[E, N/16, ceil(K/8), 16, 4]`：一个 chunk 为 +一条 cache line，每个 lane 仍然发出的是*字节*加载。现在布局改为 +`[E, N/16, ceil(K/8), 16, 4]`：一个 chunk 为 tile 内 16 列中的每一列存放 4 个连续的打包字节，按 lane 主序排列，因此 lane `l` 在 chunk 偏移 `l*4` 处读取自己的 4 个字节，sub-group 整体仍然覆盖 64 个连续字节。 -于是每个 lane 只需一次 `vec` 加载而不是四次字节加载(权重加载指令数 -降为 1/4)，sym 的符号翻转也变成每个 chunk 一次向量 XOR 而不是四次标量 XOR —— -这消除了该 kernel 中 sym 相对 asym 仅存的逐字节额外指令。group_size 为 8 的倍数时 +一个 lane 的这 4 个字节是连续的，因此恰好构成一个小端 32 位字：lane 只需发出 +一次 DWORD 加载(权重加载指令数降为 1/4)，并直接交给 `decode_int4_octet`， +两种模式下 8 个 nibble 都用原生 32 位运算取出。group_size 为 8 的倍数时 (16/32/64/128/256，即全部已发布的量化配置)每个 K 组都从 chunk 边界开始，向量 阶段覆盖整个组；其他偶数 group_size 则通过标量前导/收尾循环在同一布局上处理。 对外的 `[E, N, K/2]` 权重约定保持不变。 -**提取激活求和。** 两个 int4 GEMV 都按 +**提取激活求和(仅 asym)。** asym 的 int4 GEMV 按 `scale * (Σ a·q − zero · Σ a)` 折叠每组的 scale/zero。`Σ a` 只依赖激活行与 K 组， 与输出列无关，但此前它是在内层循环里重复计算的 —— 每个 sub-group lane 算一遍 (16 倍冗余)，每个 N-tile work-group 再算一遍 —— 每个 K 元素多付出一次浮点加法。 现在它被预先计算成一张 `[total_tokens, K/group_size]` 的 fp32 表 (`launch_act_group_sums`)，GEMV 内层循环只累加 `Σ a·q`，每组读取一个 float。 -sym 获益最大：对 sym 来说 `Σ a` 的存在纯粹是为了承载 `^ 0x88` 解码引入的常数 -zero-point 8。求和顺序的变化仅带来几个 fp32 ULP 的差异，远在 kernel 现有的 -量化容差之内。 +求和顺序的变化仅带来几个 fp32 ULP 的差异，远在 kernel 现有的量化容差之内。 + +**sym 完全跳过这一前置 pass。** `launch_act_group_sums` 是一个独立的 +`parallel_for`，在 in-order queue 上会完全串行地排在 GEMV 之前。对 decode +规模而言这笔交易并不划算：它在一个本就受访存带宽限制的循环里省下每个 K +元素一次浮点加法，却给一次 GEMV 仅几十微秒(bs1)的调用额外增加了一整次 +kernel 派发 —— 这正是让 sym 走"有偏置无符号解码"反而变*慢*的原因。现在 +sym 解码出真正的有符号 nibble，不含 zero-point 项，因此只有在 `Asym` 为真 +时才会计算该表(并派发该 kernel)。 **用 scratch 池替代每次调用的 `malloc_device`。** repack 缓冲区原本是临时的 USM 分配，每次 decode 调用都必须在一次阻塞的 `queue::wait()` 之后释放 —— 而 decode From 6cd81921310670f30ec3439a5448692eb59f18cb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:48:51 +0000 Subject: [PATCH 41/47] perf: word-native FP8 decode + decode-tuned FP8 DPAS dispatch Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 337 +++++++++++++----- .../wrapper/include/sycl_tla_moe_dequant.hpp | 133 +++++++ .../include/sycl_tla_moe_prefill_fp8_dpas.hpp | 111 ++++++ .../include/sycl_tla_moe_prefill_s4_dpas.hpp | 32 +- 4 files changed, 499 insertions(+), 114 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 9ae797dfb6..1c03de22ef 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -63,20 +63,25 @@ // ---------------------------------------------------------------------------- // FP8 decode implementation switch (runtime) // -// FP8 weight bytes can be dequantized either via inline bit manipulation or -// via the 128-entry magnitude LUT in `bestla/sycl/fp8_lut.h` (sign applied -// separately). Both paths are mathematically equivalent for finite values; -// pick whichever is faster on the target hardware. +// FP8 weight bytes can be dequantized three ways, all mathematically equivalent +// for the values a real checkpoint contains: +// - word : convert four bytes of a 32-bit weight word straight into four fp16 +// bit patterns with native DWORD field moves, folding E4M3's +// residual 2^-8 into the per-K-group scale. No memory traffic, no +// 8-bit ALU ops. +// - lut : the 128-entry magnitude LUT in `bestla/sycl/fp8_lut.h` (sign +// applied separately). +// - bits : self-contained inline bit manipulation. // -// Selection is done at runtime through the environment variable -// `ARK_FP8_DECODE_USE_LUT`: -// - unset / "1" / "true" / "on" / "yes" (case-insensitive) -> LUT path (default) -// - "0" / "false" / "off" / "no" (case-insensitive) -> inline bit-manip +// Selection is done at runtime through `ARK_FP8_DECODE_MODE` ("word" / "lut" / +// "bits", case-insensitive), defaulting to "word". The legacy +// `ARK_FP8_DECODE_USE_LUT` variable still works when set explicitly and keeps +// its old meaning (truthy -> lut, falsy -> bits). // -// The env var is read once on the host (cached) and passed as a template -// parameter into the SYCL kernel, so there is no per-element runtime branch. -// The actual primitives live in `sycl_tla_moe_dequant.hpp` (shared with the -// mixed-input prefill path); this file just re-exports them via `using`. +// The env var is read on the host and passed as a template parameter into the +// SYCL kernel, so there is no per-element runtime branch. The actual primitives +// live in `sycl_tla_moe_dequant.hpp` (shared with the mixed-input prefill +// path); this file just re-exports them via `using`. // ---------------------------------------------------------------------------- #if defined(ARK_XPU) && defined(ARK_SYCL_TLA) @@ -146,7 +151,7 @@ class MoEDecodeKernelInt8; template class MoEDecodeKernelInt2; -template +template class MoEDecodeKernelFP8; // ---------------------------------------------------------------------------- @@ -156,16 +161,21 @@ class MoEDecodeKernelFP8; // keep the in-kernel call sites (`decode_fp8<...>(byte)`) and the host-side // `fp8_decode_use_lut()` lookup inside `moe_decode_detail` working unchanged. // ---------------------------------------------------------------------------- +using moe_dequant::Fp8DecodeMode; using moe_dequant::decode_fp8; using moe_dequant::decode_fp8_e4m3_bits; using moe_dequant::decode_fp8_e4m3_lut; using moe_dequant::decode_fp8_e5m2_bits; using moe_dequant::decode_fp8_e5m2_lut; +using moe_dequant::decode_fp8_half_bits; +using moe_dequant::decode_fp8_quad_half_bits; using moe_dequant::decode_int2_quad; using moe_dequant::decode_int4_octet; using moe_dequant::decode_int4_pair; using moe_dequant::decode_int8; +using moe_dequant::fp8_decode_mode; using moe_dequant::fp8_decode_use_lut; +using moe_dequant::fp8_word_scale_bias; // ---------------------------------------------------------------------------- // Build a [total_tokens] -> expert_id mapping from num_tokens_per_expert. @@ -1213,10 +1223,24 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // ---------------------------------------------------------------------------- // FP8 (E4M3 / E5M2) GEMV with group-wise scale (no zero-point). // -// Weights are 1 FP8 byte per element [E, N, K]. The byte is decoded via the -// `decode_fp8` helper, which selects between the LUT and the -// inline bit-manipulation path at compile time. The choice is driven at -// launch time by the env var `ARK_FP8_DECODE_USE_LUT` (default: ON). +// Weights are 1 FP8 byte per element [E, N, K]. How a byte becomes a float is +// chosen at launch time by `fp8_decode_mode()` and passed in as the `Mode` +// template parameter, so the hot path stays branch-free: +// +// * `kWord` (default) -- the four bytes of a 32-bit weight word are turned +// into four fp16 bit patterns by `decode_fp8_quad_half_bits`, i.e. a couple +// of native DWORD ops and no memory traffic at all. This mirrors the +// word-native `decode_int4_octet` treatment that made int4 decode fast: Xe +// ALU lanes are 32-bit, so the previous `sycl::vec` weight +// vector plus per-byte decode paid narrow-type regioning on *every* weight +// element, and the LUT variant additionally issued one load per element in +// a loop that already does only ~1 MAC per byte. +// For E4M3 the field move leaves a constant 2^-8 factor, which is folded +// into the per-K-group scale below (`fp8_word_scale_bias`), so it costs +// nothing per element. +// +// * `kLut` / `kBits` -- the original per-byte `decode_fp8` +// decoders, kept for A/B measurement and regression escape. // ---------------------------------------------------------------------------- // Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK weight @@ -1225,33 +1249,79 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // remainder, mirroring the int4/int8 paths. sycl::vec only supports widths of // 1, 2, 3, 4, 8 or 16, so CHUNK is processed in 16-wide sub-blocks. // +// In `kWord` mode the 16 weight bytes of a sub-block are read as a +// `sycl::vec` -- the same 16-byte transaction (and the same +// 16-byte alignment requirement) as the byte vector it replaces, but 32-bit +// typed, so the decode never leaves the native datapath. +// // The per-group scale is constant across the whole group, so it is NOT applied // here: this accumulates the raw dot product (sum of act * decoded_fp8) and the // caller multiplies the group total by the scale once (Σ a·(w·s) == s·Σ a·w). // For the per-expert / per-tensor scale case (group_size == K, one scale per // output row) this collapses the whole K reduction to a single scale multiply, // removing one multiply per K element on the decode hot path. -template -static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc) { +// +// Two independent partial accumulators break the single fp32 dependency chain +// so the FMA pipeline is not latency-bound (same trick as `int4_decode_chunk`); +// the caller reduces the pair. +template +static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc0, + float& acc1) { static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); constexpr int SUB = 16; using ActVec = sycl::vec; - using ByteVec = sycl::vec; #pragma unroll for (int s = 0; s < CHUNK / SUB; ++s) { const ActVec av = *reinterpret_cast(act_ptr + s * SUB); - const ByteVec wv = *reinterpret_cast(w_ptr + s * SUB); + if constexpr (Mode == Fp8DecodeMode::kWord) { + constexpr int WORDS = SUB / 4; // one 32-bit word per 4 FP8 bytes + using WordVec = sycl::vec; + const WordVec wv = *reinterpret_cast(w_ptr + s * SUB); #pragma unroll - for (int u = 0; u < SUB; ++u) { - const float w = decode_fp8(wv[u]); - const ScalarT a = sycl::bit_cast(static_cast(av[u])); - acc += static_cast(a) * w; + for (int w = 0; w < WORDS; ++w) { + uint32_t lo2, hi2; + decode_fp8_quad_half_bits(wv[w], lo2, hi2); + const uint16_t hb[4] = {static_cast(lo2), static_cast(lo2 >> 16), + static_cast(hi2), static_cast(hi2 >> 16)}; +#pragma unroll + for (int u = 0; u < 4; u += 2) { + const ScalarT a0 = sycl::bit_cast(static_cast(av[4 * w + u])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[4 * w + u + 1])); + acc0 += static_cast(a0) * static_cast(sycl::bit_cast(hb[u])); + acc1 += static_cast(a1) * static_cast(sycl::bit_cast(hb[u + 1])); + } + } + } else { + constexpr bool kUseLut = (Mode == Fp8DecodeMode::kLut); + using ByteVec = sycl::vec; + const ByteVec wv = *reinterpret_cast(w_ptr + s * SUB); +#pragma unroll + for (int u = 0; u < SUB; u += 2) { + const float w0 = decode_fp8(wv[u]); + const float w1 = decode_fp8(wv[u + 1]); + const ScalarT a0 = sycl::bit_cast(static_cast(av[u])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[u + 1])); + acc0 += static_cast(a0) * w0; + acc1 += static_cast(a1) * w1; + } } } } -template +// Single-byte decode matching `fp8_decode_chunk`'s convention: in `kWord` mode +// the returned value carries the same folded 2^-8 bias as the vector stage, so +// the scalar tail can share the group accumulator. +template +static inline float fp8_decode_scalar(uint8_t raw) { + if constexpr (Mode == Fp8DecodeMode::kWord) { + return static_cast(sycl::bit_cast(decode_fp8_half_bits(raw))); + } else { + return decode_fp8(raw); + } +} + +template void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, int K, int group_size) { if (N % N_TILE != 0) { @@ -1264,11 +1334,15 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh const int n_tiles = N / N_TILE; const int num_groups_k = K / group_size; + // Undoes the exponent re-bias the word-native decode leaves behind (1.0f for + // every other mode). Exact power of two, applied once per K-group. + constexpr float kScaleBias = + (Mode == Fp8DecodeMode::kWord) ? fp8_word_scale_bias() : 1.0f; sycl::range<2> global{static_cast(total_tokens), static_cast(n_tiles * SG_SIZE)}; sycl::range<2> local{1, static_cast(SG_SIZE)}; - q->parallel_for>( + q->parallel_for>( sycl::nd_range<2>(global, local), [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { const int token = static_cast(it.get_global_id(0)); @@ -1286,7 +1360,7 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh float acc = 0.0f; for (int g = 0; g < num_groups_k; ++g) { - const float scale = static_cast(s_row[g]); + const float scale = static_cast(s_row[g]) * kScaleBias; const int k_base = g * group_size; // Vectorized ladder mirroring the int4/int8 paths: process 32 // K-elements (32 weight bytes + vec activations) at a @@ -1297,32 +1371,57 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // load for the shipped group sizes (32/64/128/256); hoisting the // scale removes one multiply per K element, which is the dominant // cost for the per-expert / per-tensor scale case (group_size == K). - float group_acc = 0.0f; + // Two partial accumulators break the fp32 dependency chain. + float group_acc0 = 0.0f; + float group_acc1 = 0.0f; int kk = 0; constexpr int CHUNK32 = 32; const int end32 = (group_size / CHUNK32) * CHUNK32; for (; kk < end32; kk += CHUNK32) { - fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, - group_acc); + fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, + group_acc0, group_acc1); } constexpr int CHUNK16 = 16; const int end16 = kk + ((group_size - kk) / CHUNK16) * CHUNK16; for (; kk < end16; kk += CHUNK16) { - fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, - group_acc); + fp8_decode_chunk(act_row + k_base + kk, w_row + k_base + kk, + group_acc0, group_acc1); } for (; kk < group_size; ++kk) { - const uint8_t raw = w_row[k_base + kk]; - const float w = decode_fp8(raw); - group_acc += static_cast(act_row[k_base + kk]) * w; + const float w = fp8_decode_scalar(w_row[k_base + kk]); + group_acc0 += static_cast(act_row[k_base + kk]) * w; } - acc += group_acc * scale; + acc += (group_acc0 + group_acc1) * scale; } outputs[static_cast(token) * N + n_global] = static_cast(acc); }); } +// Runtime -> compile-time bridge for the decode-mode selector. Keeps the +// `moe_gemm_decode` dispatch to one branch per (act dtype, format) instead of +// re-nesting the mode selection at every call site. +template +void launch_fp8_by_mode(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, + const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, + int total_tokens, int N, int K, int group_size) { + switch (fp8_decode_mode()) { + case Fp8DecodeMode::kLut: + launch_fp8(q, activations, weights, scales, outputs, + expert_id_per_token, total_tokens, N, K, group_size); + return; + case Fp8DecodeMode::kBits: + launch_fp8(q, activations, weights, scales, outputs, + expert_id_per_token, total_tokens, N, K, group_size); + return; + case Fp8DecodeMode::kWord: + default: + launch_fp8(q, activations, weights, scales, outputs, + expert_id_per_token, total_tokens, N, K, group_size); + return; + } +} + } // namespace moe_decode_detail // ---------------------------------------------------------------------------- @@ -1396,6 +1495,52 @@ inline bool moe_decode_dpas_s4_occupancy_ok(int total_tokens, int num_experts) { return static_cast(total_tokens) >= min_tokens_per_expert * static_cast(num_experts); } +// ---------------------------------------------------------------------------- +// Env-flag helper -- `ARK_MOE_DECODE_DPAS_FP8` (default ON). When ON, FP8 +// (E4M3 / E5M2, sym) decode is routed to the decode-phase FP8 DPAS grouped +// GEMM (`moe_dpas_fp8::moe_decode_fp8_dpas_per_group_dispatch`) instead of the +// scalar FMA GEMV (`launch_fp8`). This is the FP8 twin of +// `ARK_MOE_DECODE_DPAS_S4`: same `[E, N, K]` FP8 bytes and `[E, N, K/group]` +// scales, no repack, tile picked from the `A_avg_M` ladder. +// +// Setting the var to "0" / "false" / "off" / "no" (case-insensitive) forces the +// scalar GEMV, for A/B comparison and regression escape. Shapes that fail the +// DPAS shape gate and batches that fail the tokens-per-expert occupancy gate +// (`moe_decode_dpas_fp8_occupancy_ok`, below -- this is what keeps real decode +// batches on the fast scalar GEMV) always fall back to the scalar path +// regardless of this flag. Re-read on every call so tests / benchmarks can +// toggle the path in-process. +// ---------------------------------------------------------------------------- +inline bool moe_decode_dpas_fp8_enabled() { + return moe_decode_detail::env_flag_enabled("ARK_MOE_DECODE_DPAS_FP8", true); // default ON +} + +// ---------------------------------------------------------------------------- +// Occupancy gate for the FP8 DPAS decode path. Identical reasoning to +// `moe_decode_dpas_s4_occupancy_ok`: the smallest DPAS tile the decode +// dispatch can pick (`dpas_w4a16_policy_m_8`) processes 8 token rows per +// expert, so below 8 tokens per expert on average the tile is mostly padding +// and the bandwidth-bound FP8 weights get streamed for rows that contribute +// nothing -- exactly the regime real decode batches live in. Above that the +// DPAS pipeline wins, so the threshold is where the two cross. +// +// `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` overrides the tokens-per-expert threshold; +// "0" disables the gate (always take DPAS when the shape gate allows), which is +// what the accuracy tests use to exercise the DPAS kernel on tiny shapes. +// ---------------------------------------------------------------------------- +inline bool moe_decode_dpas_fp8_occupancy_ok(int total_tokens, int num_experts) { + if (num_experts <= 0) return true; + long long min_tokens_per_expert = 8; // rows in `dpas_w4a16_policy_m_8` + const char* env = std::getenv("ARK_MOE_DECODE_DPAS_FP8_MIN_TPE"); + if (env != nullptr) { + char* end = nullptr; + long long v = std::strtoll(env, &end, 10); + if (end != env && v >= 0) min_tokens_per_expert = v; + } + if (min_tokens_per_expert == 0) return true; + return static_cast(total_tokens) >= min_tokens_per_expert * static_cast(num_experts); +} + // ---------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_DECODE_COALESCE_INT4` (default ON). When ON, the // int4 scalar-GEMV fallback (asym, or sym with the DPAS path disabled / shape @@ -1456,22 +1601,32 @@ inline bool moe_decode_coalesce_int4_amortized(int total_tokens, int num_experts // BTLA_DTYPE::S2_CLIP : packed int2 weights [E, N, K/4] (uint8), // 4 values per byte, sym/asym like int4 // BTLA_DTYPE::F8_E4M3 / F8_E5M2 : FP8 weights [E, N, K] (uint8 buffer), -// group-wise scales, no zero-points +// group-wise scales, no zero-points. Routed +// to the per-group FP8 DPAS grouped GEMM only +// when the batch fills its M tile (>= 8 +// tokens per expert on average, +// `ARK_MOE_DECODE_DPAS_FP8` default ON); a +// disabled flag, a shape-gate miss, or a +// decode-sized batch uses the scalar GEMV. // act_dtype: F16 or BF16 (must match scales/outputs dtype) // ---------------------------------------------------------------------------- inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, void* scales, void* zeros, void* outputs, int* expert_id_per_token_buf, BTLA_DTYPE act_dtype, BTLA_DTYPE weight_dtype, int N, int K, int group_size, int* num_tokens_per_expert, int num_experts, int total_tokens, bool asym) { - // The S4-sym DPAS fast path consumes `num_tokens_per_expert` directly and - // never reads `expert_id_per_token_buf`. Skipping the fill on that path + // The S4-sym and FP8 DPAS fast paths consume `num_tokens_per_expert` directly + // and never read `expert_id_per_token_buf`. Skipping the fill on those paths // removes an extra device-timeline kernel launch from the decode hot path; - // every other path (fp, int8, int2, fp8, and the scalar int4 fallback) still - // needs the per-token expert mapping. + // every other path (fp, int8, int2, and the scalar int4 / fp8 fallbacks) + // still needs the per-token expert mapping. const bool s4_dpas_fastpath = weight_dtype == BTLA_DTYPE::S4_CLIP && !asym && moe_decode_dpas_s4_enabled() && moe_decode_dpas_s4_occupancy_ok(total_tokens, num_experts) && moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size); - if (!s4_dpas_fastpath) { + const bool fp8_dpas_fastpath = (weight_dtype == BTLA_DTYPE::F8_E4M3 || weight_dtype == BTLA_DTYPE::F8_E5M2) && + !asym && moe_decode_dpas_fp8_enabled() && + moe_decode_dpas_fp8_occupancy_ok(total_tokens, num_experts) && + moe_dpas_fp8::moe_prefill_dpas_fp8_pergroup_shape_ok(N, K, group_size); + if (!s4_dpas_fastpath && !fp8_dpas_fastpath) { moe_decode_detail::fill_expert_id_per_token(q, expert_id_per_token_buf, num_tokens_per_expert, num_experts, total_tokens); } @@ -1655,63 +1810,71 @@ inline void moe_gemm_decode(sycl::queue* q, void* activations, void* weights, vo if (asym) { throw std::invalid_argument("moe_gemm_decode(fp8): asym mode is not supported"); } + if (act_dtype != BTLA_DTYPE::F16 && act_dtype != BTLA_DTYPE::BF16) { + throw std::invalid_argument("moe_gemm_decode(fp8): act_dtype must be FP16 or BF16"); + } const bool is_e4m3 = (weight_dtype == BTLA_DTYPE::F8_E4M3); - const bool use_lut = moe_decode_detail::fp8_decode_use_lut(); - if (act_dtype == BTLA_DTYPE::F16) { - if (is_e4m3) { - if (use_lut) { - moe_decode_detail::launch_fp8( + // Fast path: FP8 through the decode-phase per-group DPAS grouped GEMM. + // Falls back to the scalar GEMV when the env flag is off, when the batch is + // too small to fill the DPAS M tile (the usual decode case), or when the + // shape gate rejects the tile geometry (e.g. N%64!=0, K%32!=0, unsupported + // group_size). Reuses the `fp8_dpas_fastpath` predicate computed above + // (which also gated the `fill_expert_id_per_token` skip) so the two + // decisions cannot diverge. + if (fp8_dpas_fastpath) { + if (act_dtype == BTLA_DTYPE::F16) { + if (is_e4m3) { + moe_dpas_fp8::moe_decode_fp8_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, - total_tokens, N, K, group_size); + static_cast(scales), static_cast(outputs), num_tokens_per_expert, + num_experts, N, K, group_size, total_tokens); } else { - moe_decode_detail::launch_fp8( + moe_dpas_fp8::moe_decode_fp8_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, - total_tokens, N, K, group_size); + static_cast(scales), static_cast(outputs), num_tokens_per_expert, + num_experts, N, K, group_size, total_tokens); } } else { - if (use_lut) { - moe_decode_detail::launch_fp8( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, - total_tokens, N, K, group_size); - } else { - moe_decode_detail::launch_fp8( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, - total_tokens, N, K, group_size); - } - } - } else if (act_dtype == BTLA_DTYPE::BF16) { - using BF = sycl::ext::oneapi::bfloat16; - if (is_e4m3) { - if (use_lut) { - moe_decode_detail::launch_fp8( + using BF = sycl::ext::oneapi::bfloat16; + if (is_e4m3) { + moe_dpas_fp8::moe_decode_fp8_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, - group_size); + static_cast(scales), static_cast(outputs), num_tokens_per_expert, num_experts, N, K, + group_size, total_tokens); } else { - moe_decode_detail::launch_fp8( + moe_dpas_fp8::moe_decode_fp8_dpas_per_group_dispatch( q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, - group_size); + static_cast(scales), static_cast(outputs), num_tokens_per_expert, num_experts, N, K, + group_size, total_tokens); } + } + return; + } + if (act_dtype == BTLA_DTYPE::F16) { + if (is_e4m3) { + moe_decode_detail::launch_fp8_by_mode( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(outputs), expert_id_per_token_buf, + total_tokens, N, K, group_size); } else { - if (use_lut) { - moe_decode_detail::launch_fp8( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, - group_size); - } else { - moe_decode_detail::launch_fp8( - q, static_cast(activations), static_cast(weights), - static_cast(scales), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, - group_size); - } + moe_decode_detail::launch_fp8_by_mode( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(outputs), expert_id_per_token_buf, + total_tokens, N, K, group_size); } } else { - throw std::invalid_argument("moe_gemm_decode(fp8): act_dtype must be FP16 or BF16"); + using BF = sycl::ext::oneapi::bfloat16; + if (is_e4m3) { + moe_decode_detail::launch_fp8_by_mode( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, + group_size); + } else { + moe_decode_detail::launch_fp8_by_mode( + q, static_cast(activations), static_cast(weights), + static_cast(scales), static_cast(outputs), expert_id_per_token_buf, total_tokens, N, K, + group_size); + } } return; } diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp index ecc95f52c3..736ab0e3c0 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp @@ -10,6 +10,11 @@ // Currently extracted: // - FP8 (E4M3 / E5M2) byte->float decoders + host-side // `ARK_FP8_DECODE_USE_LUT` env-var reader (PR-A1). +// - FP8 word-native decoders (`decode_fp8_half_bits`, +// `decode_fp8_quad_half_bits`, `fp8_word_scale_bias`) + the +// `Fp8DecodeMode` selector: convert FP8 bytes to fp16 bit patterns with +// pure 32-bit field moves (no LUT load, no 8-bit ALU), folding E4M3's +// residual 2^-8 into the per-K-group scale. Used by the decode GEMV. // - INT2 / INT4 / INT8 packed-byte decoders (PR-A2): return the raw // integer field(s) prior to `(q - zp) * scale`. Both the decode (GEMV) // and prefill (mixed-input Grouped GEMM) paths call these directly, @@ -134,6 +139,98 @@ inline float decode_fp8(uint8_t byte) { } } +// ---------------------------------------------------------------------------- +// Word-native FP8 -> half decode (the `Fp8DecodeMode::kWord` path). +// +// Both LUT and inline-bits decoders above cost real work per weight byte: the +// LUT issues a memory load (plus a sign select) and the bit-manip path runs a +// branchy `ldexp` chain. On the decode hot path -- a pure GEMV that streams one +// weight byte per multiply-add -- that dequant cost is the kernel. Neither is +// necessary, because an FP8 byte is already an IEEE-style float and fp16 is a +// *superset* of both FP8 formats: the whole conversion is a bit-field move. +// +// E5M2 -> fp16: identical sign position, identical 5-bit exponent with the +// same bias 15, mantissa just needs 8 more bits -> +// h = byte << 8 +// Exact for every one of the 256 encodings, specials included +// (subnormals stay subnormal, exp==31 stays Inf/NaN). +// +// E4M3 -> fp16: 4-bit exponent, bias 7. Shifting the 7 magnitude bits up by +// 7 lands the exponent in fp16's exponent field and the 3 +// mantissa bits in the top of fp16's mantissa, which yields the +// correct value scaled by 2^(7-15) == 2^-8; the sign bit has to +// move 8 places instead of 7. Both moves collapse into one +// add + one shift, because adding the sign bit to itself +// carries it exactly one position further: +// h = (byte + (byte & 0x80)) << 7 +// The residual 2^-8 is constant, so callers fold the reciprocal +// (`fp8_word_scale_bias()` == 256.0f) into the +// per-K-group scale, i.e. it costs nothing per element. +// +// Exactness (verified exhaustively over all 256 byte values / all four +// format-mode combinations): E5M2 is bit-exact including Inf/NaN; E4M3 is +// bit-exact for all 254 finite encodings, including subnormals and both zeros. +// The two E4M3 *NaN* encodings (0x7F / 0xFF -- `torch.float8_e4m3fn` has no +// Inf) decode to +-480 instead of NaN, since fp16 has no NaN pattern reachable +// by a pure field move. auto-round FP8 checkpoints are produced by scaling to +// `finfo(float8_e4m3fn).max == 448` and clamping, so those two encodings cannot +// occur; callers that need NaN propagation can select `Fp8DecodeMode::kLut` or +// `kBits` (see `fp8_decode_mode()`). +// ---------------------------------------------------------------------------- +template +inline uint16_t decode_fp8_half_bits(uint32_t byte) { + if constexpr (IsE4M3) { + return static_cast((byte + (byte & 0x80u)) << 7); + } else { + return static_cast(byte << 8); + } +} + +// Constant the caller must fold into the per-group scale to undo the exponent +// re-bias performed by `decode_fp8_half_bits`. Exact power of two, so the fold +// is a pure exponent bump on the fp32 scale (no rounding). +template +inline constexpr float fp8_word_scale_bias() { + return IsE4M3 ? 256.0f : 1.0f; +} + +// SWAR form: decode the four FP8 bytes of one little-endian 32-bit word into +// two 32-bit words, each packing two fp16 bit patterns (low 16-bit lane holds +// the lower K index). Bit-identical to calling `decode_fp8_half_bits` on each +// byte, but the whole quad costs a handful of native DWORD ops and -- crucially +// on Xe, whose ALU lanes are 32-bit -- never touches an 8-bit-typed vector, +// which IGC has to expand into narrow-type regioning. This mirrors what +// `decode_int4_octet` does for packed nibbles. +template +inline void decode_fp8_quad_half_bits(uint32_t word, uint32_t& lo2, uint32_t& hi2) { + // Spread bytes 0/1 and 2/3 into the two 16-bit lanes of `lo` / `hi`. + const uint32_t lo = (word & 0x000000FFu) | ((word & 0x0000FF00u) << 8); + const uint32_t hi = ((word >> 16) & 0x000000FFu) | ((word >> 8) & 0x00FF0000u); + if constexpr (IsE4M3) { + // Per-lane `(b + (b & 0x80)) << 7`. A lane's value is <= 0x17F before the + // shift and <= 0xBF80 after it, so neither the add nor the shift can carry + // into the neighbouring lane. + lo2 = (lo + (lo & 0x00800080u)) << 7; + hi2 = (hi + (hi & 0x00800080u)) << 7; + } else { + lo2 = lo << 8; + hi2 = hi << 8; + } +} + +// ---------------------------------------------------------------------------- +// FP8 decode implementation selector. +// +// kWord : word-native bit-field move + folded scale bias (default; fastest, +// no memory traffic, no 8-bit ALU ops -- see above). +// kLut : 128-entry magnitude table in `bestla/sycl/fp8_lut.h`. +// kBits : self-contained inline bit manipulation. +// +// `kLut` / `kBits` are kept reachable for A/B measurement, regression escape, +// and the (checkpoint-impossible) E4M3 NaN encodings. +// ---------------------------------------------------------------------------- +enum class Fp8DecodeMode { kWord, kLut, kBits }; + // ---------------------------------------------------------------------------- // INT4 (S4_CLIP) packed-byte decode. // @@ -287,6 +384,10 @@ inline int decode_int8(uint8_t raw) { // // Read once on first call and cached in a function-local static, so it is // safe (and free) to call this on every launch. +// +// NOTE: this only chooses between the two *per-byte* decoders. The decode GEMV +// selects between {word, lut, bits} through `fp8_decode_mode()` below, which +// still honours this variable when it is set explicitly. // ---------------------------------------------------------------------------- inline bool fp8_decode_use_lut() { static const bool value = []() { @@ -300,6 +401,38 @@ inline bool fp8_decode_use_lut() { return value; } +// ---------------------------------------------------------------------------- +// Host-side selector for the FP8 decode implementation. +// +// `ARK_FP8_DECODE_MODE` = "word" | "lut" | "bits" (case-insensitive) picks a +// mode explicitly and wins over everything else. +// +// Otherwise, if the legacy `ARK_FP8_DECODE_USE_LUT` is set, it keeps its old +// meaning (`kLut` when truthy, `kBits` when falsy) so existing A/B scripts +// behave exactly as before. +// +// With neither set, the default is `kWord` -- the word-native bit-field move. +// +// Re-read on every call (not cached) so tests and benchmarks can toggle the +// path in-process; the result is passed into the kernel as a template argument, +// so there is no per-element runtime branch. +// ---------------------------------------------------------------------------- +inline Fp8DecodeMode fp8_decode_mode() { + const char* mode = std::getenv("ARK_FP8_DECODE_MODE"); + if (mode != nullptr) { + std::string s(mode); + for (char& c : s) c = static_cast(std::tolower(static_cast(c))); + if (s == "word") return Fp8DecodeMode::kWord; + if (s == "lut") return Fp8DecodeMode::kLut; + if (s == "bits") return Fp8DecodeMode::kBits; + // Unrecognised value: fall through to the legacy variable / default. + } + if (std::getenv("ARK_FP8_DECODE_USE_LUT") != nullptr) { + return fp8_decode_use_lut() ? Fp8DecodeMode::kLut : Fp8DecodeMode::kBits; + } + return Fp8DecodeMode::kWord; +} + } // namespace moe_dequant } // namespace ark diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp index b67694e80b..a7bf88d2e2 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp @@ -80,9 +80,12 @@ #pragma once #include +#include +#include #include #include #include +#include #ifdef ARK_XPU #include @@ -152,6 +155,39 @@ struct cute_scalar { template using cute_scalar_t = typename cute_scalar::type; +// --------------------------------------------------------------------------- +// Persistent per-queue atomic work-group counter. +// +// The grouped-GEMM launchers below need a single `int32_t` device slot as a +// global work-group counter (`atomicAdd`). The kernel self-initialises it to 0 +// at launch (group 0 / lane 0 does `atm.store(0)`), so the host never has to +// reset it between calls. Allocating it per dispatch with `sycl::malloc_device` +// and releasing it with `sycl::free` costs two queue synchronizations, which is +// pure overhead on the decode hot path where the GEMM itself is only tens of +// microseconds. +// +// Instead, hand out one persistent buffer per queue and reuse it across calls. +// This is safe because every launcher call is synchronous (`event.wait()` in +// `MoEGEMMLauncher`), so two launches can never share the buffer concurrently. +// Buffers live until process exit (one `int32_t` per queue), matching the +// singleton lifetime already used by `EventManager`. The S4 header re-exports +// this helper rather than defining its own, so both paths share one cache. +// --------------------------------------------------------------------------- +inline int32_t* get_persistent_atomic_buffer(sycl::queue* q) { + static std::mutex mtx; + static std::unordered_map cache; + std::lock_guard lock(mtx); + auto it = cache.find(q); + if (it != cache.end()) return it->second; + int32_t* buf = sycl::malloc_device(1, *q); + if (buf == nullptr) { + throw std::runtime_error( + "moe_dpas_fp8: failed to allocate persistent atomic buffer"); + } + cache.emplace(q, buf); + return buf; +} + // --------------------------------------------------------------------------- // Policy classes (ported verbatim from vllm-xpu-kernels // `gemm_xe2_policy.hpp`, renamed to `dpas_*` to avoid collision with any @@ -1041,6 +1077,81 @@ void moe_prefill_fp8_dpas_per_group_dispatch( sycl::free(atomic_buffer, *q); } +// --------------------------------------------------------------------------- +// Host-side driver: per-K-group FP8, *decode* phase. +// +// Same math, same mainloop and same weight/scale layout as +// `moe_prefill_fp8_dpas_per_group_dispatch` above -- only the tile selection +// and the atomic-buffer lifetime differ, for two reasons that are specific to +// the decode regime (a handful of tokens spread over many experts): +// +// 1. Finer small-M ladder. The reference `w8a16` dispatch in +// vllm-xpu-kernels bottoms out at the 16-row tile (`m_16`), while its +// `w4a16` dispatch has an extra 8-row bucket. Decode `A_avg_M` is far +// below 16, so the missing rung means half of every M tile is padding and +// the (bandwidth-bound) FP8 weights are streamed for rows that contribute +// nothing. `dpas_w4a16_policy_m_8` carries no 4-bit-specific types -- it +// is purely a `WGTile`/`SGLayout` shape (8x64x32) -- so the FP8 mainloop +// reuses it verbatim, closing that gap. This mirrors what the S4 decode +// dispatch already does. +// 2. Persistent atomic counter. The prefill dispatch allocates and frees the +// work-group counter per call; each of those forces a queue sync. At +// prefill sizes that is noise, at decode sizes it is a large fraction of +// the total. Use the per-queue persistent slot instead. +// +// The upper rungs are pulled in to match the S4 decode ladder +// (`m_8` -> `m_16` -> `m_32` -> wide) rather than the prefill one, whose +// `<= 512 -> m_32` rung is tuned for prefill-sized batches. +// +// Numerically identical to the prefill dispatch for every input; only the tile +// geometry changes, so `test_moe_prefill_accuracy.py::test_accuracy_fp8` +// tolerances apply unchanged. Inherits this header's +// NEEDS-HARDWARE-VALIDATION status. +// --------------------------------------------------------------------------- + +template +void moe_decode_fp8_dpas_per_group_dispatch( + sycl::queue* q, const ScalarT* activations, const uint8_t* weights_NK, + const ScalarT* scales, ScalarT* outputs, const int* num_tokens_per_expert, + int E, int N, int K, int group_size, int total_tokens) { + if (E == 0 || N == 0 || K == 0 || total_tokens == 0) return; + if (K % group_size != 0) { + throw std::invalid_argument( + "moe_decode_fp8_dpas(per-group): K must be a multiple of group_size"); + } + + compat::set_default_queue(*q); + + using ElementB = std::conditional_t; + using ElementA = cute_scalar_t; + const auto* activations_ca = + reinterpret_cast(activations); + const auto* scales_ca = reinterpret_cast(scales); + auto* outputs_ca = reinterpret_cast(outputs); + + const int A_avg_M = total_tokens / E; + + int32_t* atomic_buffer = get_persistent_atomic_buffer(q); + +#define ARK_DPAS_DECODE_PG_LAUNCH(policy) \ + MoEGEMMLauncher<'R', 'C', policy, ScaleMode::kPerGroup>( \ + *q, activations_ca, reinterpret_cast(weights_NK), \ + scales_ca, static_cast(nullptr), outputs_ca, N, K, \ + num_tokens_per_expert, E, group_size, atomic_buffer); + + if (A_avg_M <= 4) { + ARK_DPAS_DECODE_PG_LAUNCH(dpas_w4a16_policy_m_8); + } else if (A_avg_M <= 8) { + ARK_DPAS_DECODE_PG_LAUNCH(dpas_w8a16_policy_m_16); + } else if (A_avg_M <= 128) { + ARK_DPAS_DECODE_PG_LAUNCH(dpas_w8a16_policy_m_32); + } else { + ARK_DPAS_DECODE_PG_LAUNCH(dpas_w8a16_policy); + } +#undef ARK_DPAS_DECODE_PG_LAUNCH +} + // --------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_PREFILL_DPAS_FP8` (default ON per plan). // diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp index ba2bb9ca6e..0a1c6a923c 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_s4_dpas.hpp @@ -148,34 +148,12 @@ using ::ark::moe_dpas_fp8::make_moe_tensor; // --------------------------------------------------------------------------- // Persistent per-queue atomic work-group counter. // -// The grouped-GEMM launcher below needs a single `int32_t` device slot as a -// global work-group counter (`atomicAdd`). The kernel self-initialises it to 0 -// at launch (group 0 / lane 0 does `atm.store(0)`), so the host never has to -// reset it between calls. Previously each dispatch call allocated this slot -// with `sycl::malloc_device` and released it with `sycl::free`; both operations -// force a queue synchronization, which is pure overhead on the decode hot path -// where the GEMM itself is only tens of microseconds. -// -// Instead, hand out one persistent buffer per queue and reuse it across calls. -// This is safe because every launcher call is synchronous (`event.wait()` in -// `MoEGEMMLauncher_s4`), so two launches can never share the buffer -// concurrently. Buffers live until process exit (one `int32_t` per queue), -// matching the singleton lifetime already used by `EventManager`. +// Shared with the FP8 path -- see `moe_dpas_fp8::get_persistent_atomic_buffer` +// for the rationale (one `int32_t` device slot per queue reused across calls, +// instead of a `sycl::malloc_device` / `sycl::free` pair per dispatch, each of +// which forces a queue synchronization on the decode hot path). // --------------------------------------------------------------------------- -inline int32_t* get_persistent_atomic_buffer(sycl::queue* q) { - static std::mutex mtx; - static std::unordered_map cache; - std::lock_guard lock(mtx); - auto it = cache.find(q); - if (it != cache.end()) return it->second; - int32_t* buf = sycl::malloc_device(1, *q); - if (buf == nullptr) { - throw std::runtime_error( - "moe_dpas_s4: failed to allocate persistent atomic buffer"); - } - cache.emplace(q, buf); - return buf; -} +using ::ark::moe_dpas_fp8::get_persistent_atomic_buffer; // --------------------------------------------------------------------------- // Variant B -- per-K-group S4 (sym) mainloop. From 3423da8f6df4d8cdbe9b0120138d91ac3e148053 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:53:01 +0000 Subject: [PATCH 42/47] test: FP8 decode-mode + DPAS parity/perf coverage; docs: FP8 decode section (EN+CN) Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/test/README_MOE_PREFILL_PERF.md | 110 +++++++++++++++++ .../ark/test/README_MOE_PREFILL_PERF_CN.md | 98 +++++++++++++++ auto_round_extension/ark/test/test_moe.py | 112 ++++++++++++++++++ .../ark/test/test_moe_decode_perf.py | 108 ++++++++++++++++- 4 files changed, 427 insertions(+), 1 deletion(-) 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 090bb3e1cf..5d5177426a 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -465,6 +465,116 @@ which forces `ARK_MOE_PREFILL_DPAS_S4=1` + exclusively exercised, at the same production shapes as `test_accuracy_int4`, with tolerance `rtol=atol=1e-1`. +## FP8 Decode Paths (`sycl_tla_moe_decode.hpp`) + +int4-sym decode is now at target, and the same two levers that got it +there apply to FP8: get the dequant off the byte-typed datapath, and stop +paying setup cost per decode call. On top of that, the FP8 MoE dispatch +from vllm-xpu-kernels is mirrored into a decode-specialised entry point. + +**Word-native FP8 decode (`ARK_FP8_DECODE_MODE`, default `word`).** The +decode GEMV does roughly one multiply-add per weight byte, so the dequant +*is* the kernel. Both legacy decoders paid real work per byte: `lut` +issues a memory load per weight element into the 128-entry magnitude table +plus a sign select, and `bits` runs a branchy `ldexp` chain. Both also +indexed an 8-bit-typed `sycl::vec`, which Xe's 32-bit ALU +lanes cannot address directly, so IGC expands it into narrow-type +regioning — exactly the problem `decode_int4_octet` fixed for nibbles. + +None of that work is necessary, because an FP8 byte is already an +IEEE-style float and fp16 is a *superset* of both FP8 formats: the whole +conversion is a bit-field move. + +| Format | fp16 bit pattern | Exactness | +| ------ | ---------------- | --------- | +| E5M2 | `byte << 8` | Bit-exact for all 256 encodings — same sign position, same 5-bit exponent, same bias 15. Subnormals stay subnormal, `exp==31` stays Inf/NaN. | +| E4M3 | `(byte + (byte & 0x80)) << 7` | Bit-exact for all 254 finite encodings (normals, subnormals, both zeros), yielding the true value × `2^-8`. | + +E4M3's 4-bit exponent has bias 7 against fp16's bias 15, so the field move +leaves a constant `2^-8` factor; `fp8_word_scale_bias()` (`256.0f`) +is folded into the per-K-group scale, an exact power of two applied once +per group, so it costs nothing per element. Adding the sign bit to itself +carries it exactly one position further, which is why the sign move and +the magnitude move collapse into one add plus one shift. + +The kernel reads the weights as `sycl::vec` — the same +16-byte transaction and the same 16-byte alignment requirement as the byte +vector it replaces — and `decode_fp8_quad_half_bits` turns each 32-bit word +into four fp16 bit patterns in a handful of native DWORD ops (SWAR, no +cross-lane carry). Two partial accumulators break the fp32 dependency +chain, as in `int4_decode_chunk`. Both primitives live in +`sycl_tla_moe_dequant.hpp`, and both were verified exhaustively over all +256 byte values in both formats. + +**E4M3 NaN caveat.** The two E4M3 NaN encodings (`0x7F` / `0xFF`; +`torch.float8_e4m3fn` has no Inf) decode to ±480 instead of NaN, since fp16 +has no NaN pattern reachable by a pure field move. auto-round FP8 +checkpoints are produced by scaling to `finfo(float8_e4m3fn).max == 448` +and clamping, so those two encodings cannot occur. Callers that need NaN +propagation can select `ARK_FP8_DECODE_MODE=lut` or `=bits`. + +**FP8 DPAS decode dispatch.** `moe_decode_fp8_dpas_per_group_dispatch` +(`sycl_tla_moe_prefill_fp8_dpas.hpp`, `ARK_MOE_DECODE_DPAS_FP8` default ON) +is the FP8 twin of the S4 decode dispatch: same mainloop, same `[E, N, K]` +FP8 bytes + `[E, N, K/group]` scales, no repack. It differs from the +prefill dispatch in two decode-specific ways. + +*Finer small-M ladder.* The reference `w8a16` dispatch in vllm-xpu-kernels +bottoms out at the 16-row tile, while its `w4a16` dispatch has an extra +8-row bucket. Decode `A_avg_M` sits far below 16, so the missing rung means +half of every M tile is padding and the bandwidth-bound FP8 weights get +streamed for rows that contribute nothing. `dpas_w4a16_policy_m_8` carries +no 4-bit-specific types — it is purely an `8×64×32` `WGTile` / `SGLayout` +shape — so the FP8 mainloop reuses it verbatim, closing that gap: + +| `A_avg_M` bucket | WG tile (M×N×K) | Policy | +| ---------------- | --------------- | ------ | +| `≤ 4` | `8×64×32` | `dpas_w4a16_policy_m_8` | +| `≤ 8` | `16×64×32` | `dpas_w8a16_policy_m_16` | +| `≤ 128` | `32×64×32` | `dpas_w8a16_policy_m_32` | +| `> 128` | `128×128×16` | `dpas_w8a16_policy` | + +The upper rungs match the S4 *decode* ladder rather than the FP8 prefill +one, whose `≤ 512 → m_32` rung is tuned for prefill-sized batches. + +*Persistent atomic counter.* The prefill dispatch allocates the +work-group counter with `sycl::malloc_device` and releases it with +`sycl::free` on every call; each of those forces a queue synchronization. +At prefill sizes that is noise, at decode sizes — where the GEMM itself is +only tens of microseconds and one call is issued per generated token — it +is a large fraction of the total. The decode dispatch uses a persistent +per-queue slot instead (`get_persistent_atomic_buffer`, now shared with the +S4 header so both paths use one cache). Taking the fast path also skips the +`fill_expert_id_per_token` pre-pass, since the DPAS dispatch consumes +`num_tokens_per_expert` directly — one fewer kernel launch on the decode +timeline. **Status: NEEDS-HARDWARE-VALIDATION** (this header is an +untested port). + +**Occupancy gate — real decode batches stay on the scalar GEMV.** Same +reasoning as int4-sym: the smallest tile the decode ladder can pick +processes 8 token rows per expert, so below 8 tokens per expert on average +the tile is mostly padding. That is exactly the decode regime (MiniMax-M2, +192 experts: 0.04–1.3 tokens per expert), so FP8 decode is routed to the +scalar GEMV unless the batch supplies at least 8 tokens per expert. +`ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` overrides the threshold; `0` disables the +gate, which is what the parity and A/B perf tests set. Shapes that fail the +per-group shape gate (`N%64==0`, `K%32==0`, `K%group_size==0`, +`group_size ∈ {32,64,128,256}`) always fall back to the scalar GEMV. + +| Env var | Default | Effect | +| ------- | ------- | ------ | +| `ARK_FP8_DECODE_MODE` | `word` | FP8 decode implementation for the scalar GEMV: `word` (bit-field move + folded scale bias), `lut` (128-entry magnitude table), `bits` (inline bit manipulation). | +| `ARK_FP8_DECODE_USE_LUT` | unset | Legacy selector, still honoured when set explicitly and when `ARK_FP8_DECODE_MODE` is unset/unrecognised: truthy → `lut`, falsy → `bits`. Also still drives the mixed-input prefill path. | +| `ARK_MOE_DECODE_DPAS_FP8` | ON | Route FP8 decode to the per-group DPAS grouped GEMM when the shape and occupancy gates pass; `0` forces the scalar GEMV. | +| `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | Minimum tokens per expert before the DPAS path is taken; `0` disables the gate (what the parity/A-B tests set). | + +Perf A/B rows are `test_moe_decode_perf.py::test_perf_fp8_word_vs_lut` +(`speedup` is `lut / word`) and `::test_perf_fp8_dpas_vs_scalar` +(`speedup` is `scalar / dpas`). Correctness is covered by +`test_moe.py::test_decode_fp8_modes_match` (all three decoders agree, and +each tracks the dequant reference) and +`::test_decode_fp8_dpas_matches_scalar`. + ## FP8 per-expert (per-tensor) perf tests `test_perf_fp8_per_tensor` benchmarks the Variant A DPAS path against 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 201a2bb49b..27c4badc83 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 @@ -366,6 +366,104 @@ token 数(16–128)都远低于该门控(8 × 192 个专家 == 1536 个 token) `ARK_MOE_PREFILL_DPAS_INT8=0`,专门验证单遍 mainloop 路径,形状矩阵与 `test_accuracy_int4` 一致,容差 `rtol=atol=1e-1`。 +## FP8 Decode 路径 (`sycl_tla_moe_decode.hpp`) + +int4-sym decode 的性能已经达标,把它推到达标的两个手段同样适用于 FP8: +让 dequant 离开按字节的数据通路,以及不要在每次 decode 调用里重复付出 +启动开销。在此之上,还把 vllm-xpu-kernels 的 FP8 MoE dispatch 镜像成一个 +decode 专用入口。 + +**Word-native FP8 解码 (`ARK_FP8_DECODE_MODE`, 默认 `word`)。** decode +GEMV 每读一个权重字节大约只做一次乘加,所以 dequant *就是* kernel 本身。 +两条旧解码路径每个字节都要付出真实开销:`lut` 每个权重元素都要向 128 项 +幅值表发一次访存再做一次符号选择,`bits` 则要跑一串带分支的 `ldexp`。 +两者还都索引了 8-bit 类型的 `sycl::vec`,而 Xe 的 ALU 通道是 +32-bit 的、无法直接寻址它,于是 IGC 只能展开成窄类型 regioning —— +正是 `decode_int4_octet` 为 nibble 解决过的那个问题。 + +这些工作其实都不必要:FP8 字节本身就是一个 IEEE 风格的浮点数,而 fp16 是 +两种 FP8 格式的*超集*,整个转换就是一次位域搬移。 + +| 格式 | fp16 位模式 | 精确性 | +| ---- | ----------- | ------ | +| E5M2 | `byte << 8` | 对全部 256 种编码逐位精确 —— 符号位位置相同、5 位指数相同、bias 同为 15。次正规数仍是次正规数,`exp==31` 仍是 Inf/NaN。 | +| E4M3 | `(byte + (byte & 0x80)) << 7` | 对全部 254 种有限编码(正规数、次正规数、两个零)逐位精确,得到真值 × `2^-8`。 | + +E4M3 的 4 位指数 bias 为 7,而 fp16 的 bias 是 15,所以位域搬移会留下一个 +常数因子 `2^-8`;`fp8_word_scale_bias()`(`256.0f`)被折叠进 +per-K-group 的 scale,是一个精确的 2 的幂、每组只乘一次,因此对单个元素而言 +零开销。把符号位加到它自身上,恰好会把它再进位一格,这就是符号搬移与幅值 +搬移能合并成一次加法加一次移位的原因。 + +kernel 以 `sycl::vec` 读取权重 —— 与它替换掉的字节向量是同一次 +16 字节访存、同样的 16 字节对齐要求 —— 再由 `decode_fp8_quad_half_bits` 用 +少量原生 DWORD 运算把每个 32 位字变成四个 fp16 位模式(SWAR,不会跨 lane +进位)。两个部分累加器打断 fp32 依赖链,与 `int4_decode_chunk` 的做法一致。 +两个原语都放在 `sycl_tla_moe_dequant.hpp`,并已对两种格式的全部 256 个字节 +值做过穷举验证。 + +**E4M3 NaN 注意事项。** E4M3 的两个 NaN 编码(`0x7F` / `0xFF`; +`torch.float8_e4m3fn` 没有 Inf)会解码成 ±480 而不是 NaN,因为纯位域搬移 +到不了 fp16 的任何 NaN 模式。auto-round 的 FP8 checkpoint 是按 +`finfo(float8_e4m3fn).max == 448` 缩放并 clamp 得到的,所以这两个编码不可能 +出现。需要 NaN 传播的调用方可以选择 `ARK_FP8_DECODE_MODE=lut` 或 `=bits`。 + +**FP8 DPAS decode dispatch。** `moe_decode_fp8_dpas_per_group_dispatch` +(`sycl_tla_moe_prefill_fp8_dpas.hpp`,`ARK_MOE_DECODE_DPAS_FP8` 默认 ON) +是 S4 decode dispatch 的 FP8 对应物:同一套 mainloop、同样的 `[E, N, K]` +FP8 字节 + `[E, N, K/group]` scale、无需 repack。它与 prefill dispatch 有 +两点 decode 专属的差异。 + +*更细的 small-M 阶梯。* vllm-xpu-kernels 的参考 `w8a16` dispatch 最小只到 +16 行 tile,而它的 `w4a16` dispatch 多一个 8 行档位。decode 的 `A_avg_M` +远低于 16,缺这一档意味着每个 M tile 有一半是 padding,而受带宽约束的 FP8 +权重要为这些毫无贡献的行反复搬运。`dpas_w4a16_policy_m_8` 不含任何 4-bit +专用类型 —— 它纯粹是一个 `8×64×32` 的 `WGTile` / `SGLayout` 形状 —— +所以 FP8 mainloop 可以原样复用它,补上这一档: + +| `A_avg_M` 档位 | WG tile (M×N×K) | Policy | +| -------------- | --------------- | ------ | +| `≤ 4` | `8×64×32` | `dpas_w4a16_policy_m_8` | +| `≤ 8` | `16×64×32` | `dpas_w8a16_policy_m_16` | +| `≤ 128` | `32×64×32` | `dpas_w8a16_policy_m_32` | +| `> 128` | `128×128×16` | `dpas_w8a16_policy` | + +上面几档对齐的是 S4 的 *decode* 阶梯,而不是 FP8 prefill 的那条 —— +后者的 `≤ 512 → m_32` 档是按 prefill 规模的 batch 调过的。 + +*常驻 atomic 计数器。* prefill dispatch 每次调用都用 `sycl::malloc_device` +分配 work-group 计数器、再用 `sycl::free` 释放,这两个操作各会强制一次队列 +同步。在 prefill 规模下这只是噪声,但在 decode 规模下 —— GEMM 本身只有几十 +微秒、且每生成一个 token 就要发一次调用 —— 它占总时间的比例相当可观。 +decode dispatch 改用每队列常驻的一个 slot(`get_persistent_atomic_buffer`, +现已与 S4 头文件共享,两条路径共用一份 cache)。走上这条快路径时还会跳过 +`fill_expert_id_per_token` 前置 pass,因为 DPAS dispatch 直接消费 +`num_tokens_per_expert` —— decode 时间线上少一次 kernel 启动。 +**状态:NEEDS-HARDWARE-VALIDATION**(该头文件是未经硬件验证的移植)。 + +**占用率门控 —— 真实 decode batch 仍走 scalar GEMV。** 理由与 int4-sym +相同:decode 阶梯能选到的最小 tile 每个专家处理 8 行 token,所以平均每专家 +不足 8 个 token 时,tile 大部分是 padding。这正是 decode 的场景(MiniMax-M2, +192 个专家:每专家 0.04–1.3 个 token),因此除非 batch 平均每专家至少提供 +8 个 token,FP8 decode 一律走 scalar GEMV。 +`ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` 可覆盖该阈值;`0` 关闭门控,这也是对齐 +用例与 A/B 性能用例所设置的值。未通过 per-group 形状门控 +(`N%64==0`、`K%32==0`、`K%group_size==0`、 +`group_size ∈ {32,64,128,256}`)的形状始终回退到 scalar GEMV。 + +| Env 变量 | 默认值 | 作用 | +| -------- | ------ | ---- | +| `ARK_FP8_DECODE_MODE` | `word` | scalar GEMV 的 FP8 解码实现:`word`(位域搬移 + 折叠 scale bias)、`lut`(128 项幅值表)、`bits`(内联位运算)。 | +| `ARK_FP8_DECODE_USE_LUT` | 未设置 | 旧的选择开关;当它被显式设置、且 `ARK_FP8_DECODE_MODE` 未设置或取值无法识别时仍然生效:truthy → `lut`,falsy → `bits`。它同时仍然驱动 mixed-input prefill 路径。 | +| `ARK_MOE_DECODE_DPAS_FP8` | ON | 形状与占用率门控都通过时,把 FP8 decode 路由到 per-group DPAS grouped GEMM;`0` 强制走 scalar GEMV。 | +| `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | 走 DPAS 路径所需的最小每专家 token 数;`0` 关闭门控(对齐/A-B 用例所设)。 | + +性能 A/B 行是 `test_moe_decode_perf.py::test_perf_fp8_word_vs_lut` +(`speedup` 为 `lut / word`)与 `::test_perf_fp8_dpas_vs_scalar` +(`speedup` 为 `scalar / dpas`)。正确性由 +`test_moe.py::test_decode_fp8_modes_match`(三种解码器互相一致,且各自都 +对齐 dequant 参考)与 `::test_decode_fp8_dpas_matches_scalar` 覆盖。 + ## FP8 per-expert (per-tensor) 性能测试 `test_perf_fp8_per_tensor` 提供 Variant A DPAS 路径的性能表格,对应 diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index e1eb623531..005643da8b 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -1025,6 +1025,118 @@ def test_decode_fp8(self, dtype, fp8_dtype, group_size): atol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 torch.testing.assert_close(out, ref, rtol=rtol, atol=atol) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + @pytest.mark.parametrize("group_size", [32, 128]) + def test_decode_fp8_modes_match(self, monkeypatch, dtype, fp8_dtype, group_size): + """FP8 decode: the word-native decoder (``ARK_FP8_DECODE_MODE=word``, + the default) must be numerically identical to the LUT and inline-bits + decoders. + + ``word`` converts each FP8 byte to an fp16 bit pattern with a pure + bit-field move and folds E4M3's residual ``2**-8`` into the per-K-group + scale (an exact power of two), so on the finite encodings a real + checkpoint contains all three modes decode to exactly the same value. + Only the fp32 accumulation order differs (``word`` uses two partial + accumulators), hence ``bitwise=False`` and a tight -- not exact -- + tolerance. + """ + num_experts = 4 + tokens_per_expert = [1, 0, 2, 1] + total_tokens = sum(tokens_per_expert) + N, K = 256, 256 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + group_size=group_size, + asym=False, + ) + + # Keep every mode on the scalar GEMV so this compares decoders only. + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.delenv("ARK_FP8_DECODE_USE_LUT", raising=False) + + monkeypatch.setenv("ARK_FP8_DECODE_MODE", "word") + out_word = _run() + monkeypatch.setenv("ARK_FP8_DECODE_MODE", "lut") + out_lut = _run() + monkeypatch.setenv("ARK_FP8_DECODE_MODE", "bits") + out_bits = _run() + + assert out_word.shape == (total_tokens, N) + torch.testing.assert_close(out_word, out_lut, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(out_word, out_bits, rtol=1e-3, atol=1e-3) + + # Every mode must still track the dequant reference. + dequant = _dequant_fp8(packed, scales, group_size, dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + rtol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 + atol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 + torch.testing.assert_close(out_word, ref, rtol=rtol, atol=atol) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + @pytest.mark.parametrize("group_size", [32, 128]) + def test_decode_fp8_dpas_matches_scalar(self, monkeypatch, dtype, fp8_dtype, group_size): + """FP8 decode: the per-group DPAS path (``ARK_MOE_DECODE_DPAS_FP8=1``, + the default) must match both the scalar GEMV fallback + (``ARK_MOE_DECODE_DPAS_FP8=0``) and the dequant->bmm reference. + + Shapes satisfy the DPAS per-group shape gate (N%64==0, K%32==0, + group_size in {32,64,128,256}) and ``ARK_MOE_DECODE_DPAS_FP8_MIN_TPE=0`` + disables the tokens-per-expert occupancy gate (these tiny token counts + would otherwise be routed to the scalar GEMV, which is faster there) so + the DPAS fast path is actually taken. + """ + num_experts = 4 + tokens_per_expert = [1, 0, 2, 1] + total_tokens = sum(tokens_per_expert) + N, K = 320, 256 # N%64==0, K%32==0 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "1") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8_MIN_TPE", "0") + out_dpas = _run() + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + out_scalar = _run() + + dequant = _dequant_fp8(packed, scales, group_size, dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + rtol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 + atol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 + + assert out_dpas.shape == (total_tokens, N) + assert out_scalar.shape == (total_tokens, N) + torch.testing.assert_close(out_dpas, ref, rtol=rtol, atol=atol) + torch.testing.assert_close(out_scalar, ref, rtol=rtol, atol=atol) + torch.testing.assert_close(out_dpas, out_scalar, rtol=rtol, atol=atol) + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index 60861b407e..bd4da41a25 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -695,7 +695,113 @@ def test_perf_fp8(self, dtype, fp8_dtype): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) - def test_perf_fp8_per_tensor(self, dtype, fp8_dtype): + def test_perf_fp8_word_vs_lut(self, monkeypatch, dtype, fp8_dtype): + """FP8 scalar-GEMV A/B: word-native decode (``ARK_FP8_DECODE_MODE=word``, + the default) vs the 128-entry magnitude LUT (``=lut``, the old default). + + ``speedup`` is ``lut / word`` (the word-native decoder is the "ark" + column). The decode GEMV does roughly one multiply-add per weight byte, + so the dequant *is* the kernel: the LUT path issues a memory load per + weight element and both legacy paths index an 8-bit-typed + ``sycl::vec``, which Xe's 32-bit ALU lanes cannot address + directly. The word-native path reads the same bytes as + ``sycl::vec`` and turns each 32-bit word into four fp16 bit + patterns with a couple of native DWORD ops, folding E4M3's residual + ``2**-8`` into the per-K-group scale. This is the same treatment that + made int4-sym decode fast (see ``decode_int4_octet``). + + The FP8 DPAS fast path is disabled so both columns run the scalar GEMV + -- the only path the decode-mode flag affects. + """ + group_size = 128 + _print_header( + f"FP8 {str(fp8_dtype).split('.')[-1]} word vs lut decode (group_size={group_size}, " + f"act={str(dtype).split('.')[-1]}) -- LUT GEMV (baseline) vs word-native GEMV (ark)" + ) + for label, E, tpe, N, K in DECODE_SHAPES: + if K % group_size != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.delenv("ARK_FP8_DECODE_USE_LUT", raising=False) + monkeypatch.setenv("ARK_FP8_DECODE_MODE", "lut") + lut_ms = _xpu_time_ms(_run) + monkeypatch.setenv("ARK_FP8_DECODE_MODE", "word") + word_ms = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_FP8_DECODE_MODE", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) + _print_row(label, N, K, total_tokens, lut_ms, word_ms) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + def test_perf_fp8_dpas_vs_scalar(self, monkeypatch, dtype, fp8_dtype): + """FP8 decode: the per-group DPAS grouped GEMM + (``ARK_MOE_DECODE_DPAS_FP8=1``) vs the scalar GEMV (``=0``). + + ``speedup`` is ``scalar / dpas`` (the DPAS path is the "ark" column). + Only shapes that clear the DPAS shape gate are timed. + ``ARK_MOE_DECODE_DPAS_FP8_MIN_TPE=0`` disables the tokens-per-expert + occupancy gate so the DPAS column really runs DPAS (by default these + decode-sized batches are routed to the scalar GEMV). + + This is the FP8 twin of ``test_perf_int4_sym_dpas_vs_scalar`` and + exists for the same reason: to locate the batch size where the DPAS + pipeline overtakes the scalar GEMV, which is what the default + occupancy threshold encodes. On MiniMax-M2 decode shapes (192 experts, + 0.04-1.3 tokens/expert) the DPAS M tile is starved -- even the 8-row + ``dpas_w4a16_policy_m_8`` bucket this decode ladder adds on top of the + reference ``w8a16`` ladder -- so the scalar column is expected to win + there. + """ + group_size = 128 + _print_header( + f"FP8 {str(fp8_dtype).split('.')[-1]} DPAS vs scalar (group_size={group_size}, " + f"act={str(dtype).split('.')[-1]}) -- scalar GEMV (baseline) vs FP8 DPAS (ark)" + ) + for label, E, tpe, N, K in DECODE_SHAPES: + if K % group_size != 0 or N % 64 != 0 or K % 32 != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + scalar_ms = _xpu_time_ms(_run) + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "1") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8_MIN_TPE", "0") + dpas_ms = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8_MIN_TPE", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) + _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) """Perf: FP8 per-expert (per-tensor) scale for the decode path. The C++ decode kernel does NOT expose a native ``[E]`` per-tensor From c6f67a131161551711badb309657576b675207ed Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:07:58 +0000 Subject: [PATCH 43/47] fix: restore test_perf_fp8_per_tensor definition; make fp8_decode_mode allocation-free Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_dequant.hpp | 21 +++++++++++++------ .../ark/test/test_moe_decode_perf.py | 4 ++++ 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp index 736ab0e3c0..7ba9c8fc61 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_dequant.hpp @@ -415,16 +415,25 @@ inline bool fp8_decode_use_lut() { // // Re-read on every call (not cached) so tests and benchmarks can toggle the // path in-process; the result is passed into the kernel as a template argument, -// so there is no per-element runtime branch. +// so there is no per-element runtime branch. The comparison is done in place +// rather than via `std::string` because decode issues one call per generated +// token, and a heap allocation per lookup on that path buys nothing (same +// reasoning as `moe_decode_detail::env_flag_enabled`). // ---------------------------------------------------------------------------- inline Fp8DecodeMode fp8_decode_mode() { const char* mode = std::getenv("ARK_FP8_DECODE_MODE"); if (mode != nullptr) { - std::string s(mode); - for (char& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "word") return Fp8DecodeMode::kWord; - if (s == "lut") return Fp8DecodeMode::kLut; - if (s == "bits") return Fp8DecodeMode::kBits; + auto iequals = [](const char* value, const char* lowercase_literal) { + const char* a = value; + const char* b = lowercase_literal; + for (; *a != '\0' && *b != '\0'; ++a, ++b) { + if (static_cast(std::tolower(static_cast(*a))) != *b) return false; + } + return *a == '\0' && *b == '\0'; + }; + if (iequals(mode, "word")) return Fp8DecodeMode::kWord; + if (iequals(mode, "lut")) return Fp8DecodeMode::kLut; + if (iequals(mode, "bits")) return Fp8DecodeMode::kBits; // Unrecognised value: fall through to the legacy variable / default. } if (std::getenv("ARK_FP8_DECODE_USE_LUT") != nullptr) { diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index bd4da41a25..1efdc6ccc1 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -802,6 +802,10 @@ def _run(): monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8_MIN_TPE", raising=False) monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) _print_row(label, N, K, total_tokens, scalar_ms, dpas_ms) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + def test_perf_fp8_per_tensor(self, dtype, fp8_dtype): """Perf: FP8 per-expert (per-tensor) scale for the decode path. The C++ decode kernel does NOT expose a native ``[E]`` per-tensor From 7102f0eca367cf29e8259ab7395281aa43b76547 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:38:05 +0000 Subject: [PATCH 44/47] perf: add K-split lane-parallel FP8 MoE decode GEMV Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 199 +++++++++++++++++- auto_round_extension/ark/test/test_moe.py | 100 +++++++++ .../ark/test/test_moe_decode_perf.py | 56 +++++ 3 files changed, 348 insertions(+), 7 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 1c03de22ef..c8af8e40c8 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -34,6 +34,12 @@ // so no cross-lane reduction is needed and activation reads are coalesced across // the sub-group through the L1 cache. // +// The FP8 scalar path additionally offers a K-split mapping (one sub-group per +// output element, lanes splitting K, `ARK_MOE_DECODE_FP8_KSPLIT`, default ON): +// it trades a sub-group reduction for fully coalesced weight loads and 16x the +// thread count, which is what the memory-bound decode GEMV is short of. See the +// block comment above `launch_fp8_ksplit`. +// // Copyright (C) 2026 Intel Corporation // SPDX-License-Identifier: Apache-2.0 @@ -154,6 +160,9 @@ class MoEDecodeKernelInt2; template class MoEDecodeKernelFP8; +template +class MoEDecodeKernelFP8KSplit; + // ---------------------------------------------------------------------------- // FP8 weight dequantization primitives + host-side env-var reader live in // `sycl_tla_moe_dequant.hpp` so the prefill (mixed-input Grouped GEMM) and @@ -1241,6 +1250,10 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // // * `kLut` / `kBits` -- the original per-byte `decode_fp8` // decoders, kept for A/B measurement and regression escape. +// +// Two lane mappings share those decoders: the legacy per-work-item GEMV +// (`launch_fp8`) and the K-split GEMV (`launch_fp8_ksplit`, default, see its +// block comment). `launch_fp8_by_mode` picks between them. // ---------------------------------------------------------------------------- // Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK weight @@ -1398,26 +1411,198 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh }); } +// ---------------------------------------------------------------------------- +// FP8 K-split (lane-parallel) decode GEMV. +// +// `launch_fp8` above maps one *work-item* to one output element, so a lane +// walks a whole `[n_global, K]` weight row on its own. Two things follow from +// that mapping, and both cost real bandwidth on a kernel that does ~1 MAC per +// weight byte: +// +// 1. Weight loads are not coalesced. Lane `l` and lane `l+1` of a sub-group +// read bytes that are `K` apart, so every 16-byte load instruction turns +// into 16 scattered cache-line requests. The lines are eventually fully +// consumed (each lane walks its own row sequentially), so no DRAM byte is +// wasted -- but the memory controller sees `16 x resident sub-groups` +// independent streams instead of one per thread, which is exactly the +// access pattern DRAM row buffers handle worst. +// 2. The grid is small. Decode runs `total_tokens * N / 16` sub-groups; for +// a batch-1 MiniMax-M2 step (8 tokens, N=1536) that is 768 SIMD16 +// threads, below the thread slots of a BMG-class GPU, so there are not +// enough outstanding loads in flight to cover DRAM latency. +// +// This kernel transposes the lane mapping: a whole *sub-group* cooperates on +// one output element and the lanes split K. Lane `l` owns the `KSPLIT_CH` +// consecutive K elements at `l * KSPLIT_CH` inside each `KSPLIT_STEP`-wide +// K-tile, so per instruction the sub-group covers `KSPLIT_STEP` *contiguous* +// weight bytes (256 B -- four full cache lines) and `2 * KSPLIT_STEP` +// contiguous activation bytes. Each thread now walks a single sequential +// stream, and the grid grows by `SG_SIZE` (12288 sub-groups for the batch-1 +// step above), which is what puts enough requests in flight. The per-lane +// partial sums are reduced once at the end with `reduce_over_group` -- a +// handful of shuffles per output element against `K` multiply-adds. +// +// The int4 fallback solves the same coalescing problem by repacking the packed +// weights into an N-tiled layout (`launch_int4_coalesced`), which costs a full +// pass over the weight tensor and is therefore gated on a token-count +// amortization heuristic. FP8 weights are one byte per element and already +// K-contiguous, so K-splitting the lane mapping gets the same coalescing with +// no repack, no scratch buffer and no extra kernel launch. +// +// Scale handling: a lane's chunk is `KSPLIT_CH` consecutive K elements +// starting at a multiple of `KSPLIT_CH`, so with `group_size` a power of two +// that is >= `KSPLIT_CH` (the shape gate below) the chunk always sits inside a +// single K-group and its scale index is `k0 >> log2(group_size)` -- one shift, +// no integer division in the hot loop. The scale is applied per chunk instead +// of once per group; that is one extra multiply per `KSPLIT_CH` elements and +// keeps the `Sigma a * (w * s) == s * Sigma a * w` fold exact-per-group, +// including the folded `2^-8` E4M3 word-decode bias. +// ---------------------------------------------------------------------------- + +// K elements a lane owns per step. 16 FP8 bytes = one 16-byte weight load and +// one `vec` (32-byte) activation load per lane, i.e. exactly the +// transactions `fp8_decode_chunk` already issues, so the alignment contract is +// unchanged. +constexpr int KSPLIT_CH = 16; +// K elements a sub-group covers per step: the contiguous span its 16 lanes +// read in one instruction. +constexpr int KSPLIT_STEP = SG_SIZE * KSPLIT_CH; +// Sub-groups per work-group. Each owns one output column, so a work-group +// covers `N_TILE` consecutive columns and `N % N_TILE == 0` (already required +// by every decode path) is enough to tile N exactly. +constexpr int KSPLIT_WG_SGS = N_TILE; + +// ---------------------------------------------------------------------------- +// Env-flag helper -- `ARK_MOE_DECODE_FP8_KSPLIT` (default ON). When ON, the FP8 +// scalar decode GEMV uses the K-split kernel above; setting the var to "0" / +// "false" / "off" / "no" (case-insensitive) forces the legacy per-lane-strided +// `launch_fp8`, for A/B comparison and regression escape. Re-read on every call +// so tests and benchmarks can toggle the path in-process. +// ---------------------------------------------------------------------------- +inline bool moe_decode_fp8_ksplit_enabled() { + return env_flag_enabled("ARK_MOE_DECODE_FP8_KSPLIT", true); // default ON +} + +// Shape gate for the K-split kernel. `group_size` must be a power of two of at +// least `KSPLIT_CH` so that (a) a lane's chunk never straddles a K-group +// boundary and (b) the group index is a shift rather than an integer division +// on the hot path. Every shipped FP8 quant config (32 / 64 / 128 / 256) passes; +// anything else keeps the legacy GEMV, which handles arbitrary group sizes. +inline bool moe_decode_fp8_ksplit_shape_ok(int N, int K, int group_size) { + if (N % N_TILE != 0) return false; + if (group_size < KSPLIT_CH) return false; + if ((group_size & (group_size - 1)) != 0) return false; // not a power of two + if (K % group_size != 0) return false; + return true; +} + +template +void launch_fp8_ksplit(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, + ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, int K, + int group_size) { + if (!moe_decode_fp8_ksplit_shape_ok(N, K, group_size)) { + throw std::invalid_argument("moe_gemm_decode(fp8): K-split GEMV called on an unsupported shape"); + } + if (total_tokens == 0) return; + + const int num_groups_k = K / group_size; + int log2_group = 0; + while ((1 << log2_group) < group_size) ++log2_group; + // Undoes the exponent re-bias the word-native decode leaves behind (1.0f for + // every other mode). Exact power of two, applied once per lane chunk. + constexpr float kScaleBias = (Mode == Fp8DecodeMode::kWord) ? fp8_word_scale_bias() : 1.0f; + + // One sub-group per (token, output column); `KSPLIT_WG_SGS` of them per + // work-group so the dispatcher sees `N / N_TILE` work-groups per token + // instead of `N` single-sub-group ones. + sycl::range<2> global{static_cast(total_tokens), static_cast(N) * SG_SIZE}; + sycl::range<2> local{1, static_cast(KSPLIT_WG_SGS * SG_SIZE)}; + + q->parallel_for>( + sycl::nd_range<2>(global, local), + [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { + const auto sg = it.get_sub_group(); + const int token = static_cast(it.get_global_id(0)); + const int local_id = static_cast(it.get_local_id(1)); + // The work-group is one row of `KSPLIT_WG_SGS * SG_SIZE` work-items, so + // sub-group index and lane index are just the halves of the local id. + const int lane = local_id % SG_SIZE; + const int n_global = static_cast(it.get_group(1)) * KSPLIT_WG_SGS + local_id / SG_SIZE; + + const int expert = expert_id_per_token[token]; + const ScalarT* act_row = activations + static_cast(token) * K; + const uint8_t* w_row = + weights + (static_cast(expert) * N + static_cast(n_global)) * K; + const ScalarT* s_row = + scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + + // Each lane accumulates the scaled partial dot product of the chunks + // it owns; `fp8_decode_chunk` keeps two partial accumulators per chunk + // so the fp32 dependency chain stays broken. + float acc = 0.0f; + int k0 = lane * KSPLIT_CH; + // Two chunks per iteration: their loads are independent, so the pair + // doubles the number of weight requests a thread keeps in flight. + for (; k0 + KSPLIT_STEP + KSPLIT_CH <= K; k0 += 2 * KSPLIT_STEP) { + float a0 = 0.0f, a1 = 0.0f, b0 = 0.0f, b1 = 0.0f; + fp8_decode_chunk(act_row + k0, w_row + k0, a0, a1); + fp8_decode_chunk(act_row + k0 + KSPLIT_STEP, + w_row + k0 + KSPLIT_STEP, b0, b1); + const float s0 = static_cast(s_row[k0 >> log2_group]) * kScaleBias; + const float s1 = static_cast(s_row[(k0 + KSPLIT_STEP) >> log2_group]) * kScaleBias; + acc += (a0 + a1) * s0 + (b0 + b1) * s1; + } + // Remainder (at most one chunk per lane, plus the lanes whose chunk + // falls past K when K < KSPLIT_STEP -- those simply contribute 0). + for (; k0 < K; k0 += KSPLIT_STEP) { + float p0 = 0.0f, p1 = 0.0f; + fp8_decode_chunk(act_row + k0, w_row + k0, p0, p1); + acc += (p0 + p1) * (static_cast(s_row[k0 >> log2_group]) * kScaleBias); + } + + const float total = sycl::reduce_over_group(sg, acc, sycl::plus{}); + if (lane == 0) { + outputs[static_cast(token) * N + n_global] = static_cast(total); + } + }); +} + // Runtime -> compile-time bridge for the decode-mode selector. Keeps the // `moe_gemm_decode` dispatch to one branch per (act dtype, format) instead of -// re-nesting the mode selection at every call site. +// re-nesting the mode selection at every call site. The K-split vs legacy +// choice is made here as well, so all three decode modes run the same kernel +// structure and `word` / `lut` / `bits` stay comparable to one another. +template +void launch_fp8_dispatch(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, + const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, + int total_tokens, int N, int K, int group_size, bool ksplit) { + if (ksplit) { + launch_fp8_ksplit(q, activations, weights, scales, outputs, expert_id_per_token, + total_tokens, N, K, group_size); + } else { + launch_fp8(q, activations, weights, scales, outputs, expert_id_per_token, + total_tokens, N, K, group_size); + } +} + template void launch_fp8_by_mode(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, int K, int group_size) { + const bool ksplit = moe_decode_fp8_ksplit_enabled() && moe_decode_fp8_ksplit_shape_ok(N, K, group_size); switch (fp8_decode_mode()) { case Fp8DecodeMode::kLut: - launch_fp8(q, activations, weights, scales, outputs, - expert_id_per_token, total_tokens, N, K, group_size); + launch_fp8_dispatch( + q, activations, weights, scales, outputs, expert_id_per_token, total_tokens, N, K, group_size, ksplit); return; case Fp8DecodeMode::kBits: - launch_fp8(q, activations, weights, scales, outputs, - expert_id_per_token, total_tokens, N, K, group_size); + launch_fp8_dispatch( + q, activations, weights, scales, outputs, expert_id_per_token, total_tokens, N, K, group_size, ksplit); return; case Fp8DecodeMode::kWord: default: - launch_fp8(q, activations, weights, scales, outputs, - expert_id_per_token, total_tokens, N, K, group_size); + launch_fp8_dispatch( + q, activations, weights, scales, outputs, expert_id_per_token, total_tokens, N, K, group_size, ksplit); return; } } diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 005643da8b..790c7e7aba 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -1137,6 +1137,106 @@ def _run(): torch.testing.assert_close(out_scalar, ref, rtol=rtol, atol=atol) torch.testing.assert_close(out_dpas, out_scalar, rtol=rtol, atol=atol) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + @pytest.mark.parametrize("group_size", [32, 128]) + def test_decode_fp8_ksplit_matches_strided(self, monkeypatch, dtype, fp8_dtype, group_size): + """FP8 decode: the K-split GEMV (``ARK_MOE_DECODE_FP8_KSPLIT=1``, the + default) must match the legacy per-work-item GEMV + (``ARK_MOE_DECODE_FP8_KSPLIT=0``) and the dequant reference. + + The two kernels compute the same dot products with a different lane + mapping: the legacy one gives a whole K row to one work-item, the + K-split one gives one output element to a whole sub-group and splits K + across its 16 lanes, then reduces. Only the fp32 summation order + differs, so the two outputs are compared at the same tight tolerance + used by ``test_decode_fp8_modes_match``. + + K is deliberately larger than one sub-group step (16 lanes x 16 + elements = 256) so the kernel's unrolled main loop *and* its remainder + loop are both exercised; ``group_size`` covers a group narrower and a + group wider than a lane's 16-element chunk. + """ + num_experts = 4 + tokens_per_expert = [1, 0, 2, 1] + total_tokens = sum(tokens_per_expert) + N, K = 256, 640 # K = 2*256 + 128 -> main loop plus a partial step + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + group_size=group_size, + asym=False, + ) + + # Compare the two scalar-GEMV lane mappings only: keep DPAS off. + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "1") + out_ksplit = _run() + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "0") + out_strided = _run() + + dequant = _dequant_fp8(packed, scales, group_size, dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + rtol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 + atol = 1e-1 if fp8_dtype == torch.float8_e5m2 else 5e-2 + + assert out_ksplit.shape == (total_tokens, N) + torch.testing.assert_close(out_ksplit, out_strided, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(out_ksplit, ref, rtol=rtol, atol=atol) + torch.testing.assert_close(out_strided, ref, rtol=rtol, atol=atol) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("group_size", [48, 96]) + def test_decode_fp8_ksplit_non_pow2_group_falls_back(self, monkeypatch, dtype, group_size): + """FP8 decode: a non-power-of-two ``group_size`` must still be correct. + + The K-split GEMV indexes the scale array with a shift, which is only + valid when ``group_size`` is a power of two of at least 16 (a lane's + 16-element chunk must sit inside one K-group). Other group sizes take + the legacy GEMV; this pins that fallback so a future gate change cannot + silently start feeding them to the shift-indexed kernel. + + ``group_size`` stays a multiple of 16 because the shared vectorized + inner loop reads 16-element blocks from the start of every K-group. + """ + num_experts = 4 + tokens_per_expert = [1, 0, 2, 1] + total_tokens = sum(tokens_per_expert) + N, K = 256, 480 # divisible by both 48 and 96 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, torch.float8_e4m3fn) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "1") + out = ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + group_size=group_size, + asym=False, + ) + + dequant = _dequant_fp8(packed, scales, group_size, dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + assert out.shape == (total_tokens, N) + torch.testing.assert_close(out, ref, rtol=5e-2, atol=5e-2) + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index 1efdc6ccc1..a9b6b0723c 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -748,6 +748,62 @@ def _run(): monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) _print_row(label, N, K, total_tokens, lut_ms, word_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + def test_perf_fp8_ksplit_vs_strided(self, monkeypatch, dtype, fp8_dtype): + """FP8 scalar-GEMV A/B: K-split lane mapping + (``ARK_MOE_DECODE_FP8_KSPLIT=1``, the default) vs the legacy + per-work-item mapping (``=0``). + + ``speedup`` is ``strided / ksplit`` (the K-split kernel is the "ark" + column). Decode does ~1 multiply-add per weight byte, so the kernel is + bound by how fast the weight tile streams in, not by arithmetic. The + legacy mapping gives each work-item its own ``[n, K]`` weight row, so + the 16 lanes of a sub-group read bytes ``K`` apart and one load + instruction touches 16 cache lines; it also launches only + ``total_tokens * N / 16`` threads, too few to keep enough loads in + flight to hide DRAM latency at batch 1. The K-split mapping gives one + output element to a whole sub-group and splits K across its lanes, so a + load instruction covers 256 contiguous weight bytes and the thread + count grows 16x, at the cost of one sub-group reduction per output + element. + + The FP8 DPAS fast path is disabled so both columns run the scalar GEMV. + """ + group_size = 128 + _print_header( + f"FP8 {str(fp8_dtype).split('.')[-1]} K-split vs strided GEMV (group_size={group_size}, " + f"act={str(dtype).split('.')[-1]}) -- per-work-item GEMV (baseline) vs K-split GEMV (ark)" + ) + for label, E, tpe, N, K in DECODE_SHAPES: + if K % group_size != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "0") + strided_ms = _xpu_time_ms(_run) + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "1") + ksplit_ms = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_MOE_DECODE_FP8_KSPLIT", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) + _print_row(label, N, K, total_tokens, strided_ms, ksplit_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) def test_perf_fp8_dpas_vs_scalar(self, monkeypatch, dtype, fp8_dtype): From e43b98f7d59c3244329944aae53d5cc0e2312f76 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:44:42 +0000 Subject: [PATCH 45/47] docs: describe the K-split FP8 decode GEMV (EN + CN) Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../wrapper/include/sycl_tla_moe_decode.hpp | 10 ++-- .../ark/test/README_MOE_PREFILL_PERF.md | 50 ++++++++++++++++++- .../ark/test/README_MOE_PREFILL_PERF_CN.md | 44 +++++++++++++++- 3 files changed, 97 insertions(+), 7 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index c8af8e40c8..8793da4098 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -1474,7 +1474,7 @@ constexpr int KSPLIT_WG_SGS = N_TILE; // ---------------------------------------------------------------------------- // Env-flag helper -- `ARK_MOE_DECODE_FP8_KSPLIT` (default ON). When ON, the FP8 -// scalar decode GEMV uses the K-split kernel above; setting the var to "0" / +// scalar decode GEMV uses the K-split kernel below; setting the var to "0" / // "false" / "off" / "no" (case-insensitive) forces the legacy per-lane-strided // `launch_fp8`, for A/B comparison and regression escape. Re-read on every call // so tests and benchmarks can toggle the path in-process. @@ -1488,8 +1488,12 @@ inline bool moe_decode_fp8_ksplit_enabled() { // boundary and (b) the group index is a shift rather than an integer division // on the hot path. Every shipped FP8 quant config (32 / 64 / 128 / 256) passes; // anything else keeps the legacy GEMV, which handles arbitrary group sizes. +// `K >= KSPLIT_STEP` additionally keeps every lane of the sub-group busy -- +// below that some lanes own no chunk at all and only pay the reduction, which +// is the one regime where splitting K cannot pay for itself. inline bool moe_decode_fp8_ksplit_shape_ok(int N, int K, int group_size) { if (N % N_TILE != 0) return false; + if (K < KSPLIT_STEP) return false; if (group_size < KSPLIT_CH) return false; if ((group_size & (group_size - 1)) != 0) return false; // not a power of two if (K % group_size != 0) return false; @@ -1552,8 +1556,8 @@ void launch_fp8_ksplit(sycl::queue* q, const ScalarT* activations, const uint8_t const float s1 = static_cast(s_row[(k0 + KSPLIT_STEP) >> log2_group]) * kScaleBias; acc += (a0 + a1) * s0 + (b0 + b1) * s1; } - // Remainder (at most one chunk per lane, plus the lanes whose chunk - // falls past K when K < KSPLIT_STEP -- those simply contribute 0). + // Remainder: the lanes whose last chunk does not have a partner a + // full step away. At most one chunk per lane given the shape gate. for (; k0 < K; k0 += KSPLIT_STEP) { float p0 = 0.0f, p1 = 0.0f; fp8_decode_chunk(act_row + k0, w_row + k0, p0, p1); 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 5d5177426a..c9d8524912 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -513,6 +513,48 @@ checkpoints are produced by scaling to `finfo(float8_e4m3fn).max == 448` and clamping, so those two encodings cannot occur. Callers that need NaN propagation can select `ARK_FP8_DECODE_MODE=lut` or `=bits`. +**K-split lane mapping (`ARK_MOE_DECODE_FP8_KSPLIT`, default ON).** Once +the dequant is a couple of DWORD ops, the scalar GEMV is purely a +bandwidth problem: at ~1 multiply-add per weight byte it can only run as +fast as the expert tile streams in. The original mapping gave one output +element to one *work-item*, so a lane walked a whole `[n, K]` weight row on +its own. Two costs follow: + +* **Uncoalesced weight loads.** Lanes `l` and `l+1` of a sub-group read + bytes `K` apart, so each 16-byte load instruction is split into 16 + cache-line requests. No DRAM byte is wasted (each lane consumes its lines + as it walks the row) but the memory controller sees 16 independent + streams per sub-group, the pattern DRAM row buffers handle worst. +* **Too few threads.** The grid is `total_tokens × N / 16` sub-groups — + 768 SIMD16 threads for a MiniMax-M2 batch-1 step (8 tokens, N=1536), + below the thread slots of a BMG-class GPU, so there are never enough + loads in flight to hide DRAM latency. + +`launch_fp8_ksplit` transposes the mapping: one *sub-group* per output +element, with the 16 lanes splitting K. Lane `l` owns the 16 consecutive K +elements at `l*16` inside each 256-element step, so one instruction covers +256 **contiguous** weight bytes (four full cache lines) and 512 contiguous +activation bytes, each thread walks a single sequential stream, and the +thread count grows 16× (12288 sub-groups for that batch-1 step). The price +is one `reduce_over_group` per output element — a handful of shuffles +against `K` multiply-adds — and 16× more activation traffic out of L1, +which has ample headroom at this arithmetic intensity. + +This is the same problem the int4 fallback solves by repacking weights +into an N-tiled layout (`ARK_MOE_DECODE_COALESCE_INT4`), which costs a +full extra pass over the weight tensor and a scratch buffer. FP8 weights +are one byte per element and already K-contiguous, so K-splitting the lane +mapping gets the same coalescing with no repack, no scratch and no extra +kernel launch. + +The kernel indexes the scale array with a shift, so the shape gate +requires a power-of-two `group_size ≥ 16` (every shipped FP8 config — 32 / +64 / 128 / 256 — passes) plus `N%16==0`, `K%group_size==0` and `K ≥ 256` (so +every lane of the sub-group owns at least one chunk); anything else keeps the +legacy GEMV, which handles arbitrary group sizes. All three +`ARK_FP8_DECODE_MODE` decoders run under both mappings, so the mode A/B +stays apples-to-apples. **Status: NEEDS-HARDWARE-VALIDATION.** + **FP8 DPAS decode dispatch.** `moe_decode_fp8_dpas_per_group_dispatch` (`sycl_tla_moe_prefill_fp8_dpas.hpp`, `ARK_MOE_DECODE_DPAS_FP8` default ON) is the FP8 twin of the S4 decode dispatch: same mainloop, same `[E, N, K]` @@ -567,12 +609,16 @@ per-group shape gate (`N%64==0`, `K%32==0`, `K%group_size==0`, | `ARK_FP8_DECODE_USE_LUT` | unset | Legacy selector, still honoured when set explicitly and when `ARK_FP8_DECODE_MODE` is unset/unrecognised: truthy → `lut`, falsy → `bits`. Also still drives the mixed-input prefill path. | | `ARK_MOE_DECODE_DPAS_FP8` | ON | Route FP8 decode to the per-group DPAS grouped GEMM when the shape and occupancy gates pass; `0` forces the scalar GEMV. | | `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | Minimum tokens per expert before the DPAS path is taken; `0` disables the gate (what the parity/A-B tests set). | +| `ARK_MOE_DECODE_FP8_KSPLIT` | ON | Scalar-GEMV lane mapping: one sub-group per output element with the lanes splitting K (coalesced weight loads, 16× the threads); `0` forces the legacy one-work-item-per-output-element GEMV. Shapes outside the gate (power-of-two `group_size ≥ 16`, `N%16==0`, `K%group_size==0`, `K ≥ 256`) always use the legacy mapping. | Perf A/B rows are `test_moe_decode_perf.py::test_perf_fp8_word_vs_lut` -(`speedup` is `lut / word`) and `::test_perf_fp8_dpas_vs_scalar` +(`speedup` is `lut / word`), `::test_perf_fp8_ksplit_vs_strided` +(`speedup` is `strided / ksplit`) and `::test_perf_fp8_dpas_vs_scalar` (`speedup` is `scalar / dpas`). Correctness is covered by `test_moe.py::test_decode_fp8_modes_match` (all three decoders agree, and -each tracks the dequant reference) and +each tracks the dequant reference), +`::test_decode_fp8_ksplit_matches_strided` (both lane mappings agree, plus +a non-power-of-two `group_size` fallback case) and `::test_decode_fp8_dpas_matches_scalar`. ## FP8 per-expert (per-tensor) perf tests 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 27c4badc83..e965f0e612 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 @@ -408,6 +408,42 @@ kernel 以 `sycl::vec` 读取权重 —— 与它替换掉的字节 `finfo(float8_e4m3fn).max == 448` 缩放并 clamp 得到的,所以这两个编码不可能 出现。需要 NaN 传播的调用方可以选择 `ARK_FP8_DECODE_MODE=lut` 或 `=bits`。 +**K-split lane 映射(`ARK_MOE_DECODE_FP8_KSPLIT`,默认 ON)。** 当 dequant +只剩几条 DWORD 运算之后,scalar GEMV 就是一个纯粹的带宽问题:每个权重字节 +大约只做一次乘加,所以它最快只能跑到专家 tile 的搬运速度。原来的映射把一个 +输出元素交给一个 *work-item*,于是一个 lane 要独自走完整条 `[n, K]` 权重行。 +由此带来两笔开销: + +* **权重访存不合并。** 同一 sub-group 中 lane `l` 与 lane `l+1` 读到的字节 + 相距 `K`,因此每条 16 字节的 load 指令都会被拆成 16 个 cache line 请求。 + DRAM 字节并没有浪费(每个 lane 会沿着自己的行把这些 line 用完),但内存 + 控制器看到的是每个 sub-group 16 条互相独立的数据流 —— 这正是 DRAM row + buffer 最不擅长的访问模式。 +* **线程太少。** grid 只有 `total_tokens × N / 16` 个 sub-group —— + MiniMax-M2 batch-1 一步(8 个 token,N=1536)只有 768 个 SIMD16 线程, + 低于 BMG 级 GPU 的线程槽数量,飞行中的 load 永远不足以掩盖 DRAM 延迟。 + +`launch_fp8_ksplit` 把映射转置过来:一个 *sub-group* 负责一个输出元素,由它 +的 16 个 lane 切分 K。lane `l` 在每个 256 元素的步长内拥有起点为 `l*16` 的 +16 个连续 K 元素,于是一条指令覆盖 256 字节**连续**权重(四条完整 cache +line)和 512 字节连续激活,每个线程只走一条顺序数据流,线程数则提升 16× +(上述 batch-1 场景为 12288 个 sub-group)。代价是每个输出元素一次 +`reduce_over_group` —— 相对 `K` 次乘加只是几条 shuffle —— 以及 16× 的 L1 +激活流量,而在这样的计算密度下 L1 有充足余量。 + +int4 的回退路径解决的是同一个问题,办法是把权重 repack 成 N-tiled 布局 +(`ARK_MOE_DECODE_COALESCE_INT4`),那需要额外完整扫一遍权重张量并占用 +scratch 显存。FP8 权重每元素一个字节、本来就是 K 连续的,所以只切分 lane +映射就能拿到同样的合并访存,无需 repack、无需 scratch、也不多一次 kernel +启动。 + +该 kernel 用移位来索引 scale 数组,因此形状门控要求 `group_size` 是 ≥ 16 的 +2 的幂(已发布的 FP8 配置 —— 32 / 64 / 128 / 256 —— 全部满足),另外还要 +`N%16==0`、`K%group_size==0` 以及 `K ≥ 256`(保证 sub-group 的每个 lane 至少 +分到一个 chunk);其余情况继续走老的 GEMV,它支持任意 group size。三种 `ARK_FP8_DECODE_MODE` 解码器在两种映射下都能运行,所以 +decode mode 的 A/B 依然是同口径对比。 +**状态:NEEDS-HARDWARE-VALIDATION。** + **FP8 DPAS decode dispatch。** `moe_decode_fp8_dpas_per_group_dispatch` (`sycl_tla_moe_prefill_fp8_dpas.hpp`,`ARK_MOE_DECODE_DPAS_FP8` 默认 ON) 是 S4 decode dispatch 的 FP8 对应物:同一套 mainloop、同样的 `[E, N, K]` @@ -457,12 +493,16 @@ decode dispatch 改用每队列常驻的一个 slot(`get_persistent_atomic_buffe | `ARK_FP8_DECODE_USE_LUT` | 未设置 | 旧的选择开关;当它被显式设置、且 `ARK_FP8_DECODE_MODE` 未设置或取值无法识别时仍然生效:truthy → `lut`,falsy → `bits`。它同时仍然驱动 mixed-input prefill 路径。 | | `ARK_MOE_DECODE_DPAS_FP8` | ON | 形状与占用率门控都通过时,把 FP8 decode 路由到 per-group DPAS grouped GEMM;`0` 强制走 scalar GEMV。 | | `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | 走 DPAS 路径所需的最小每专家 token 数;`0` 关闭门控(对齐/A-B 用例所设)。 | +| `ARK_MOE_DECODE_FP8_KSPLIT` | ON | scalar GEMV 的 lane 映射:一个 sub-group 负责一个输出元素、由 lane 切分 K(访存合并,线程数 ×16);`0` 强制走老的「一个 work-item 一个输出元素」GEMV。未通过门控(`group_size` 为 ≥ 16 的 2 的幂、`N%16==0`、`K%group_size==0`、`K ≥ 256`)的形状始终使用老映射。 | 性能 A/B 行是 `test_moe_decode_perf.py::test_perf_fp8_word_vs_lut` -(`speedup` 为 `lut / word`)与 `::test_perf_fp8_dpas_vs_scalar` +(`speedup` 为 `lut / word`)、`::test_perf_fp8_ksplit_vs_strided` +(`speedup` 为 `strided / ksplit`)与 `::test_perf_fp8_dpas_vs_scalar` (`speedup` 为 `scalar / dpas`)。正确性由 `test_moe.py::test_decode_fp8_modes_match`(三种解码器互相一致,且各自都 -对齐 dequant 参考)与 `::test_decode_fp8_dpas_matches_scalar` 覆盖。 +对齐 dequant 参考)、`::test_decode_fp8_ksplit_matches_strided`(两种 lane +映射一致,并覆盖非 2 的幂 `group_size` 的回退)与 +`::test_decode_fp8_dpas_matches_scalar` 覆盖。 ## FP8 per-expert (per-tensor) 性能测试 From 51328f366531f58e197691722cdb3e25173282c5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:48:24 +0000 Subject: [PATCH 46/47] perf: N-block the FP8 K-split decode GEMV and drop the per-call routing sync Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/auto_round_kernel/__init__.py | 58 +++- .../wrapper/include/sycl_tla_moe_decode.hpp | 285 +++++++++++++----- .../ark/test/README_MOE_PREFILL_PERF.md | 54 +++- .../ark/test/README_MOE_PREFILL_PERF_CN.md | 48 ++- auto_round_extension/ark/test/test_moe.py | 116 +++++++ .../ark/test/test_moe_decode_perf.py | 63 ++++ 6 files changed, 528 insertions(+), 96 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 0a485bd7c8..9ead9ff265 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -1610,7 +1610,10 @@ def moe_gemm_decode( ``K_packed == K``. ``weight_bits`` is ignored; ``asym`` must be ``False`` (no zero-points for FP8). num_tokens_per_expert: ``[E]`` int32. Sum must equal - ``activations.shape[0]``. + ``activations.shape[0]``; this is a caller contract. It is checked + eagerly only when the tensor lives on the host, or when + ``ARK_MOE_VALIDATE_ROUTING`` is set -- summing a device tensor + costs a blocking device-to-host sync on the decode hot path. scales: ``[E, N, K // group_size]`` in activations dtype. Required for all quantized paths (int8/int4/int2/fp8); must be ``None`` for unquantized weights. @@ -1686,6 +1689,43 @@ def moe_decode_release_scratch() -> None: lib.moe_decode_release_scratch() +def moe_routing_validation_enabled() -> bool: + """Whether ``num_tokens_per_expert`` is checked against ``total_tokens``. + + The check needs the *sum* of the routing table, which for a table that + already lives on the device costs a reduction kernel plus a blocking + device-to-host copy -- a full pipeline flush on every call. Decode issues + one call per generated token, so that sync lands directly in the + token-latency path (and inside the timed region of the decode benchmarks), + where it is worth tens of microseconds against kernels that take ~150us. + + So the check runs unconditionally for host-side (CPU) routing tables, where + it is free, and is skipped for device tables unless + ``ARK_MOE_VALIDATE_ROUTING`` is set to a truthy value. The C++ side does not + need the host value: it consumes the device pointer directly and derives + ``expert_id_per_token`` on-device, clamped to ``num_experts - 1``. + + Truthy values (case-insensitive): anything other than "0", "false", "off", + "no". Unset means disabled (no sync). + """ + env = os.environ.get("ARK_MOE_VALIDATE_ROUTING") + if env is None: + return False + return env.strip().lower() not in ("0", "false", "off", "no") + + +def _check_routing_total(num_tokens_per_expert: torch.Tensor, total_tokens: int) -> None: + """Check ``sum(num_tokens_per_expert) == total_tokens`` without a device sync. + + See :func:`moe_routing_validation_enabled` for when the check is skipped. + """ + if num_tokens_per_expert.device.type != "cpu" and not moe_routing_validation_enabled(): + return + expected_total = int(num_tokens_per_expert.sum().item()) + if expected_total != total_tokens: + raise ValueError(f"Sum of num_tokens_per_expert ({expected_total}) != total_tokens ({total_tokens})") + + def _validate_moe_quant_args( activations: torch.Tensor, weights: torch.Tensor, @@ -1704,6 +1744,10 @@ def _validate_moe_quant_args( kernel-call site: ``(activations, weights, scales, zeros, num_tokens_per_expert, weight_dtype, total_tokens, N, K, num_experts)``. + + The caller owns the contract that ``num_tokens_per_expert`` sums to + ``activations.shape[0]``; see :func:`moe_routing_validation_enabled` for how + that is (or is not) enforced. """ if activations.device.type != "xpu": raise NotImplementedError(f"{api_name} is only supported on XPU") @@ -1817,9 +1861,7 @@ def _validate_moe_quant_args( if N % 16 != 0: raise ValueError(f"N must be a multiple of 16 (got {N})") - expected_total = int(num_tokens_per_expert.sum().item()) - if expected_total != total_tokens: - raise ValueError(f"Sum of num_tokens_per_expert ({expected_total}) != total_tokens ({total_tokens})") + _check_routing_total(num_tokens_per_expert, total_tokens) return (activations, weights, scales, zeros, num_tokens_per_expert, weight_dtype, total_tokens, N, K, num_experts) @@ -1875,9 +1917,7 @@ def moe_gemm( raise ValueError(f"num_tokens_per_expert length {num_tokens_per_expert.shape[0]} != num_experts {num_experts}") # Validate total tokens - expected_total = int(num_tokens_per_expert.sum().item()) - if expected_total != total_tokens: - raise ValueError(f"Sum of num_tokens_per_expert ({expected_total}) != total_tokens ({total_tokens})") + _check_routing_total(num_tokens_per_expert, total_tokens) lib = get_lib(activations) stream = get_stream(activations) @@ -2110,7 +2150,7 @@ def moe_gemm_prefill( ``[E, N, K]`` -- callers providing already-``[E, K, N]`` weights (as ``moe_gemm`` requires) should call ``moe_gemm`` directly. num_tokens_per_expert: ``[E]`` int32. Sum must equal - ``activations.shape[0]``. + ``activations.shape[0]`` (see :func:`moe_gemm_decode`). scales: ``[E, N, K // group_size]`` in activations dtype. Required for quantized paths; ignored (must be ``None``) for unquantized. zeros: ``[E, N, K // group_size]`` in activations dtype, required when @@ -2427,7 +2467,7 @@ def moe( weights: ``[E, N, K_packed]`` -- see :func:`moe_gemm_decode` for the quant-specific layout/dtype contract. num_tokens_per_expert: ``[E]`` int32. Sum must equal - ``activations.shape[0]``. + ``activations.shape[0]`` (see :func:`moe_gemm_decode`). scales, zeros, weight_bits, group_size, asym: forwarded to the underlying kernel; see :func:`moe_gemm_decode`. phase: dispatch mode. diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp index 8793da4098..65d1089d7e 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_decode.hpp @@ -53,6 +53,7 @@ #include #include #include +#include #include "bestla/bestla.h" #include "sycl_tla_moe_dequant.hpp" @@ -160,7 +161,7 @@ class MoEDecodeKernelInt2; template class MoEDecodeKernelFP8; -template +template class MoEDecodeKernelFP8KSplit; // ---------------------------------------------------------------------------- @@ -1256,69 +1257,92 @@ void launch_int2(sycl::queue* q, const ScalarT* activations, const uint8_t* weig // block comment). `launch_fp8_by_mode` picks between them. // ---------------------------------------------------------------------------- +// Activation / weight vector types for one 16-element FP8 sub-block, plus the +// MAC that consumes them. Splitting "load" from "multiply-accumulate" lets a +// caller issue several independent loads before any of them is consumed -- +// which is what the N-blocked K-split kernel below needs to keep more than one +// weight request per thread in flight. `fp8_decode_chunk` is a thin +// load-then-MAC wrapper over these, so both callers run identical arithmetic. +constexpr int FP8_SUB = 16; + +using Fp8ActVec16 = sycl::vec; + +// `kWord` mode reads the 16 weight bytes as four 32-bit words so the decode +// never leaves the native datapath; the other modes read them as bytes. Either +// way it is the same single 16-byte transaction with the same 16-byte alignment +// requirement. +template +using Fp8WeightVec16 = + std::conditional_t, sycl::vec>; + +template +static inline Fp8ActVec16 load_fp8_act_vec16(const ScalarT* act_ptr) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + return *reinterpret_cast(act_ptr); +} + +template +static inline Fp8WeightVec16 load_fp8_weight_vec16(const uint8_t* w_ptr) { + return *reinterpret_cast*>(w_ptr); +} + +// Accumulate `FP8_SUB` products of an already-loaded activation / weight pair. +// Two independent partial accumulators break the single fp32 dependency chain +// so the FMA pipeline is not latency-bound (same trick as `int4_decode_chunk`); +// the caller reduces the pair. +template +static inline void fp8_mac_vec16(const Fp8ActVec16& av, const Fp8WeightVec16& wv, float& acc0, float& acc1) { + if constexpr (Mode == Fp8DecodeMode::kWord) { + constexpr int WORDS = FP8_SUB / 4; // one 32-bit word per 4 FP8 bytes +#pragma unroll + for (int w = 0; w < WORDS; ++w) { + uint32_t lo2, hi2; + decode_fp8_quad_half_bits(wv[w], lo2, hi2); + const uint16_t hb[4] = {static_cast(lo2), static_cast(lo2 >> 16), + static_cast(hi2), static_cast(hi2 >> 16)}; +#pragma unroll + for (int u = 0; u < 4; u += 2) { + const ScalarT a0 = sycl::bit_cast(static_cast(av[4 * w + u])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[4 * w + u + 1])); + acc0 += static_cast(a0) * static_cast(sycl::bit_cast(hb[u])); + acc1 += static_cast(a1) * static_cast(sycl::bit_cast(hb[u + 1])); + } + } + } else { + constexpr bool kUseLut = (Mode == Fp8DecodeMode::kLut); +#pragma unroll + for (int u = 0; u < FP8_SUB; u += 2) { + const float w0 = decode_fp8(wv[u]); + const float w1 = decode_fp8(wv[u + 1]); + const ScalarT a0 = sycl::bit_cast(static_cast(av[u])); + const ScalarT a1 = sycl::bit_cast(static_cast(av[u + 1])); + acc0 += static_cast(a0) * w0; + acc1 += static_cast(a1) * w1; + } + } +} + // Vectorized inner accumulation over CHUNK consecutive K elements (CHUNK weight // bytes + a vec activation block). Templated on CHUNK so the // caller can run a wide (32) stage first and a narrower (16) stage for the // remainder, mirroring the int4/int8 paths. sycl::vec only supports widths of // 1, 2, 3, 4, 8 or 16, so CHUNK is processed in 16-wide sub-blocks. // -// In `kWord` mode the 16 weight bytes of a sub-block are read as a -// `sycl::vec` -- the same 16-byte transaction (and the same -// 16-byte alignment requirement) as the byte vector it replaces, but 32-bit -// typed, so the decode never leaves the native datapath. -// // The per-group scale is constant across the whole group, so it is NOT applied // here: this accumulates the raw dot product (sum of act * decoded_fp8) and the // caller multiplies the group total by the scale once (Σ a·(w·s) == s·Σ a·w). // For the per-expert / per-tensor scale case (group_size == K, one scale per // output row) this collapses the whole K reduction to a single scale multiply, // removing one multiply per K element on the decode hot path. -// -// Two independent partial accumulators break the single fp32 dependency chain -// so the FMA pipeline is not latency-bound (same trick as `int4_decode_chunk`); -// the caller reduces the pair. template static inline void fp8_decode_chunk(const ScalarT* act_ptr, const uint8_t* w_ptr, float& acc0, float& acc1) { static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); - static_assert(CHUNK % 16 == 0, "CHUNK must be a multiple of 16"); - constexpr int SUB = 16; - using ActVec = sycl::vec; -#pragma unroll - for (int s = 0; s < CHUNK / SUB; ++s) { - const ActVec av = *reinterpret_cast(act_ptr + s * SUB); - if constexpr (Mode == Fp8DecodeMode::kWord) { - constexpr int WORDS = SUB / 4; // one 32-bit word per 4 FP8 bytes - using WordVec = sycl::vec; - const WordVec wv = *reinterpret_cast(w_ptr + s * SUB); -#pragma unroll - for (int w = 0; w < WORDS; ++w) { - uint32_t lo2, hi2; - decode_fp8_quad_half_bits(wv[w], lo2, hi2); - const uint16_t hb[4] = {static_cast(lo2), static_cast(lo2 >> 16), - static_cast(hi2), static_cast(hi2 >> 16)}; -#pragma unroll - for (int u = 0; u < 4; u += 2) { - const ScalarT a0 = sycl::bit_cast(static_cast(av[4 * w + u])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[4 * w + u + 1])); - acc0 += static_cast(a0) * static_cast(sycl::bit_cast(hb[u])); - acc1 += static_cast(a1) * static_cast(sycl::bit_cast(hb[u + 1])); - } - } - } else { - constexpr bool kUseLut = (Mode == Fp8DecodeMode::kLut); - using ByteVec = sycl::vec; - const ByteVec wv = *reinterpret_cast(w_ptr + s * SUB); + static_assert(CHUNK % FP8_SUB == 0, "CHUNK must be a multiple of 16"); #pragma unroll - for (int u = 0; u < SUB; u += 2) { - const float w0 = decode_fp8(wv[u]); - const float w1 = decode_fp8(wv[u + 1]); - const ScalarT a0 = sycl::bit_cast(static_cast(av[u])); - const ScalarT a1 = sycl::bit_cast(static_cast(av[u + 1])); - acc0 += static_cast(a0) * w0; - acc1 += static_cast(a1) * w1; - } - } + for (int s = 0; s < CHUNK / FP8_SUB; ++s) { + fp8_mac_vec16(load_fp8_act_vec16(act_ptr + s * FP8_SUB), + load_fp8_weight_vec16(w_ptr + s * FP8_SUB), acc0, acc1); } } @@ -1457,6 +1481,10 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // of once per group; that is one extra multiply per `KSPLIT_CH` elements and // keeps the `Sigma a * (w * s) == s * Sigma a * w` fold exact-per-group, // including the folded `2^-8` E4M3 word-decode bias. +// +// On top of that mapping the sub-group also blocks N: it owns `NCOLS` +// consecutive output columns and reuses one activation load across all of them +// (see `moe_decode_fp8_ksplit_ncols`). // ---------------------------------------------------------------------------- // K elements a lane owns per step. 16 FP8 bytes = one 16-byte weight load and @@ -1467,9 +1495,8 @@ constexpr int KSPLIT_CH = 16; // K elements a sub-group covers per step: the contiguous span its 16 lanes // read in one instruction. constexpr int KSPLIT_STEP = SG_SIZE * KSPLIT_CH; -// Sub-groups per work-group. Each owns one output column, so a work-group -// covers `N_TILE` consecutive columns and `N % N_TILE == 0` (already required -// by every decode path) is enough to tile N exactly. +// Sub-groups per work-group. Each owns `NCOLS` output columns, so a work-group +// covers `N_TILE * NCOLS` consecutive columns. constexpr int KSPLIT_WG_SGS = N_TILE; // ---------------------------------------------------------------------------- @@ -1500,11 +1527,55 @@ inline bool moe_decode_fp8_ksplit_shape_ok(int N, int K, int group_size) { return true; } -template +// ---------------------------------------------------------------------------- +// N-blocking factor: output columns a sub-group owns. +// +// With one column per sub-group the hot loop issues, per 16-byte weight chunk, +// one weight message *and* one 32-byte activation message -- half the traffic a +// thread requests is the activation row, which every column of that token +// re-reads. Giving a sub-group NCOLS consecutive columns loads the activation +// chunk once and reuses it for all NCOLS weight chunks, so +// +// activation messages per weight chunk: 1 -> 1 / NCOLS +// independent weight loads in flight: 2 -> 2 * NCOLS +// +// The first effect cuts request-queue pressure; the second raises memory-level +// parallelism, which is what a pure-streaming GEMV is actually limited by (the +// measured kernel sits well below peak DRAM bandwidth, so it is latency- and +// message-bound, not bandwidth-bound). The cost is NCOLS times the live weight +// vectors and accumulators, so the factor is kept small. +// +// A work-group still holds `KSPLIT_WG_SGS` sub-groups, so it now covers +// `KSPLIT_WG_SGS * NCOLS` columns and N must divide by that. NCOLS == 1 +// reproduces the previous kernel instruction-for-instruction. +// `ARK_MOE_DECODE_FP8_KSPLIT_NCOLS` overrides the default (accepted values 1, 2 +// and 4); anything else, or a factor the shape cannot tile, falls back to the +// largest valid smaller power of two. +// ---------------------------------------------------------------------------- +constexpr int KSPLIT_NCOLS_DEFAULT = 2; +constexpr int KSPLIT_NCOLS_MAX = 4; + +inline int moe_decode_fp8_ksplit_ncols(int N) { + int ncols = KSPLIT_NCOLS_DEFAULT; + const char* env = std::getenv("ARK_MOE_DECODE_FP8_KSPLIT_NCOLS"); + if (env != nullptr) { + char* end = nullptr; + const long long v = std::strtoll(env, &end, 10); + if (end != env && v >= 1 && v <= KSPLIT_NCOLS_MAX && (v & (v - 1)) == 0) { + ncols = static_cast(v); + } + } + // A work-group covers `KSPLIT_WG_SGS * ncols` columns; shrink until it tiles. + while (ncols > 1 && (N % (KSPLIT_WG_SGS * ncols)) != 0) ncols /= 2; + return ncols; +} + +template void launch_fp8_ksplit(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, int K, int group_size) { - if (!moe_decode_fp8_ksplit_shape_ok(N, K, group_size)) { + static_assert(NCOLS >= 1 && (NCOLS & (NCOLS - 1)) == 0, "NCOLS must be a power of two"); + if (!moe_decode_fp8_ksplit_shape_ok(N, K, group_size) || (N % (KSPLIT_WG_SGS * NCOLS)) != 0) { throw std::invalid_argument("moe_gemm_decode(fp8): K-split GEMV called on an unsupported shape"); } if (total_tokens == 0) return; @@ -1516,13 +1587,13 @@ void launch_fp8_ksplit(sycl::queue* q, const ScalarT* activations, const uint8_t // every other mode). Exact power of two, applied once per lane chunk. constexpr float kScaleBias = (Mode == Fp8DecodeMode::kWord) ? fp8_word_scale_bias() : 1.0f; - // One sub-group per (token, output column); `KSPLIT_WG_SGS` of them per - // work-group so the dispatcher sees `N / N_TILE` work-groups per token - // instead of `N` single-sub-group ones. - sycl::range<2> global{static_cast(total_tokens), static_cast(N) * SG_SIZE}; + // One sub-group per (token, NCOLS output columns); `KSPLIT_WG_SGS` of them + // per work-group so the dispatcher sees `N / (N_TILE * NCOLS)` work-groups + // per token instead of `N` single-sub-group ones. + sycl::range<2> global{static_cast(total_tokens), static_cast(N / NCOLS) * SG_SIZE}; sycl::range<2> local{1, static_cast(KSPLIT_WG_SGS * SG_SIZE)}; - q->parallel_for>( + q->parallel_for>( sycl::nd_range<2>(global, local), [=](sycl::nd_item<2> it) [[intel::reqd_sub_group_size(SG_SIZE)]] { const auto sg = it.get_sub_group(); @@ -1531,46 +1602,98 @@ void launch_fp8_ksplit(sycl::queue* q, const ScalarT* activations, const uint8_t // The work-group is one row of `KSPLIT_WG_SGS * SG_SIZE` work-items, so // sub-group index and lane index are just the halves of the local id. const int lane = local_id % SG_SIZE; - const int n_global = static_cast(it.get_group(1)) * KSPLIT_WG_SGS + local_id / SG_SIZE; + const int n_base = + (static_cast(it.get_group(1)) * KSPLIT_WG_SGS + local_id / SG_SIZE) * NCOLS; const int expert = expert_id_per_token[token]; const ScalarT* act_row = activations + static_cast(token) * K; - const uint8_t* w_row = - weights + (static_cast(expert) * N + static_cast(n_global)) * K; - const ScalarT* s_row = - scales + (static_cast(expert) * N + static_cast(n_global)) * num_groups_k; + const size_t row0 = (static_cast(expert) * N + static_cast(n_base)); + const uint8_t* w_rows[NCOLS]; + const ScalarT* s_rows[NCOLS]; +#pragma unroll + for (int c = 0; c < NCOLS; ++c) { + w_rows[c] = weights + (row0 + static_cast(c)) * K; + s_rows[c] = scales + (row0 + static_cast(c)) * num_groups_k; + } // Each lane accumulates the scaled partial dot product of the chunks - // it owns; `fp8_decode_chunk` keeps two partial accumulators per chunk - // so the fp32 dependency chain stays broken. - float acc = 0.0f; + // it owns, for each of its NCOLS columns; `fp8_mac_vec16` keeps two + // partial accumulators per chunk so the fp32 dependency chain stays + // broken. NCOLS is a compile-time constant, so `acc` and the staged + // weight vectors below stay in registers. + float acc[NCOLS]; +#pragma unroll + for (int c = 0; c < NCOLS; ++c) acc[c] = 0.0f; + int k0 = lane * KSPLIT_CH; // Two chunks per iteration: their loads are independent, so the pair - // doubles the number of weight requests a thread keeps in flight. + // doubles the number of weight requests a thread keeps in flight. All + // 2 * NCOLS weight loads are issued before the first is consumed. for (; k0 + KSPLIT_STEP + KSPLIT_CH <= K; k0 += 2 * KSPLIT_STEP) { - float a0 = 0.0f, a1 = 0.0f, b0 = 0.0f, b1 = 0.0f; - fp8_decode_chunk(act_row + k0, w_row + k0, a0, a1); - fp8_decode_chunk(act_row + k0 + KSPLIT_STEP, - w_row + k0 + KSPLIT_STEP, b0, b1); - const float s0 = static_cast(s_row[k0 >> log2_group]) * kScaleBias; - const float s1 = static_cast(s_row[(k0 + KSPLIT_STEP) >> log2_group]) * kScaleBias; - acc += (a0 + a1) * s0 + (b0 + b1) * s1; + const Fp8ActVec16 av0 = load_fp8_act_vec16(act_row + k0); + const Fp8ActVec16 av1 = load_fp8_act_vec16(act_row + k0 + KSPLIT_STEP); + Fp8WeightVec16 wv0[NCOLS], wv1[NCOLS]; +#pragma unroll + for (int c = 0; c < NCOLS; ++c) { + wv0[c] = load_fp8_weight_vec16(w_rows[c] + k0); + wv1[c] = load_fp8_weight_vec16(w_rows[c] + k0 + KSPLIT_STEP); + } + const int g0 = k0 >> log2_group; + const int g1 = (k0 + KSPLIT_STEP) >> log2_group; +#pragma unroll + for (int c = 0; c < NCOLS; ++c) { + float a0 = 0.0f, a1 = 0.0f, b0 = 0.0f, b1 = 0.0f; + fp8_mac_vec16(av0, wv0[c], a0, a1); + fp8_mac_vec16(av1, wv1[c], b0, b1); + const float s0 = static_cast(s_rows[c][g0]) * kScaleBias; + const float s1 = static_cast(s_rows[c][g1]) * kScaleBias; + acc[c] += (a0 + a1) * s0 + (b0 + b1) * s1; + } } // Remainder: the lanes whose last chunk does not have a partner a // full step away. At most one chunk per lane given the shape gate. for (; k0 < K; k0 += KSPLIT_STEP) { - float p0 = 0.0f, p1 = 0.0f; - fp8_decode_chunk(act_row + k0, w_row + k0, p0, p1); - acc += (p0 + p1) * (static_cast(s_row[k0 >> log2_group]) * kScaleBias); + const Fp8ActVec16 av = load_fp8_act_vec16(act_row + k0); + const int g = k0 >> log2_group; +#pragma unroll + for (int c = 0; c < NCOLS; ++c) { + float p0 = 0.0f, p1 = 0.0f; + fp8_mac_vec16(av, load_fp8_weight_vec16(w_rows[c] + k0), p0, p1); + acc[c] += (p0 + p1) * (static_cast(s_rows[c][g]) * kScaleBias); + } } - const float total = sycl::reduce_over_group(sg, acc, sycl::plus{}); - if (lane == 0) { - outputs[static_cast(token) * N + n_global] = static_cast(total); +#pragma unroll + for (int c = 0; c < NCOLS; ++c) { + const float total = sycl::reduce_over_group(sg, acc[c], sycl::plus{}); + if (lane == 0) { + outputs[static_cast(token) * N + n_base + c] = static_cast(total); + } } }); } +// Runtime NCOLS -> compile-time NCOLS bridge. +template +void launch_fp8_ksplit_by_ncols(sycl::queue* q, const ScalarT* activations, const uint8_t* weights, + const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, + int total_tokens, int N, int K, int group_size) { + switch (moe_decode_fp8_ksplit_ncols(N)) { + case 4: + launch_fp8_ksplit(q, activations, weights, scales, outputs, expert_id_per_token, + total_tokens, N, K, group_size); + return; + case 2: + launch_fp8_ksplit(q, activations, weights, scales, outputs, expert_id_per_token, + total_tokens, N, K, group_size); + return; + default: + launch_fp8_ksplit(q, activations, weights, scales, outputs, expert_id_per_token, + total_tokens, N, K, group_size); + return; + } +} + // Runtime -> compile-time bridge for the decode-mode selector. Keeps the // `moe_gemm_decode` dispatch to one branch per (act dtype, format) instead of // re-nesting the mode selection at every call site. The K-split vs legacy @@ -1581,8 +1704,8 @@ void launch_fp8_dispatch(sycl::queue* q, const ScalarT* activations, const uint8 const ScalarT* scales, ScalarT* outputs, const int* expert_id_per_token, int total_tokens, int N, int K, int group_size, bool ksplit) { if (ksplit) { - launch_fp8_ksplit(q, activations, weights, scales, outputs, expert_id_per_token, - total_tokens, N, K, group_size); + launch_fp8_ksplit_by_ncols(q, activations, weights, scales, outputs, + expert_id_per_token, total_tokens, N, K, group_size); } else { launch_fp8(q, activations, weights, scales, outputs, expert_id_per_token, total_tokens, N, K, group_size); 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 c9d8524912..8655e53a63 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -555,6 +555,48 @@ legacy GEMV, which handles arbitrary group sizes. All three `ARK_FP8_DECODE_MODE` decoders run under both mappings, so the mode A/B stays apples-to-apples. **Status: NEEDS-HARDWARE-VALIDATION.** +**N-blocking inside the K-split kernel (`ARK_MOE_DECODE_FP8_KSPLIT_NCOLS`, +default 2).** With one output column per sub-group the hot loop issues, per +16-byte weight chunk, one weight message *and* one 32-byte activation +message — half of what a thread requests is the activation row, which every +column of that token re-reads — and only two weight loads are ever in +flight. Giving a sub-group `NCOLS` consecutive columns loads the activation +chunk once and reuses it for all of them: + +| | `NCOLS=1` | `NCOLS=n` | +| --- | --- | --- | +| activation messages per weight chunk | 1 | 1/n | +| independent weight loads in flight | 2 | 2n | + +The first effect cuts request-queue pressure; the second raises +memory-level parallelism, which is what a streaming GEMV sitting well below +peak DRAM bandwidth is actually limited by. The cost is `n` times the live +weight vectors and accumulators, so past some point the kernel spills — hence +the small ladder (1, 2, 4) and the conservative default. + +A work-group still holds 16 sub-groups, so it now covers `16 * NCOLS` +columns; an `N` that cannot be tiled at the requested factor falls back to +the largest valid smaller power of two on the host side (`N=1536` and +`N=3072` tile at every factor). The lane → K-chunk mapping, the per-chunk +scale fold and the final `reduce_over_group` are untouched, so the +arithmetic per output element is unchanged and `NCOLS=1` reproduces the +previous kernel exactly. `test_perf_fp8_ksplit_ncols_sweep` prints all three +factors per shape so the default can be set from measured data. +**Status: NEEDS-HARDWARE-VALIDATION.** + +**Routing-table validation (`ARK_MOE_VALIDATE_ROUTING`, default OFF).** The +Python entry point used to check `sum(num_tokens_per_expert) == total_tokens` +on every call. For a routing table that already lives on the device that +sum means a reduction kernel plus a *blocking* device-to-host copy, i.e. a +full pipeline flush — on a decode step whose kernel takes ~150 µs, and once +per generated token. It also lands inside the timed region of every decode +benchmark, because the queue is idle when the timing event is recorded. +The sum is now a caller contract (the C++ side never needed the host value: +it consumes the device pointer and derives `expert_id_per_token` on-device, +clamped to `num_experts - 1`); set `ARK_MOE_VALIDATE_ROUTING=1` to restore +the eager check when debugging a router. Host-side (CPU) routing tables are +still checked unconditionally, since summing those is free. + **FP8 DPAS decode dispatch.** `moe_decode_fp8_dpas_per_group_dispatch` (`sycl_tla_moe_prefill_fp8_dpas.hpp`, `ARK_MOE_DECODE_DPAS_FP8` default ON) is the FP8 twin of the S4 decode dispatch: same mainloop, same `[E, N, K]` @@ -610,15 +652,21 @@ per-group shape gate (`N%64==0`, `K%32==0`, `K%group_size==0`, | `ARK_MOE_DECODE_DPAS_FP8` | ON | Route FP8 decode to the per-group DPAS grouped GEMM when the shape and occupancy gates pass; `0` forces the scalar GEMV. | | `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | Minimum tokens per expert before the DPAS path is taken; `0` disables the gate (what the parity/A-B tests set). | | `ARK_MOE_DECODE_FP8_KSPLIT` | ON | Scalar-GEMV lane mapping: one sub-group per output element with the lanes splitting K (coalesced weight loads, 16× the threads); `0` forces the legacy one-work-item-per-output-element GEMV. Shapes outside the gate (power-of-two `group_size ≥ 16`, `N%16==0`, `K%group_size==0`, `K ≥ 256`) always use the legacy mapping. | +| `ARK_MOE_DECODE_FP8_KSPLIT_NCOLS` | `2` | Output columns one sub-group owns in the K-split GEMV (1, 2 or 4). Higher values reuse one activation load across more columns and keep more weight loads in flight, at the cost of more live registers. An `N` that `16 * NCOLS` cannot tile falls back to the largest valid smaller power of two. | +| `ARK_MOE_VALIDATE_ROUTING` | OFF | Eagerly check `sum(num_tokens_per_expert) == activations.shape[0]` for device-resident routing tables. The check costs a blocking device-to-host sync per call, so it is opt-in; CPU-resident tables are always checked. | Perf A/B rows are `test_moe_decode_perf.py::test_perf_fp8_word_vs_lut` (`speedup` is `lut / word`), `::test_perf_fp8_ksplit_vs_strided` -(`speedup` is `strided / ksplit`) and `::test_perf_fp8_dpas_vs_scalar` -(`speedup` is `scalar / dpas`). Correctness is covered by +(`speedup` is `strided / ksplit`), `::test_perf_fp8_ksplit_ncols_sweep` +(`speedup` is `NCOLS=1 / best NCOLS`, with all factors printed) and +`::test_perf_fp8_dpas_vs_scalar` (`speedup` is `scalar / dpas`). +Correctness is covered by `test_moe.py::test_decode_fp8_modes_match` (all three decoders agree, and each tracks the dequant reference), `::test_decode_fp8_ksplit_matches_strided` (both lane mappings agree, plus -a non-power-of-two `group_size` fallback case) and +a non-power-of-two `group_size` fallback case), +`::test_decode_fp8_ksplit_ncols_match` (every blocking factor agrees, plus +an untileable-`N` fallback case) and `::test_decode_fp8_dpas_matches_scalar`. ## FP8 per-expert (per-tensor) perf tests 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 e965f0e612..c33590c6a5 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 @@ -444,6 +444,43 @@ scratch 显存。FP8 权重每元素一个字节、本来就是 K 连续的,所 decode mode 的 A/B 依然是同口径对比。 **状态:NEEDS-HARDWARE-VALIDATION。** +**K-split kernel 内的 N 分块(`ARK_MOE_DECODE_FP8_KSPLIT_NCOLS`,默认 2)。** +当一个 sub-group 只负责一个输出列时,热循环中每读一个 16 字节权重 chunk, +既要发一条权重访存,又要发一条 32 字节的激活访存 —— 线程请求的数据里有一半 +是激活行,而该 token 的每一列都会重复读它 —— 并且飞行中的权重 load 始终只 +有两条。让一个 sub-group 负责 `NCOLS` 个连续列,激活 chunk 只需读一次就能 +被所有列复用: + +| | `NCOLS=1` | `NCOLS=n` | +| --- | --- | --- | +| 每个权重 chunk 的激活访存条数 | 1 | 1/n | +| 飞行中的独立权重 load | 2 | 2n | + +前者降低请求队列压力,后者提升 memory-level parallelism —— 对于一个远低于 +DRAM 峰值带宽的流式 GEMV,后者才是真正的瓶颈。代价是活跃的权重向量与 +累加器变成 `n` 倍,超过某个点 kernel 就会 spill,所以只提供 1、2、4 这个 +很短的阶梯,并且默认值取得保守。 + +一个 work-group 仍然是 16 个 sub-group,因此它现在覆盖 `16 * NCOLS` 列; +若 `N` 无法按所请求的因子切分,host 侧会回退到最大的、合法的更小 2 的幂 +(`N=1536` 与 `N=3072` 在所有因子下都能整除)。lane → K chunk 的映射、 +每个 chunk 的 scale 折叠以及最后的 `reduce_over_group` 都没有改动,因此 +单个输出元素的算术完全不变,`NCOLS=1` 与改动前的 kernel 完全一致。 +`test_perf_fp8_ksplit_ncols_sweep` 会逐形状打印全部三个因子的耗时,便于用 +实测数据确定默认值。 +**状态:NEEDS-HARDWARE-VALIDATION。** + +**路由表校验(`ARK_MOE_VALIDATE_ROUTING`,默认 OFF)。** Python 入口原先 +在每次调用时都会检查 `sum(num_tokens_per_expert) == total_tokens`。当路由表 +本身就在设备上时,这个求和意味着一次 reduction kernel 外加一次**阻塞式**的 +device-to-host 拷贝,也就是一次完整的流水线 flush —— 而 decode 一步的 kernel +本身只有约 150 µs,并且每生成一个 token 就要付一次。它同样落在 decode +benchmark 的计时区间内,因为记录计时 event 时队列正好是空的。 +现在这个求和关系是调用方契约(C++ 侧本来就不需要 host 上的值:它直接使用 +设备指针,并在设备上推导 `expert_id_per_token`,且会 clamp 到 +`num_experts - 1`);调试 router 时可设置 `ARK_MOE_VALIDATE_ROUTING=1` 恢复 +即时校验。位于 host(CPU)上的路由表仍然始终校验,因为对它们求和是免费的。 + **FP8 DPAS decode dispatch。** `moe_decode_fp8_dpas_per_group_dispatch` (`sycl_tla_moe_prefill_fp8_dpas.hpp`,`ARK_MOE_DECODE_DPAS_FP8` 默认 ON) 是 S4 decode dispatch 的 FP8 对应物:同一套 mainloop、同样的 `[E, N, K]` @@ -494,14 +531,19 @@ decode dispatch 改用每队列常驻的一个 slot(`get_persistent_atomic_buffe | `ARK_MOE_DECODE_DPAS_FP8` | ON | 形状与占用率门控都通过时,把 FP8 decode 路由到 per-group DPAS grouped GEMM;`0` 强制走 scalar GEMV。 | | `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | 走 DPAS 路径所需的最小每专家 token 数;`0` 关闭门控(对齐/A-B 用例所设)。 | | `ARK_MOE_DECODE_FP8_KSPLIT` | ON | scalar GEMV 的 lane 映射:一个 sub-group 负责一个输出元素、由 lane 切分 K(访存合并,线程数 ×16);`0` 强制走老的「一个 work-item 一个输出元素」GEMV。未通过门控(`group_size` 为 ≥ 16 的 2 的幂、`N%16==0`、`K%group_size==0`、`K ≥ 256`)的形状始终使用老映射。 | +| `ARK_MOE_DECODE_FP8_KSPLIT_NCOLS` | `2` | K-split GEMV 中一个 sub-group 负责的输出列数(1、2 或 4)。取值越大,一次激活 load 被复用的列越多、飞行中的权重 load 越多,代价是活跃寄存器更多。若 `16 * NCOLS` 无法整除 `N`,会回退到最大的、合法的更小 2 的幂。 | +| `ARK_MOE_VALIDATE_ROUTING` | OFF | 即时校验 `sum(num_tokens_per_expert) == activations.shape[0]`(针对位于设备上的路由表)。该校验每次调用都要付一次阻塞式 device-to-host 同步,因此改为按需开启;位于 CPU 上的路由表始终校验。 | 性能 A/B 行是 `test_moe_decode_perf.py::test_perf_fp8_word_vs_lut` (`speedup` 为 `lut / word`)、`::test_perf_fp8_ksplit_vs_strided` -(`speedup` 为 `strided / ksplit`)与 `::test_perf_fp8_dpas_vs_scalar` -(`speedup` 为 `scalar / dpas`)。正确性由 +(`speedup` 为 `strided / ksplit`)、`::test_perf_fp8_ksplit_ncols_sweep` +(`speedup` 为 `NCOLS=1 / 最优 NCOLS`,并打印全部因子)与 +`::test_perf_fp8_dpas_vs_scalar`(`speedup` 为 `scalar / dpas`)。正确性由 `test_moe.py::test_decode_fp8_modes_match`(三种解码器互相一致,且各自都 对齐 dequant 参考)、`::test_decode_fp8_ksplit_matches_strided`(两种 lane -映射一致,并覆盖非 2 的幂 `group_size` 的回退)与 +映射一致,并覆盖非 2 的幂 `group_size` 的回退)、 +`::test_decode_fp8_ksplit_ncols_match`(各分块因子结果一致,并覆盖 `N` +无法整除时的回退)与 `::test_decode_fp8_dpas_matches_scalar` 覆盖。 ## FP8 per-expert (per-tensor) 性能测试 diff --git a/auto_round_extension/ark/test/test_moe.py b/auto_round_extension/ark/test/test_moe.py index 790c7e7aba..fd6d04c346 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -872,6 +872,30 @@ def test_decode_validation_errors(self): asym=True, ) + def test_routing_sum_validation_is_opt_in(self, monkeypatch): + """A mismatched routing table is only rejected when asked for. + + Summing ``num_tokens_per_expert`` means pulling a device tensor back to + the host, which stalls the whole queue -- far too expensive to pay once + per generated token. The check is therefore opt-in via + ``ARK_MOE_VALIDATE_ROUTING``; otherwise the sum being + ``activations.shape[0]`` is a caller contract. + """ + activations = torch.randn(2, 128, dtype=torch.float16, device="xpu") + weights = torch.randn(2, 32, 128, dtype=torch.float16, device="xpu") + # Sums to 3, not to activations.shape[0] == 2. + bad_ntpe = torch.tensor([2, 1], dtype=torch.int32, device="xpu") + + monkeypatch.setenv("ARK_MOE_VALIDATE_ROUTING", "1") + assert ark.moe_routing_validation_enabled() is True + with pytest.raises(ValueError, match="num_tokens_per_expert"): + ark.moe_gemm_decode(activations, weights, bad_ntpe, weight_bits=16) + + monkeypatch.delenv("ARK_MOE_VALIDATE_ROUTING", raising=False) + assert ark.moe_routing_validation_enabled() is False + monkeypatch.setenv("ARK_MOE_VALIDATE_ROUTING", "0") + assert ark.moe_routing_validation_enabled() is False + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("group_size", [32, 128]) def test_decode_int8_sym(self, dtype, group_size): @@ -1237,6 +1261,98 @@ def test_decode_fp8_ksplit_non_pow2_group_falls_back(self, monkeypatch, dtype, g assert out.shape == (total_tokens, N) torch.testing.assert_close(out, ref, rtol=5e-2, atol=5e-2) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("group_size", [32, 128]) + def test_decode_fp8_ksplit_ncols_match(self, monkeypatch, dtype, group_size): + """FP8 decode: every N-blocking factor must produce the same result. + + ``ARK_MOE_DECODE_FP8_KSPLIT_NCOLS`` sets how many consecutive output + columns one sub-group owns. The factor only changes which columns share + an activation load -- the per-output arithmetic (lane -> K chunk + mapping, per-chunk scale fold, final ``reduce_over_group``) is + unchanged -- so all factors must agree with each other to within + fp-contraction noise, and with the dequant reference. + + The N chosen here (256) is divisible by ``16 * 4``, so no factor is + silently reduced by the host-side tiling fallback, and K is not a + multiple of the 256-element sub-group step so the remainder loop is + exercised for every factor. + """ + num_experts = 4 + tokens_per_expert = [1, 0, 2, 1] + total_tokens = sum(tokens_per_expert) + N, K = 256, 640 # 256 % (16*4) == 0; K = 2*256 + 128 + + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, torch.float8_e4m3fn) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "1") + + outs = {} + for ncols in (1, 2, 4): + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT_NCOLS", str(ncols)) + outs[ncols] = ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + group_size=group_size, + asym=False, + ) + + dequant = _dequant_fp8(packed, scales, group_size, dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + # One ulp of the output dtype at these magnitudes; any indexing bug + # (wrong column, dropped or double-counted K chunk) is orders of + # magnitude larger than this. + tight = 1e-3 if dtype == torch.float16 else 8e-3 + for ncols, out in outs.items(): + assert out.shape == (total_tokens, N), f"ncols={ncols}" + torch.testing.assert_close(out, outs[1], rtol=tight, atol=tight, msg=f"ncols={ncols} != ncols=1") + torch.testing.assert_close(out, ref, rtol=5e-2, atol=5e-2) + + @pytest.mark.parametrize("dtype", [torch.float16]) + def test_decode_fp8_ksplit_ncols_falls_back_on_untileable_n(self, monkeypatch, dtype): + """FP8 decode: an N that a blocking factor cannot tile must still work. + + A work-group covers ``16 * NCOLS`` columns, so ``N = 48`` only tiles + with ``NCOLS == 1`` and ``N = 160`` only with ``NCOLS <= 2``. The host + selector must shrink the requested factor instead of launching a grid + that walks past the last column. + """ + num_experts = 3 + tokens_per_expert = [2, 0, 1] + total_tokens = sum(tokens_per_expert) + group_size = 32 + for N in (48, 160): + K = 512 + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(num_experts, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(num_experts, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, torch.float8_e4m3fn) + num_tokens_per_expert = torch.tensor(tokens_per_expert, dtype=torch.int32, device="xpu") + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "1") + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT_NCOLS", "4") + out = ark.moe_gemm_decode( + activations, + packed, + num_tokens_per_expert, + scales=scales, + group_size=group_size, + asym=False, + ) + + dequant = _dequant_fp8(packed, scales, group_size, dtype) + ref = _moe_decode_reference(activations, dequant, num_tokens_per_expert) + assert out.shape == (total_tokens, N), f"N={N}" + torch.testing.assert_close(out, ref, rtol=5e-2, atol=5e-2) + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/auto_round_extension/ark/test/test_moe_decode_perf.py b/auto_round_extension/ark/test/test_moe_decode_perf.py index a9b6b0723c..ad348e3b78 100644 --- a/auto_round_extension/ark/test/test_moe_decode_perf.py +++ b/auto_round_extension/ark/test/test_moe_decode_perf.py @@ -804,6 +804,69 @@ def _run(): monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) _print_row(label, N, K, total_tokens, strided_ms, ksplit_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_perf_fp8_ksplit_ncols_sweep(self, monkeypatch, dtype): + """FP8 K-split GEMV: sweep the N-blocking factor. + + ``ARK_MOE_DECODE_FP8_KSPLIT_NCOLS`` sets how many consecutive output + columns one sub-group owns. With one column per sub-group, half of what + a thread requests is the activation row -- which every column of that + token re-reads -- and only two weight loads are ever in flight. Owning + ``NCOLS`` columns loads the activation chunk once for all of them and + puts ``2 * NCOLS`` independent weight loads in flight, which is what a + latency-bound streaming GEMV needs; the cost is ``NCOLS`` times the live + weight registers, so past some point the kernel spills. + + The ``ark`` column is the best factor found and ``baseline`` is + ``NCOLS=1`` (the pre-blocking kernel), so ``speedup`` is the win from + blocking alone. The per-factor timings are printed underneath so the + default (``KSPLIT_NCOLS_DEFAULT`` in ``sycl_tla_moe_decode.hpp``) can be + set from measured data rather than from the register-pressure estimate + it currently reflects. + """ + group_size = 128 + fp8_dtype = torch.float8_e4m3fn + factors = (1, 2, 4) + _print_header( + f"FP8 {str(fp8_dtype).split('.')[-1]} K-split N-blocking sweep " + f"(group_size={group_size}, act={str(dtype).split('.')[-1]}) " + f"-- NCOLS=1 (baseline) vs best NCOLS (ark)" + ) + for label, E, tpe, N, K in DECODE_SHAPES: + if K % group_size != 0: + continue + total_tokens = sum(tpe) + activations = torch.randn(total_tokens, K, dtype=dtype, device="xpu") + w_float = (torch.randn(E, N, K, dtype=torch.float32, device="xpu") * 0.1).to(dtype) + scales = torch.empty(E, N, K // group_size, dtype=dtype, device="xpu") + packed = _pack_fp8(w_float, scales, group_size, fp8_dtype) + ntpe = torch.tensor(tpe, dtype=torch.int32, device="xpu") + + def _run(): + return ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + group_size=group_size, + asym=False, + ) + + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_FP8", "0") + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT", "1") + per_factor = {} + for ncols in factors: + monkeypatch.setenv("ARK_MOE_DECODE_FP8_KSPLIT_NCOLS", str(ncols)) + per_factor[ncols] = _xpu_time_ms(_run) + monkeypatch.delenv("ARK_MOE_DECODE_FP8_KSPLIT_NCOLS", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_FP8_KSPLIT", raising=False) + monkeypatch.delenv("ARK_MOE_DECODE_DPAS_FP8", raising=False) + + best = min(per_factor, key=per_factor.get) + _print_row(label, N, K, total_tokens, per_factor[1], per_factor[best]) + detail = " ".join(f"NCOLS={n}: {per_factor[n]:.4f}ms" for n in factors) + print(f"{'':<18}{'':>7}{'':>7}{'':>8} {detail} best=NCOLS={best}") + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) def test_perf_fp8_dpas_vs_scalar(self, monkeypatch, dtype, fp8_dtype): From 03ca877400e26ec8f5ef140424a862c0b330341c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:32:40 +0000 Subject: [PATCH 47/47] perf: raise the MoE auto-dispatch decode cutoff from 32 to 128 tokens Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com> --- .../ark/auto_round_kernel/__init__.py | 23 ++++++----- .../ark/test/README_MOE_PREFILL_PERF.md | 26 ++++++++++++- .../ark/test/README_MOE_PREFILL_PERF_CN.md | 22 +++++++++-- .../ark/test/test_moe_unified.py | 38 ++++++++++++++++++- 4 files changed, 93 insertions(+), 16 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 9ead9ff265..f9f147583a 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -2405,15 +2405,18 @@ def _native_fp8_prefill_enabled() -> bool: # is hardware-dependent and can be overridden via # `ARK_MOE_AUTO_DECODE_MAX_TOKENS`. # -# Empirically (see `test_perf_int4_sym_dpas_vs_scalar`) the shared S4 DPAS -# grouped-GEMM that the prefill path uses already beats the scalar-GEMV decode -# kernel by ~2x once total tokens reach 256 (bs32), while it loses (~0.5x) only -# at the single-stream bs1 extreme (8 tokens). The crossover therefore sits far -# below 256, so the default is kept small: the scalar decode kernel is only -# preferred for the tiny single-/few-stream case where every expert sees well -# under one DPAS tile row. Mirrors vLLM-xpu-kernels' `w4a16` dispatch, which -# buckets on average tokens-per-expert rather than a large total-token cutoff. -_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS = 32 +# The cutoff used to be 32, a deliberately conservative value picked while the +# decode GEMV was still the bottleneck: back then only the tiny single-/few- +# stream case (every expert well under one DPAS tile row) was worth keeping off +# the prefill grouped-GEMM. The decode GEMV has since reached its bandwidth +# target for FP8 as well as int4-sym (K-split lane mapping + N-blocking inside +# the K-split kernel, and no per-call routing sync), so it now stays ahead of +# the grouped-GEMM over the whole small-batch range rather than only at the +# bs1 extreme, and the cutoff moves up to 128 total tokens accordingly. +# Batches above that still hand enough rows to each expert to fill the DPAS M +# tile, which is where the prefill path wins. Mirrors vLLM-xpu-kernels' `w4a16` +# dispatch, which likewise keeps the GEMV for the low tokens-per-expert regime. +_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS = 128 _MOE_VALID_PHASES = ("auto", "decode", "prefill") @@ -2482,7 +2485,7 @@ def moe( Use when the model knows it is in the prefill phase. decode_threshold: Total-token threshold for ``"auto"`` mode. If not provided, uses ``ARK_MOE_AUTO_DECODE_MAX_TOKENS`` when set to a - valid positive integer, otherwise defaults to 32. Explicit + valid positive integer, otherwise defaults to 128. Explicit argument values take precedence over the environment variable. Returns: 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 8655e53a63..fce2ebcb14 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -471,6 +471,11 @@ int4-sym decode is now at target, and the same two levers that got it there apply to FP8: get the dequant off the byte-typed datapath, and stop paying setup cost per decode call. On top of that, the FP8 MoE dispatch from vllm-xpu-kernels is mirrored into a decode-specialised entry point. +Both levers have since landed and **FP8 decode is at target too** — the +word-native dequant, the K-split lane mapping with its N-blocking, and the +removal of the per-call routing sync. That is what moved the unified +`ark.moe(phase="auto")` cutoff from 32 to 128 total tokens (see +*Auto-dispatch cutoff* below). **Word-native FP8 decode (`ARK_FP8_DECODE_MODE`, default `word`).** The decode GEMV does roughly one multiply-add per weight byte, so the dequant @@ -553,7 +558,8 @@ requires a power-of-two `group_size ≥ 16` (every shipped FP8 config — 32 / every lane of the sub-group owns at least one chunk); anything else keeps the legacy GEMV, which handles arbitrary group sizes. All three `ARK_FP8_DECODE_MODE` decoders run under both mappings, so the mode A/B -stays apples-to-apples. **Status: NEEDS-HARDWARE-VALIDATION.** +stays apples-to-apples. **Status: hardware-validated** — this mapping is +what put FP8 decode at target. **N-blocking inside the K-split kernel (`ARK_MOE_DECODE_FP8_KSPLIT_NCOLS`, default 2).** With one output column per sub-group the hot loop issues, per @@ -582,7 +588,7 @@ scale fold and the final `reduce_over_group` are untouched, so the arithmetic per output element is unchanged and `NCOLS=1` reproduces the previous kernel exactly. `test_perf_fp8_ksplit_ncols_sweep` prints all three factors per shape so the default can be set from measured data. -**Status: NEEDS-HARDWARE-VALIDATION.** +**Status: hardware-validated at the shipped default (`NCOLS=2`).** **Routing-table validation (`ARK_MOE_VALIDATE_ROUTING`, default OFF).** The Python entry point used to check `sum(num_tokens_per_expert) == total_tokens` @@ -645,10 +651,26 @@ gate, which is what the parity and A/B perf tests set. Shapes that fail the per-group shape gate (`N%64==0`, `K%32==0`, `K%group_size==0`, `group_size ∈ {32,64,128,256}`) always fall back to the scalar GEMV. +**Auto-dispatch cutoff (`ARK_MOE_AUTO_DECODE_MAX_TOKENS`, default 128).** +`ark.moe(phase="auto")` routes to `moe_gemm_decode` when +`activations.shape[0] <= cutoff` and to `moe_gemm_prefill` otherwise. The +cutoff was 32 while the decode GEMV was still the bottleneck: only the tiny +single-/few-stream case was worth keeping off the prefill grouped GEMM. Now +that the FP8 decode GEMV is at target (as int4-sym already was) the GEMV +stays ahead across the whole small-batch range rather than just at the bs1 +extreme, so the cutoff is 128 total tokens; above that each expert receives +enough rows to fill the DPAS M tile, which is where the grouped GEMM wins. +The `decode_threshold=` keyword overrides it per call and takes precedence +over the env var, and `phase="decode"` / `phase="prefill"` bypass the +heuristic entirely. Dispatch parity is covered by +`test_moe_unified.py::TestMoeUnifiedDispatch`, which pins both the cutoff +boundary (128 tokens still decode) and the overrides. + | Env var | Default | Effect | | ------- | ------- | ------ | | `ARK_FP8_DECODE_MODE` | `word` | FP8 decode implementation for the scalar GEMV: `word` (bit-field move + folded scale bias), `lut` (128-entry magnitude table), `bits` (inline bit manipulation). | | `ARK_FP8_DECODE_USE_LUT` | unset | Legacy selector, still honoured when set explicitly and when `ARK_FP8_DECODE_MODE` is unset/unrecognised: truthy → `lut`, falsy → `bits`. Also still drives the mixed-input prefill path. | +| `ARK_MOE_AUTO_DECODE_MAX_TOKENS` | `128` | Total-token cutoff used by `ark.moe(phase="auto")`: at or below it the call goes to `moe_gemm_decode`, above it to `moe_gemm_prefill`. Non-positive/unparsable values fall back to the default; the `decode_threshold=` keyword wins over both. | | `ARK_MOE_DECODE_DPAS_FP8` | ON | Route FP8 decode to the per-group DPAS grouped GEMM when the shape and occupancy gates pass; `0` forces the scalar GEMV. | | `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | Minimum tokens per expert before the DPAS path is taken; `0` disables the gate (what the parity/A-B tests set). | | `ARK_MOE_DECODE_FP8_KSPLIT` | ON | Scalar-GEMV lane mapping: one sub-group per output element with the lanes splitting K (coalesced weight loads, 16× the threads); `0` forces the legacy one-work-item-per-output-element GEMV. Shapes outside the gate (power-of-two `group_size ≥ 16`, `N%16==0`, `K%group_size==0`, `K ≥ 256`) always use the legacy mapping. | 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 c33590c6a5..d0633117a7 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 @@ -371,7 +371,10 @@ token 数(16–128)都远低于该门控(8 × 192 个专家 == 1536 个 token) int4-sym decode 的性能已经达标,把它推到达标的两个手段同样适用于 FP8: 让 dequant 离开按字节的数据通路,以及不要在每次 decode 调用里重复付出 启动开销。在此之上,还把 vllm-xpu-kernels 的 FP8 MoE dispatch 镜像成一个 -decode 专用入口。 +decode 专用入口。这两个手段现已全部落地,**FP8 decode 的性能同样达标** —— +word-native dequant、带 N 分块的 K-split lane 映射,以及去掉每次调用的路由表 +同步。正因如此,统一入口 `ark.moe(phase="auto")` 的分发阈值才从 32 提高到 +128 个 token(见下文*自动分发阈值*)。 **Word-native FP8 解码 (`ARK_FP8_DECODE_MODE`, 默认 `word`)。** decode GEMV 每读一个权重字节大约只做一次乘加,所以 dequant *就是* kernel 本身。 @@ -442,7 +445,7 @@ scratch 显存。FP8 权重每元素一个字节、本来就是 K 连续的,所 `N%16==0`、`K%group_size==0` 以及 `K ≥ 256`(保证 sub-group 的每个 lane 至少 分到一个 chunk);其余情况继续走老的 GEMV,它支持任意 group size。三种 `ARK_FP8_DECODE_MODE` 解码器在两种映射下都能运行,所以 decode mode 的 A/B 依然是同口径对比。 -**状态:NEEDS-HARDWARE-VALIDATION。** +**状态:已通过硬件验证** —— 正是这个映射把 FP8 decode 推到达标。 **K-split kernel 内的 N 分块(`ARK_MOE_DECODE_FP8_KSPLIT_NCOLS`,默认 2)。** 当一个 sub-group 只负责一个输出列时,热循环中每读一个 16 字节权重 chunk, @@ -468,7 +471,7 @@ DRAM 峰值带宽的流式 GEMV,后者才是真正的瓶颈。代价是活跃的 单个输出元素的算术完全不变,`NCOLS=1` 与改动前的 kernel 完全一致。 `test_perf_fp8_ksplit_ncols_sweep` 会逐形状打印全部三个因子的耗时,便于用 实测数据确定默认值。 -**状态:NEEDS-HARDWARE-VALIDATION。** +**状态:已在发布默认值(`NCOLS=2`)下通过硬件验证。** **路由表校验(`ARK_MOE_VALIDATE_ROUTING`,默认 OFF)。** Python 入口原先 在每次调用时都会检查 `sum(num_tokens_per_expert) == total_tokens`。当路由表 @@ -524,10 +527,23 @@ decode dispatch 改用每队列常驻的一个 slot(`get_persistent_atomic_buffe (`N%64==0`、`K%32==0`、`K%group_size==0`、 `group_size ∈ {32,64,128,256}`)的形状始终回退到 scalar GEMV。 +**自动分发阈值(`ARK_MOE_AUTO_DECODE_MAX_TOKENS`,默认 128)。** +`ark.moe(phase="auto")` 在 `activations.shape[0] <= 阈值` 时分发到 +`moe_gemm_decode`,否则分发到 `moe_gemm_prefill`。该阈值原先是 32 —— 那时 +decode GEMV 仍是瓶颈,只有极小的单流/少流场景才值得不走 prefill grouped +GEMM。如今 FP8 decode GEMV 也已达标(int4-sym 此前就已达标),GEMV 在整个 +小 batch 区间都保持领先,而不再只是在 bs1 这一极端上占优,因此阈值提高到 +128 个 token;超过之后每个专家分到的行数足以填满 DPAS 的 M tile,那正是 +grouped GEMM 占优的区间。`decode_threshold=` 关键字可按调用覆盖该阈值, +优先级高于环境变量;`phase="decode"` / `phase="prefill"` 则完全跳过该启发式。 +分发行为由 `test_moe_unified.py::TestMoeUnifiedDispatch` 覆盖,其中同时锁定了 +阈值边界(128 个 token 仍走 decode)与两种覆盖方式。 + | Env 变量 | 默认值 | 作用 | | -------- | ------ | ---- | | `ARK_FP8_DECODE_MODE` | `word` | scalar GEMV 的 FP8 解码实现:`word`(位域搬移 + 折叠 scale bias)、`lut`(128 项幅值表)、`bits`(内联位运算)。 | | `ARK_FP8_DECODE_USE_LUT` | 未设置 | 旧的选择开关;当它被显式设置、且 `ARK_FP8_DECODE_MODE` 未设置或取值无法识别时仍然生效:truthy → `lut`,falsy → `bits`。它同时仍然驱动 mixed-input prefill 路径。 | +| `ARK_MOE_AUTO_DECODE_MAX_TOKENS` | `128` | `ark.moe(phase="auto")` 使用的总 token 阈值:小于等于它走 `moe_gemm_decode`,大于它走 `moe_gemm_prefill`。非正数或无法解析的取值会回退到默认值;`decode_threshold=` 关键字优先级高于两者。 | | `ARK_MOE_DECODE_DPAS_FP8` | ON | 形状与占用率门控都通过时,把 FP8 decode 路由到 per-group DPAS grouped GEMM;`0` 强制走 scalar GEMV。 | | `ARK_MOE_DECODE_DPAS_FP8_MIN_TPE` | `8` | 走 DPAS 路径所需的最小每专家 token 数;`0` 关闭门控(对齐/A-B 用例所设)。 | | `ARK_MOE_DECODE_FP8_KSPLIT` | ON | scalar GEMV 的 lane 映射:一个 sub-group 负责一个输出元素、由 lane 切分 K(访存合并,线程数 ×16);`0` 强制走老的「一个 work-item 一个输出元素」GEMV。未通过门控(`group_size` 为 ≥ 16 的 2 的幂、`N%16==0`、`K%group_size==0`、`K ≥ 256`)的形状始终使用老映射。 | diff --git a/auto_round_extension/ark/test/test_moe_unified.py b/auto_round_extension/ark/test/test_moe_unified.py index 2d5eb232fe..ad29f11589 100644 --- a/auto_round_extension/ark/test/test_moe_unified.py +++ b/auto_round_extension/ark/test/test_moe_unified.py @@ -93,6 +93,9 @@ def _unified_skip_reason() -> str: # --------------------------------------------------------------------------- _AUTO_DECODE_SHAPE = dict(num_experts=4, tokens_per_expert=[4, 4, 4, 4], N=128, K=256) # total_tokens=16 +# Sits exactly on the default cutoff (128 total tokens), so it pins the +# boundary of `_MOE_AUTO_DECODE_MAX_TOTAL_TOKENS` (dispatch is `<=`). +_AUTO_DECODE_BOUNDARY_SHAPE = dict(num_experts=4, tokens_per_expert=[32, 32, 32, 32], N=128, K=256) # total=128 _AUTO_PREFILL_SHAPE = dict(num_experts=4, tokens_per_expert=[80, 80, 80, 80], N=128, K=256) # total_tokens=320 @@ -193,10 +196,43 @@ def test_auto_picks_decode_for_small_total_tokens(self): group_size=group_size, asym=False, ) - # total tokens = 16 (<= default threshold 32) -> dispatched to decode + # total tokens = 16 (<= default threshold 128) -> dispatched to decode # -> output must be bit-identical to moe_gemm_decode. torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) + def test_auto_picks_decode_at_default_threshold(self): + shape = _AUTO_DECODE_BOUNDARY_SHAPE + total_tokens = sum(shape["tokens_per_expert"]) + E, N, K = shape["num_experts"], shape["N"], shape["K"] + group_size = 128 + dtype = torch.float16 + + activations, packed, scales, _ = _make_int4_sym(E, N, K, group_size, dtype, total_tokens) + ntpe = torch.tensor(shape["tokens_per_expert"], dtype=torch.int32, device="xpu") + + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + ) + out_decode = ark.moe_gemm_decode( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + ) + # total tokens = 128 == default threshold, and the dispatch is `<=`, + # so this still routes to decode. + torch.testing.assert_close(out_auto, out_decode, rtol=0, atol=0) + def test_auto_picks_prefill_for_large_total_tokens(self): shape = _AUTO_PREFILL_SHAPE total_tokens = sum(shape["tokens_per_expert"])