diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index b66e180d8..f9f147583 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. @@ -1670,6 +1673,59 @@ 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 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, @@ -1688,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") @@ -1801,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) @@ -1859,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) @@ -2094,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 @@ -2331,8 +2387,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 +2397,52 @@ 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`. +# +# 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") +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 +2454,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. @@ -2382,23 +2470,23 @@ 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. - * ``"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 128. Explicit + argument values take precedence over the environment variable. Returns: ``[total_tokens, N]`` in the activations dtype. Bit-identical to the @@ -2408,14 +2496,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/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index d9c9c7c03..dc701e0da 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.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp index cff5626de..e5a2b59df 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp @@ -6,8 +6,10 @@ #pragma once +#include #include #include +#include #include #include @@ -53,27 +55,32 @@ namespace moe_detail { using namespace cute; using namespace MoE; -// Helper to choose TiledMMA based on element types -template -auto choose_tiled_mma(TA* A, TB* B) { +// Helper to choose TiledMMA for a given work-group tile / sub-group layout. +// +// The MMA atom (``XE_DPAS_TT<8, float, ...>``) is fixed; only the work-group +// tile (``WGTile``) and the sub-group tiling (``SGLayout``) vary between tile +// policies. Because every bf16/fp16 policy below keeps the same number of +// sub-group rows in M (8) the per-sub-group tile stays 32x64x32, so the same +// 2D block copy atoms remain valid across all of them. +template +auto choose_tiled_mma() { using TA_non_CV = cutlass::platform::remove_cv_t; using TB_non_CV = cutlass::platform::remove_cv_t; auto op = XE_DPAS_TT<8, float, TA_non_CV, TB_non_CV>{}; - using WGTile = Shape<_256, _128, _32>; // 256x128 WG tile size - using SGLayout = Layout, Stride<_2, _1, _0>>; // 8x2 SG tiling, n-major - using MMA = typename TiledMMAHelper, Layout, SGLayout>::TiledMMA; return MMA{}; } -// Unique kernel name tag -template +// Unique kernel name tag. The tile policy (WGTile / SGLayout) is part of the +// tag so each policy specialization produces a distinct SYCL kernel name. +template class MoEGemmKernel; // MOE GEMM launcher template -template +template void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const ElementB* weights, const ElementS* scales, ElementD* outputs, const int gemm_n, const int gemm_k, int* num_rows_per_expert_device, const int num_experts) { @@ -86,7 +93,7 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen auto dummy_group_problem_shape = cutlass::gemm::GroupProblemShape>{1, &dummy_problem_shape, nullptr}; - using TileShape = Shape<_256, _128, _32>; + using TileShape = WGTile; using ClusterShape = Shape<_1, _1, _1>; auto scheduler_params = PersistentTileSchedulerXeMoE::to_underlying_arguments( @@ -97,7 +104,7 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen scheduler_params, dummy_group_problem_shape, TileShape{}, ClusterShape{}, hw_info, PersistentTileSchedulerXeMoE::Arguments{1, RasterOrderOptions::AlongN}); - auto mma = choose_tiled_mma(activations, weights); + auto mma = choose_tiled_mma(); auto MaxThreadsPerWorkgroup = size(mma); dim3 local_range{static_cast(MaxThreadsPerWorkgroup), 1, 1}; @@ -110,7 +117,7 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen syclex::properties kernel_props{syclex::sub_group_size<16>, intelex::grf_size<256>}; - auto event = q->parallel_for>( + auto event = q->parallel_for>( sycl::nd_range<3>(global, local), kernel_props, [=](auto) { MoE::MoEGEMM, XE_LOAD_2D_VNNI<16, 32, 16, 16>, XE_STORE_2D<16, 8, 32>, 'R', 'R', 'R'>(activations, weights, scales, outputs, mma, num_rows_per_expert_device, num_experts, gemm_n, @@ -121,6 +128,65 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen event.wait(); } +// Whether the N-based tile-policy heuristic is enabled (default on). +// +// Set ``ARK_MOE_GEMM_FIXED_TILE`` to a truthy value ("1"/"true"/"on"/"yes") +// to always use the historical fixed 256x128 (8x2) tile regardless of N. +// This provides an escape hatch should a specific device regress with the +// wider tiles. +inline bool moe_gemm_fixed_tile() { + const char* env = std::getenv("ARK_MOE_GEMM_FIXED_TILE"); + if (env == nullptr) { + return false; + } + std::string v(env); + for (auto& c : v) { + c = static_cast(std::tolower(static_cast(c))); + } + return !(v == "0" || v == "false" || v == "off" || v == "no" || v.empty()); +} + +// Select the work-group tile policy from the output width ``N`` and dispatch, +// mirroring the ``w16a16`` large-M heuristic in vllm-xpu-kernels grouped GEMM: +// +// * N <= 64 -> 256x64x32, SGLayout 8x1 +// * N <= 512 -> 256x128x32, SGLayout 8x2 (historical default) +// * N > 512 -> 256x256x32, SGLayout 8x4 +// +// Prefill routes many tokens per expert (large M), so the taller/wider N tile +// increases sub-group utilization and reduces the number of work-group tiles +// launched for the large-N up/down projections. All three policies share the +// same per-sub-group tile (32x64x32), so the copy atoms in +// ``moe_gemm_launcher`` remain valid. +template +void moe_gemm_dispatch(sycl::queue* q, const Element* activations, const Element* weights, const Element* scales, + Element* outputs, const int gemm_n, const int gemm_k, int* num_rows_per_expert_device, + const int num_experts) { + using N64 = Shape<_256, _64, _32>; + using SG64 = Layout, Stride<_1, _1, _0>>; + using N128 = Shape<_256, _128, _32>; + using SG128 = Layout, Stride<_2, _1, _0>>; + using N256 = Shape<_256, _256, _32>; + using SG256 = Layout, Stride<_4, _1, _0>>; + + if (moe_gemm_fixed_tile()) { + moe_gemm_launcher<'R', 'R', N128, SG128, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + return; + } + + if (gemm_n <= 64) { + moe_gemm_launcher<'R', 'R', N64, SG64, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + } else if (gemm_n <= 512) { + moe_gemm_launcher<'R', 'R', N128, SG128, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + } else { + moe_gemm_launcher<'R', 'R', N256, SG256, Element, Element, Element, Element>( + q, activations, weights, scales, outputs, gemm_n, gemm_k, num_rows_per_expert_device, num_experts); + } +} + } // namespace moe_detail // Public MOE GEMM API @@ -129,7 +195,7 @@ inline void moe_gemm(sycl::queue* q, void* activations, void* weights, void* sca switch (dtype) { case BTLA_DTYPE::BF16: { using Element = cutlass::bfloat16_t; - moe_detail::moe_gemm_launcher<'R', 'R', Element, Element, Element, Element>( + moe_detail::moe_gemm_dispatch( q, static_cast(activations), static_cast(weights), static_cast(scales), static_cast(outputs), N, K, num_tokens_per_expert, num_experts); @@ -137,7 +203,7 @@ inline void moe_gemm(sycl::queue* q, void* activations, void* weights, void* sca } case BTLA_DTYPE::F16: { using Element = cutlass::half_t; - moe_detail::moe_gemm_launcher<'R', 'R', Element, Element, Element, Element>( + moe_detail::moe_gemm_dispatch( q, static_cast(activations), static_cast(weights), static_cast(scales), static_cast(outputs), N, K, num_tokens_per_expert, num_experts); diff --git a/auto_round_extension/ark/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 1b8a4feb0..65d1089d7 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,13 @@ // 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, 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]) @@ -28,17 +34,34 @@ // 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 #pragma once +#include #include #include +#include +#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 @@ -47,20 +70,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) @@ -71,6 +99,41 @@ 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 +// 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) // ---------------------------------------------------------------------------- @@ -80,15 +143,27 @@ class MoEDecodeKernelFP; template class MoEDecodeKernelInt4; +template +class MoEDecodeKernelInt4Coalesced; + +template +class MoEDecodeRepackInt4; + +template +class MoEDecodeActGroupSum; + template class MoEDecodeKernelInt8; template class MoEDecodeKernelInt2; -template +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 @@ -96,15 +171,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. @@ -194,6 +275,171 @@ 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 (asym int4 only). +// +// 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. +// +// 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. +// ---------------------------------------------------------------------------- +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. // @@ -204,6 +450,84 @@ 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 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: +// +// asym: (word >> 4j) & 0xF +// sym : (int)(word << (28 - 4j)) >> 28 +// +// 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 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 +// 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) { + 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. 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 WordVec = sycl::vec; +#pragma unroll + for (int s = 0; s < CHUNK / SUB; ++s) { + const ActVec av = *reinterpret_cast(act_ptr + s * SUB); + const WordVec wv = *reinterpret_cast(w_ptr + s * (SUB / 2)); +#pragma unroll + 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]); + } + } + } +} + 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, @@ -223,6 +547,18 @@ 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. 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)}; @@ -241,65 +577,64 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig 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]); - float zero = 0.0f; - if constexpr (Asym) { - 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. + // + // 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). 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; - 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, acc_q0, + acc_q1); } - // Scalar tail for group_size not divisible by CHUNK. + 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); + } + // 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; - 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 += static_cast(act_row[k_base + kk]) * w0; - acc += static_cast(act_row[k_base + kk + 1]) * 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) { + acc += scale * ((acc_q0 + acc_q1) - static_cast(z_row[g]) * a_sum_row[g]); + } else { + acc += scale * (acc_q0 + acc_q1); } } @@ -307,6 +642,309 @@ void launch_int4(sycl::queue* q, const ScalarT* activations, const uint8_t* weig }); } +// ---------------------------------------------------------------------------- +// 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() { + return env_flag_enabled("ARK_MOE_DECODE_INT4_REPACK_CACHE", false); // default OFF -- see comment above +} + +// ---------------------------------------------------------------------------- +// 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, 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 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 +// 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 +// 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 independent of routing. +// ---------------------------------------------------------------------------- +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) + // 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_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 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. 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)}; + 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, 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; + [[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 + // 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]; + [[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; + if constexpr (Asym) { + 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]); + const int k_base = g * group_size; + // Per-token split accumulators for the raw integer-weighted dot + // 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) { + acc_q0[m] = 0.0f; + acc_q1[m] = 0.0f; + } + + // 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) { + 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 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; + 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(kb_base + kb, k_base + 2 * kb); + } + // 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 + const uint32_t word = *reinterpret_cast( + w_tile + static_cast(b_abs / PACK_VEC) * chunk_stride + + static_cast(lane) * PACK_VEC); + int qv[8]; + decode_int4_octet(word, qv); +#pragma unroll + for (int u = 0; u < PACK_VEC; ++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(kb_base + kb, k_base + 2 * kb); + } + + 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]); + } + } + } + + 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]); + } + } + }); +} + // ---------------------------------------------------------------------------- // INT8 (S8) GEMV with group-wise dequantization. // @@ -317,6 +955,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, @@ -365,39 +1039,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; } @@ -416,6 +1083,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, @@ -469,47 +1184,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) { @@ -542,12 +1233,132 @@ 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. +// +// 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. // ---------------------------------------------------------------------------- -template + +// 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. +// +// 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& acc0, + float& acc1) { + static_assert(sizeof(ScalarT) == sizeof(uint16_t), "ScalarT must be a 16-bit floating type"); + static_assert(CHUNK % FP8_SUB == 0, "CHUNK must be a multiple of 16"); +#pragma unroll + 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); + } +} + +// 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) { @@ -560,11 +1371,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)); @@ -582,45 +1397,502 @@ 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: 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. + // 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). + // Two partial accumulators break the fp32 dependency chain. + float group_acc0 = 0.0f; + float group_acc1 = 0.0f; 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, + 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_acc0, group_acc1); } 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 = fp8_decode_scalar(w_row[k_base + kk]); + group_acc0 += static_cast(act_row[k_base + kk]) * w; } + acc += (group_acc0 + group_acc1) * scale; } outputs[static_cast(token) * N + n_global] = static_cast(acc); }); } +// ---------------------------------------------------------------------------- +// 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. +// +// 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 +// 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 `NCOLS` output columns, so a work-group +// covers `N_TILE * NCOLS` consecutive columns. +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 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. +// ---------------------------------------------------------------------------- +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. +// `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; + return true; +} + +// ---------------------------------------------------------------------------- +// 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) { + 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; + + 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, 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>( + 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_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 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, 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. All + // 2 * NCOLS weight loads are issued before the first is consumed. + for (; k0 + KSPLIT_STEP + KSPLIT_CH <= K; k0 += 2 * KSPLIT_STEP) { + 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) { + 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); + } + } + +#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 +// 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_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); + } +} + +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_dispatch( + q, activations, weights, scales, outputs, expert_id_per_token, total_tokens, N, K, group_size, ksplit); + return; + case Fp8DecodeMode::kBits: + 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_dispatch( + q, activations, weights, scales, outputs, expert_id_per_token, total_tokens, N, K, group_size, ksplit); + return; + } +} + } // namespace moe_decode_detail // ---------------------------------------------------------------------------- -// Public API +// 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 +// grouped GEMM (`moe_dpas_s4::moe_decode_s4_dpas_per_group_dispatch`) instead +// 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 +// 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() { + return moe_decode_detail::env_flag_enabled("ARK_MOE_DECODE_DPAS_S4", true); // default ON +} + +// ---------------------------------------------------------------------------- +// 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_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 +// 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. +// Re-read on every call so tests / benchmarks can toggle it in-process. +// ---------------------------------------------------------------------------- +inline bool moe_decode_coalesce_int4_enabled() { + return moe_decode_detail::env_flag_enabled("ARK_MOE_DECODE_COALESCE_INT4", true); // default ON +} + +// ---------------------------------------------------------------------------- +// 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) * moe_decode_detail::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: // BTLA_DTYPE::F16 / BF16 : weights stored as [E, N, K] in matching @@ -630,19 +1902,46 @@ void launch_fp8(sycl::queue* q, const ScalarT* activations, const uint8_t* weigh // 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 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), -// 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) { - 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 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, 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); + 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); + } if (weight_dtype == BTLA_DTYPE::F16 || weight_dtype == BTLA_DTYPE::BF16) { if (weight_dtype != act_dtype) { @@ -663,33 +1962,94 @@ 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, 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) { + if (act_dtype == BTLA_DTYPE::F16) { + 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_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); + } + return; + } + // 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() && + moe_decode_coalesce_int4_amortized(total_tokens, num_experts); 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 if (act_dtype == BTLA_DTYPE::BF16) { + } 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); + } } - } else { - throw std::invalid_argument("moe_gemm_decode(int4): act_dtype must be FP16 or BF16"); } return; } @@ -762,63 +2122,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 d142621fd..7ba9c8fc6 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. // @@ -211,17 +308,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; + } } } @@ -273,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 = []() { @@ -286,6 +401,47 @@ 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. 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) { + 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) { + 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 1638bf7af..a7bf88d2e 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 @@ -193,6 +229,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`. @@ -1002,7 +1067,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); @@ -1012,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_fp8_native.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_native.hpp index e3bed66e0..db8835aa1 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); } 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 2e05dfa44..8f7260bd9 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 fad8d87eb..0a1c6a923 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 @@ -129,11 +131,30 @@ 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; +// --------------------------------------------------------------------------- +// Persistent per-queue atomic work-group counter. +// +// 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). +// --------------------------------------------------------------------------- +using ::ark::moe_dpas_fp8::get_persistent_atomic_buffer; + // --------------------------------------------------------------------------- // Variant B -- per-K-group S4 (sym) mainloop. // @@ -548,8 +569,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; @@ -637,11 +660,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>( \ @@ -649,16 +670,105 @@ 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 +} + +// --------------------------------------------------------------------------- +// Decode-specialized S4 (sym) per-K-group DPAS dispatch. +// +// 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 (which uses the identical ladder). +// +// `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_force_m8() { + const char* env = std::getenv("ARK_MOE_DECODE_S4_DPAS_M8"); + 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 == "1" || s == "true" || s == "on" || s == "yes") return true; + return false; +} - sycl::free(atomic_buffer, *q); +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) { + // 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); + 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 = get_persistent_atomic_buffer(q); + + // 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, + num_tokens_per_expert, E, group_size, atomic_buffer); } // --------------------------------------------------------------------------- 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 ff97deaa1..fce2ebcb1 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,173 @@ 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. + +**S4 DPAS decode path** — the *decode* phase (`sycl_tla_moe_decode.hpp`, +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). + +**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`) 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. + +**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. 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. 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 +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` + @@ -298,6 +465,232 @@ 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. +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 +*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`. + +**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: 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 +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: 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` +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]` +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. + +**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. | +| `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`), `::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), +`::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 `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 732cfec70..d0633117a 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,12 +222,346 @@ 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 浪费。 + +**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。它与 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=1` 会强制使用旧的钉死 8 行 tile 以便 A/B +对比(数值完全相同,仅 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`)。`ARK_MOE_DECODE_DPAS_S4_MIN_TPE` 可覆盖该 +"每专家 token 数" 阈值;设为 `0` 则关闭门控(只要形状门控通过就走 DPAS), +精度测试与 DPAS/标量 对比性能测试即使用该设置。 + +**基于 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 仍然发出的是*字节*加载。现在布局改为 +`[E, N/16, ceil(K/8), 16, 4]`：一个 chunk 为 +tile 内 16 列中的每一列存放 4 个连续的打包字节，按 lane 主序排列，因此 lane `l` +在 chunk 偏移 `l*4` 处读取自己的 4 个字节，sub-group 整体仍然覆盖 64 个连续字节。 +一个 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]` 权重约定保持不变。 + +**提取激活求和(仅 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。 +求和顺序的变化仅带来几个 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 +每生成一个 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` + `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 专用入口。这两个手段现已全部落地,**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 本身。 +两条旧解码路径每个字节都要付出真实开销:`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`。 + +**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 依然是同口径对比。 +**状态:已通过硬件验证** —— 正是这个映射把 FP8 decode 推到达标。 + +**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` 会逐形状打印全部三个因子的耗时,便于用 +实测数据确定默认值。 +**状态:已在发布默认值(`NCOLS=2`)下通过硬件验证。** + +**路由表校验(`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]` +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。 + +**自动分发阈值(`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`)的形状始终使用老映射。 | +| `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_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` 的回退)、 +`::test_decode_fp8_ksplit_ncols_match`(各分块因子结果一致,并覆盖 `N` +无法整除时的回退)与 +`::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 8bad20cf9..fd6d04c34 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -480,6 +480,69 @@ 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}) 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) + 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") + monkeypatch.setenv("ARK_MOE_DECODE_DPAS_S4_MIN_TPE", "0") + 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() + + 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) + # 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): num_experts = 4 @@ -511,6 +574,258 @@ 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), 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] + 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, + # 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() + + 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) + + @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, 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() + + 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) + + @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 @@ -557,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): @@ -710,6 +1049,310 @@ 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) + + @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) + + @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 cf9b9a74c..ad348e3b7 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,41 @@ 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] + +# 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 + (3072, 1536), # down-proj +] + DECODE_SHAPES = [ # (label, num_experts, tokens_per_expert, N, K) # batch=1 decode (single-stream). @@ -369,6 +404,181 @@ 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) + 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. + ``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( + 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") + 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]) + 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``. + + 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, 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). + + 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)" + ) + 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 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") + 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]) @pytest.mark.parametrize("asym", [False, True]) def test_perf_int8(self, dtype, asym): @@ -483,6 +693,235 @@ def test_perf_fp8(self, dtype, fp8_dtype): ) _print_row(label, N, K, total_tokens, base_ms, ark_ms) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) + 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_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]) + 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): + """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) + @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): diff --git a/auto_round_extension/ark/test/test_moe_unified.py b/auto_round_extension/ark/test/test_moe_unified.py index 4e769456f..ad29f1158 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,15 @@ 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=[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 def _make_int4_sym(E, N, K, group_size, dtype, total_tokens): @@ -163,8 +167,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 +196,45 @@ 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 = 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_prefill_for_large_tokens_per_expert(self): - shape = _PREFILL_SHAPE + 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"]) E, N, K = shape["num_experts"], shape["N"], shape["K"] group_size = 128 @@ -228,9 +265,67 @@ 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") + + # A threshold strictly below the shape's total tokens forces prefill. + out_auto = ark.moe( + activations, + packed, + ntpe, + scales=scales, + weight_bits=4, + group_size=group_size, + asym=False, + phase="auto", + decode_threshold=decode_total_tokens - 1, + ) + 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 +334,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 +344,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 +356,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 +427,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 +453,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 +482,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 +504,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 +526,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