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..5e28a0ecc 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 @@ -79,7 +79,18 @@ void moe_gemm_launcher(sycl::queue* q, const ElementA* activations, const Elemen const int num_experts) { compat::set_default_queue(*q); - int sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + // Query the multiprocessor (Xe-core / EU) count from the *queue's own device* + // instead of a hardcoded ordinal 0. On multi-card systems the CUTLASS / + // syclcompat device enumeration order need not match the device the caller's + // queue runs on, so `query_device_multiprocessor_count(0)` may report a + // different device's count. That mis-sizes the persistent-scheduler grid for + // the device the kernel actually executes on, leading to incorrect tile + // coverage and wrong results that only manifest when more than one device is + // visible. Deriving the count from `*q` keeps the launch consistent with the + // tensor's device. (`query_device_multiprocessor_count` is itself just + // `get_device(id).get_info()`, so this is numerically + // identical for the correct device.) + int sm_count = q->get_device().get_info(); cutlass::KernelHardwareInfo hw_info{0, sm_count}; auto dummy_problem_shape = cute::Shape{1, gemm_k, gemm_n}; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp index e6a3ececc..2105ede71 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_mixed.hpp @@ -876,16 +876,23 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v // reads packed `[E, N, K/2]` uint8_t nibbles directly and folds the // upcast into the DPAS mainloop via CuTe's `reorder(tBrB, tCrB)`, so // the B-side global traffic is halved vs. the S4->S8 upcast branch - // below. Opt-in default via `ARK_MOE_PREFILL_DPAS_S4` (default ON); + // below. Opt-in via `ARK_MOE_PREFILL_DPAS_S4=1` (default OFF); // silent fallback to the S4->S8 upcast branch (which is itself gated // by `ARK_MOE_PREFILL_DPAS_INT8`) or to the generic dequant path if // the shape gate rejects the tile geometry. // + // The single-pass mainloop decodes each packed `int4b_t` fragment into + // an `int8_t` staging fragment in registers and reuses the validated + // `int8_t -> ElementA` reorder (see `xe_gemm_s4_pergroup`), so it no + // longer routes through the interleaved + // `NumericArrayConverter` that previously + // miscomputed a fraction of outputs. int4-sym is routed here only when + // `ARK_MOE_PREFILL_DPAS_S4=1` is set (default OFF); otherwise it falls + // through to the S4->S8 upcast + INT8 DPAS branch below. + // // STATUS: NEEDS-HARDWARE-VALIDATION. See // `sycl_tla_moe_prefill_s4_dpas.hpp` for the port's provenance & the - // on-hardware TODOs (chief among them: `NumericArrayConverter - // ` availability in the pinned - // cutlass-sycl). + // remaining on-hardware TODOs. if (weight_dtype == BTLA_DTYPE::S4_CLIP && !asym && moe_dpas_s4::moe_prefill_dpas_s4_enabled() && moe_dpas_s4::moe_prefill_dpas_s4_pergroup_shape_ok(N, K, group_size) && @@ -917,15 +924,23 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v // and the DPAS mainloop then folds the per-K-group scale exactly the // same way as the S8-sym path -- reusing the packed scale tensor // unmodified. Silent fallback to the generic dequant path if the shape - // predicate rejects the tile geometry, if `asym=true`, or if the caller - // opted out via `ARK_MOE_PREFILL_DPAS_INT8=0`. + // predicate rejects the tile geometry, if `asym=true`, if the caller + // opted out via `ARK_MOE_PREFILL_DPAS_INT8=0`, or (the default) unless + // the caller opts in via `ARK_MOE_PREFILL_DPAS_LOWBIT=1`. // // For S4-sym specifically this branch is the *fallback* for the // single-pass S4 DPAS path above -- callers who disable - // `ARK_MOE_PREFILL_DPAS_S4` land here instead of on the generic - // dequant path, so the two-pass INT4->INT8 pipeline stays available - // as a runtime kill-switch until the single-pass mainloop is - // hardware-validated. + // `ARK_MOE_PREFILL_DPAS_S4` land here only when they *also* opt into + // `ARK_MOE_PREFILL_DPAS_LOWBIT=1`; otherwise int4-sym / int2-sym fall + // through to the generic bit-exact dequant path below. The two-pass + // INT4/INT2->INT8 pipeline stays available as a runtime opt-in until the + // low-bit DPAS numerics are hardware-validated. + // + // STATUS: default OFF (`ARK_MOE_PREFILL_DPAS_LOWBIT`). The upcast + + // INT8 DPAS pipeline still miscomputes a fraction of outputs on + // production-scale prefill shapes (observed: max abs diff ~70 on the + // `medium E=8`, K=14336 int4-sym + fp16 accuracy case), so it is not + // taken by default. See `moe_prefill_dpas_lowbit_enabled()`. // // The dequant workspace pointer we reinterpret as `int8_t*` is the same // caller-owned buffer used by the bf16/fp16 dequant fallback: since it @@ -934,6 +949,7 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v // safe and does not require a separate allocation. if ((weight_dtype == BTLA_DTYPE::S4_CLIP || weight_dtype == BTLA_DTYPE::S2_CLIP) && !asym && moe_dpas_int::moe_prefill_dpas_int_enabled() && + moe_dpas_int::moe_prefill_dpas_lowbit_enabled() && moe_dpas_int::moe_prefill_dpas_int_pergroup_shape_ok(N, K, group_size) && dequant_workspace != nullptr && (act_dtype == BTLA_DTYPE::F16 || act_dtype == BTLA_DTYPE::BF16)) { 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..0e9308c1b 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 @@ -733,13 +733,12 @@ CUTE_DEVICE void MoEGEMM(const ElementA* Activations, const ElementB* Weights, int group_range = item.get_group_range(1); int local_id = item.get_local_linear_id(); - if (group_id == 0 && local_id == 0) { - auto atm = sycl::atomic_ref( - atomic_buffer[0]); - atm.store(0); - } + // NOTE: `atomic_buffer[0]` is zero-initialized on the host (via a queued + // `memset` that the kernel depends on) before launch. It must NOT be reset + // here: an in-kernel `store(0)` by a single work-group has no grid-level + // synchronization, so other work-groups could read an uninitialized/stale + // value or have this late store clobber an already-advanced counter, both of + // which corrupt tile scheduling. int pre_rows = 0; int pre_tiles = 0; @@ -862,8 +861,11 @@ void MoEGEMMLauncher(sycl::queue& stream, const ElementA* activations, SGLayout>::TiledMMA; auto mma = MMA{}; + // Query the EU count from the queue's own device (see sycl_tla_moe.hpp): on + // multi-card systems a hardcoded ordinal 0 can name a different device than + // the one `stream` runs on, mis-sizing the grid and corrupting results. int sm_count = - cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + stream.get_device().get_info(); auto MaxThreadsPerWorkgroup = size(mma); static constexpr int MaxThreadsPerSM = 512; @@ -886,7 +888,14 @@ void MoEGEMMLauncher(sycl::queue& stream, const ElementA* activations, using GmemTiledCopyB = typename policy::GmemTiledCopyB; using GmemTiledCopyD = typename policy::GmemTiledCopyD; + // Zero-init the persistent-scheduler tile counter on the host before launch. + // Doing it here (rather than inside the kernel) removes the cross-work-group + // race: the kernel below is made to depend on this memset event, so every + // work-group observes a fully initialized counter. + auto init_event = stream.memset(atomic_buffer, 0, sizeof(int32_t)); + auto event = stream.submit([&](sycl::handler& cgh) { + cgh.depends_on(init_event); sycl::local_accessor local_mem(sycl::range<1>(1), cgh); cgh.parallel_for>( 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..b1394b609 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 @@ -584,13 +584,12 @@ CUTE_DEVICE void MoEGEMM_int(const ElementA* Activations, int group_range = item.get_group_range(1); int local_id = item.get_local_linear_id(); - if (group_id == 0 && local_id == 0) { - auto atm = sycl::atomic_ref( - atomic_buffer[0]); - atm.store(0); - } + // NOTE: `atomic_buffer[0]` is zero-initialized on the host (via a queued + // `memset` that the kernel depends on) before launch. It must NOT be reset + // here: an in-kernel `store(0)` by a single work-group has no grid-level + // synchronization, so other work-groups could read an uninitialized/stale + // value or have this late store clobber an already-advanced counter, both of + // which corrupt tile scheduling. int pre_rows = 0; int pre_tiles = 0; @@ -713,8 +712,11 @@ void MoEGEMMLauncher_int(sycl::queue& stream, const ElementA* activations, SGLayout>::TiledMMA; auto mma = MMA{}; + // Query the EU count from the queue's own device (see sycl_tla_moe.hpp): on + // multi-card systems a hardcoded ordinal 0 can name a different device than + // the one `stream` runs on, mis-sizing the grid and corrupting results. int sm_count = - cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + stream.get_device().get_info(); auto MaxThreadsPerWorkgroup = size(mma); static constexpr int MaxThreadsPerSM = 512; @@ -737,7 +739,14 @@ void MoEGEMMLauncher_int(sycl::queue& stream, const ElementA* activations, using GmemTiledCopyB = typename policy::GmemTiledCopyB; using GmemTiledCopyD = typename policy::GmemTiledCopyD; + // Zero-init the persistent-scheduler tile counter on the host before launch. + // Doing it here (rather than inside the kernel) removes the cross-work-group + // race: the kernel below is made to depend on this memset event, so every + // work-group observes a fully initialized counter. + auto init_event = stream.memset(atomic_buffer, 0, sizeof(int32_t)); + auto event = stream.submit([&](sycl::handler& cgh) { + cgh.depends_on(init_event); sycl::local_accessor local_mem(sycl::range<1>(1), cgh); cgh.parallel_for>( @@ -884,6 +893,37 @@ inline bool moe_prefill_dpas_int_enabled() { return true; } +// --------------------------------------------------------------------------- +// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_LOWBIT` (default OFF). +// +// Gates the low-bit (S4-sym / S2-sym) -> INT8 upcast two-pass path that +// reuses this INT8 DPAS mainloop (see `sycl_tla_moe_mixed.hpp`). Decoupled +// from `ARK_MOE_PREFILL_DPAS_INT8` (which gates the *genuine* INT8-weight +// DPAS path) so the low-bit two-pass can be toggled without disturbing the +// validated INT8 path. +// +// STATUS: default OFF. The two-pass S4->S8 upcast + INT8 DPAS path still +// miscomputes a fraction of outputs on production-scale prefill shapes +// (observed: max abs diff ~70 on the `medium E=8`, K=14336 int4-sym + fp16 +// accuracy case) -- the same defect that led `ARK_MOE_PREFILL_DPAS_S4` to be +// defaulted OFF. Until the low-bit DPAS pipeline is root-caused and fixed on +// hardware, int4-sym / int2-sym prefill falls through to the bit-exact +// generic dequant path by default; the two-pass path stays available behind +// `ARK_MOE_PREFILL_DPAS_LOWBIT=1` for continued debugging. +// +// Truthy values (case-insensitive): "1", "true", "on", "yes" enable. +// Anything else (including unset) leaves the path disabled. Re-read on every +// call so benchmarks / tests can toggle the path in-process. +// --------------------------------------------------------------------------- +inline bool moe_prefill_dpas_lowbit_enabled() { + const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_LOWBIT"); + if (env == nullptr) return false; // default OFF + 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; +} + // --------------------------------------------------------------------------- // Shape preconditions for the per-K-group INT8 dispatcher branch. Matches // the FP8 per-group predicate exactly (same policy tiles, same group_size 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..8788a4369 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 @@ -1,10 +1,22 @@ // SYCL MoE Prefill — S4 (sym) mixed-input DPAS Grouped GEMM (Variant B) // (Fork of `sycl_tla_moe_prefill_int_dpas.hpp` per-group mainloop) // -// STATUS: NEEDS-HARDWARE-VALIDATION -- untested single-pass port. Precedence -// and gating in `sycl_tla_moe_mixed.hpp` fall back to the S4->S8 upcast + -// INT8 DPAS path when this header's env gate is off, so a regression can -// be neutralised at runtime without a rebuild. +// STATUS: NEEDS-HARDWARE-VALIDATION -- single-pass port. The env gate +// `ARK_MOE_PREFILL_DPAS_S4` defaults OFF (see `moe_prefill_dpas_s4_enabled()` +// below); set `ARK_MOE_PREFILL_DPAS_S4=1` to opt in. The previously observed +// gross-outlier defect on the large-M `medium E=8`, K=14336 shape (max abs diff +// ~70) was traced to the bespoke large-M policy breaking the packed-nibble +// decode's `SG_N = 16` invariant (it used `SG_N = 64`); large-M now reuses the +// validated `m_32` geometry (see the dispatch below), matching the sub-group +// fragment shape every passing S4 shape exercises. The packed-nibble mainloop decodes each `int4b_t` +// B fragment into an `int8_t` staging fragment in registers and then reuses +// the SAME validated `int8_t -> ElementA` `reorder(...)` the INT8 per-group +// path uses (see `xe_gemm_s4_pergroup`). This avoids the interleaved +// `NumericArrayConverter` fast path, whose +// pre-shuffled-B-layout assumption previously miscomputed a fraction of +// outputs on production-scale prefill shapes. Precedence and gating in +// `sycl_tla_moe_mixed.hpp` fall back to the S4->S8 upcast + INT8 DPAS +// path whenever the gate is not explicitly enabled (the default). // --------------------------------------------------------------------------- // Design rationale // ---------------- @@ -19,13 +31,23 @@ // // This header removes the round-trip: the mainloop reads packed // `[E, N, K/2]` `uint8_t` (two nibbles per byte, sym-signed [-8, 7]) and -// upcasts to the activation dtype in registers via the same -// `cute::reorder(tBrB, tCrB)` machinery the INT8 header uses -- the only -// substantive difference is that CuTe/cutlass-sycl's -// `NumericArrayConverter` unpacks 4-bit -// fields two-per-byte from the loaded fragment. The B-side global load -// is halved (bytes, not int8 elements) so the mainloop is bandwidth- -// bound on `E * N * K / 2` instead of `E * N * K`. +// upcasts to the activation dtype in registers. Each `int4b_t` copy +// fragment is decoded into an `int8_t` staging fragment (sign-extended +// [-8, 7], reproducing `decode_int4_pair`) and then handed to the SAME +// `cute::reorder(tBrB_i8, tCrB)` machinery the INT8 header uses. The +// B-side global load is halved (bytes, not int8 elements) so the mainloop +// is bandwidth-bound on `E * N * K / 2` instead of `E * N * K`. +// +// NOTE: the decode deliberately does NOT feed the packed nibbles straight +// into `reorder(tBrB, tCrB)` via +// `NumericArrayConverter`. In the pinned +// cutlass-sycl that converter is the *interleaved* mixed-input variant -- +// it assumes the B weights were pre-shuffled into the DPAS fragment's +// interleave order. Auto-round ships sequential packing with no shuffle, +// so the fast path permuted a fraction of the B lanes and produced large +// outliers. Staging through `int8_t` keeps the value decode explicit and +// layout-preserving, deferring the register remap to the validated +// `int8_t -> ElementA` reorder. // // Design choice: "option (a) -- halve tile_k at the launcher level" // ----------------------------------------------------------------- @@ -72,17 +94,17 @@ // path in `sycl_tla_moe_mixed.hpp`. // // On-hardware open questions (must be resolved on first build): -// 1. Whether the pinned cutlass-sycl exposes a -// `NumericArrayConverter` -// specialisation. Upstream cutlass 3.5+ ships this converter under -// `cutlass/numeric_conversion.h`; cutlass-sycl may need the same -// pull-in. If the converter is missing, `reorder(tBrB, tCrB)` will -// fail to instantiate at compile time -- the failure mode is a -// "no matching function" error localised to line ~200 below. +// 1. Whether `reorder(tBrB_i8, tCrB)` -- the `int8_t -> ElementA` +// converter shared with the validated INT8 per-group mainloop -- +// instantiates cleanly here. It should: the source dtype and +// fragment layout are identical to the INT8 path (only the upstream +// decode differs), so if the INT8 header builds this one does too. +// The previous `NumericArrayConverter` +// dependency (and its interleaved-B assumption) is no longer used. // 2. Whether `XE_DPAS_TT<8, float, ElementA, ElementA>` (the atom used // by the FP8 / INT8 headers, keeping A/B as the SAME activation // dtype after the upcast) is still the right atom here. It should -// be -- once `reorder` has upcast the packed-nibble fragment to +// be -- once the staged decode + reorder has upcast the fragment to // `ElementA` the atom's operand dtypes match A exactly. // // Copyright (C) 2026 Intel Corporation @@ -134,6 +156,28 @@ using ::ark::moe_dpas_fp8::cute_scalar; using ::ark::moe_dpas_fp8::cute_scalar_t; using ::ark::moe_dpas_fp8::make_moe_tensor; +// --------------------------------------------------------------------------- +// Large-M policy for the S4 packed-nibble path. +// +// There is NO bespoke large-M policy: the `A_avg_M > 32` range reuses the +// validated `dpas_w8a16_policy_m_32` geometry (see the dispatch in +// `moe_prefill_s4_dpas_per_group_dispatch` below). The packed-nibble decode +// (`tBrB` -> `int8_t` staging -> `reorder(tBrB_i8, tCrB)`) is only correct for +// the SUB-GROUP fragment shape the `m_16` / `m_32` policies use, i.e. +// `tile_k = 32` AND `SG_N = tile_n / ATOM_N = 64 / 4 = 16` (`tile_n = 64`, +// `SGLayout<1,4,1>` -> `ATOM_N = 4`). +// +// A previous attempt introduced a wider `<128, 128, 32>` large-M policy with +// `SGLayout<4,2,1>`. That kept `tile_k = 32` but changed `ATOM_N` to 2, so +// `SG_N` became `128 / 2 = 64` (and `sg_n_strides` became 4). At that sub-group +// width the block-2d copy atom derives a B fragment TV-layout that no longer +// lines up with the element-wise `int8_t` staging decode, so a fraction of B +// lanes decode from the wrong nibble and whole output tiles come out grossly +// wrong (observed: max abs diff ~70 on the `medium E=8`, K=14336 prefill shape, +// which routes to the large-M branch because `A_avg_M = 66 > 32`). Pinning +// `tile_k = 32` alone was therefore insufficient -- `SG_N` must also stay 16. +// Large-M reuses `m_32` (more M-tiles at the same validated per-tile decode / +// reorder geometry) rather than a wider, unvalidated tile. // --------------------------------------------------------------------------- // Variant B -- per-K-group S4 (sym) mainloop. // @@ -141,11 +185,16 @@ using ::ark::moe_dpas_fp8::make_moe_tensor; // differences vs. the INT8 per-group mainloop: // // * `ElementB` is required to be `cutlass::int4b_t`. The 4-bit-per- -// element storage triggers CuTe's packed-nibble copy atom and the -// `NumericArrayConverter` specialisation -// inside `reorder(tBrB, tCrB)`, which decodes each byte into two -// sign-extended `ElementA` (bf16/fp16) values in-register. Match -// for match the encoding produced by `moe_dequant::decode_int4_pair +// element storage triggers CuTe's packed-nibble copy atom that loads +// the packed `[N, K/2]` weights into an `int4b_t` register fragment +// (`tBrB`). The mainloop then decodes each `int4b_t` field into an +// `int8_t` staging fragment (`tBrB_i8`) and reorders THAT into the +// `ElementA` (bf16/fp16) MMA fragment via the validated `int8_t -> +// ElementA` `reorder(...)` from the INT8 header -- rather than routing +// the packed nibbles through the interleaved +// `NumericArrayConverter` fast path, which +// assumes a pre-shuffled B layout auto-round does not produce. The +// per-element decode matches `moe_dequant::decode_int4_pair // ` on the auto-round side: byte low nibble decodes to // `q_lo` (K = 2i), byte high nibble to `q_hi` (K = 2i+1), both // sign-extended from [-8, 7]. @@ -268,6 +317,26 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( // per-group path. float sg_scale[sg_n_strides]; + // Scratch int8 fragment, same subgroup TV-layout as the packed int4 copy + // fragment `tBrB`. Each mainloop iteration decodes the just-loaded packed + // nibbles into this buffer (sign-extended [-8, 7]) and then hands it to the + // SAME validated `int8_t -> ElementA` `reorder(...)` the INT8 per-group path + // uses. See the reorder site below for the rationale -- this avoids the + // `NumericArrayConverter` fast path, whose + // interleaved-B-layout assumption miscomputed a fraction of outputs. + // + // NOTE: `make_fragment_like(tBrB)` cannot be used directly. `tBrB` + // is a `SubgroupTensor` of the subbyte `cutlass::int4b_t` copy fragment, and + // its work-item layout is not fully static for subbyte types -- feeding it to + // `make_fragment_like` builds a *dynamic owning* tensor, which CuTe rejects + // (`static_assert(is_static ...)` in `cute/tensor_impl.hpp`). Mirror the + // validated SAGE mainloop instead: allocate a like-fragment of the underlying + // register tensor (`tBrB.tensor()`, static) and re-wrap it as a subgroup + // tensor carrying `tBrB`'s TV-layout so the `reorder(tBrB_i8, tCrB)` below + // resolves to the subgroup-cooperative overload. + auto tBrB_i8 = make_subgroup_tensor( + make_fragment_like(tBrB.tensor()), tBrB.tv_layout()); + CUTE_UNROLL for (; k_tile_prefetch < prefetch_dist; k_tile_prefetch++) { prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); @@ -331,15 +400,33 @@ CUTE_DEVICE void xe_gemm_s4_pergroup( prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); } - // `reorder` performs the in-register `int4b_t -> ElementA` unpack - // + sign-extend + cast via `cutlass::NumericArrayConverter< - // ElementA, cutlass::int4b_t, N>`. Once `tCrB` carries bf16/fp16 - // values it is compatible with the same DPAS atom used by the FP8 - // / INT8 per-group paths. See the header preamble open-question - // (1) -- if the pinned cutlass-sycl is missing this converter - // specialisation this line is where the build fails. + // Decode the packed `int4b_t` B fragment into an `int8_t` staging + // fragment (`tBrB_i8`), then reuse the SAME `int8_t -> ElementA` + // `reorder(...)` the validated INT8 per-group mainloop uses. The + // per-element decode reproduces `decode_int4_pair`'s sign-extension + // exactly: `cutlass::int4b_t` is a signed 4-bit field, so converting + // it to `int` yields the value already sign-extended into [-8, 7]. + // + // This deliberately does NOT call `reorder(tBrB, tCrB)` directly on the + // packed nibbles. That path routes through + // `NumericArrayConverter`, which in the + // pinned cutlass-sycl is the *interleaved* mixed-input converter: it + // assumes the B weights were pre-shuffled into the DPAS fragment's + // interleave order. Auto-round ships sequential `[E, N, K/2]` packing + // with no such shuffle, so the fast path permuted a fraction of the + // B lanes and produced large outliers. Staging through `int8_t` keeps + // the value decode explicit (layout-preserving, one element per + // element of `tBrB`) and defers the register remap to the validated + // `int8_t -> ElementA` reorder -- the fragment layout of `tBrB_i8` + // matches `tBrB`, so this is bit-identical to the INT8 path modulo the + // narrower source dtype. reorder(tArA, tCrA); - reorder(tBrB, tCrB); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tBrB.size(); ++i) { + cutlass::int4b_t packed_nibble = tBrB(i); + tBrB_i8(i) = static_cast(static_cast(packed_nibble)); + } + reorder(tBrB_i8, tCrB); // HOT MAINLOOP -- MMA accumulates into `tCrC_group`. Per-N scale // is applied ONCE at the end of the group in the fold block below. @@ -423,13 +510,12 @@ CUTE_DEVICE void MoEGEMM_s4(const ElementA* Activations, int group_range = item.get_group_range(1); int local_id = item.get_local_linear_id(); - if (group_id == 0 && local_id == 0) { - auto atm = sycl::atomic_ref( - atomic_buffer[0]); - atm.store(0); - } + // NOTE: `atomic_buffer[0]` is zero-initialized on the host (via a queued + // `memset` that the kernel depends on) before launch. It must NOT be reset + // here: an in-kernel `store(0)` by a single work-group has no grid-level + // synchronization, so other work-groups could read an uninitialized/stale + // value or have this late store clobber an already-advanced counter, both of + // which corrupt tile scheduling. int pre_rows = 0; int pre_tiles = 0; @@ -537,9 +623,9 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, const int group_size, int32_t* atomic_buffer) { using ElementA_non_CV = cutlass::platform::remove_cv_t; // DPAS atom keeps its bf16/fp16 x bf16/fp16 -> fp32 shape; the S4 B - // tensor is upcast to ElementA in `reorder(tBrB, tCrB)` in the - // mainloop before entering the MMA atom. See open question #2 in - // the header preamble. + // tensor is decoded to int8 and then upcast to ElementA via + // `reorder(tBrB_i8, tCrB)` in the mainloop before entering the MMA + // atom. See open question #2 in the header preamble. auto op = XE_DPAS_TT<8, float, ElementA_non_CV, ElementA_non_CV>{}; using WGTile = typename policy::WGTile; @@ -548,8 +634,11 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, SGLayout>::TiledMMA; auto mma = MMA{}; + // Query the EU count from the queue's own device (see sycl_tla_moe.hpp): on + // multi-card systems a hardcoded ordinal 0 can name a different device than + // the one `stream` runs on, mis-sizing the grid and corrupting results. int sm_count = - cutlass::KernelHardwareInfo::query_device_multiprocessor_count(0); + stream.get_device().get_info(); auto MaxThreadsPerWorkgroup = size(mma); static constexpr int MaxThreadsPerSM = 512; @@ -572,7 +661,14 @@ void MoEGEMMLauncher_s4(sycl::queue& stream, const ElementA* activations, using GmemTiledCopyB = typename policy::GmemTiledCopyB; using GmemTiledCopyD = typename policy::GmemTiledCopyD; + // Zero-init the persistent-scheduler tile counter on the host before launch. + // Doing it here (rather than inside the kernel) removes the cross-work-group + // race: the kernel below is made to depend on this memset event, so every + // work-group observes a fully initialized counter. + auto init_event = stream.memset(atomic_buffer, 0, sizeof(int32_t)); + auto event = stream.submit([&](sycl::handler& cgh) { + cgh.depends_on(init_event); sycl::local_accessor local_mem(sycl::range<1>(1), cgh); cgh.parallel_for>( @@ -651,10 +747,21 @@ void moe_prefill_s4_dpas_per_group_dispatch( if (A_avg_M <= 8) { ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_16); - } else if (A_avg_M <= 32) { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); } else { - ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy); + // Both the mid-M (`A_avg_M <= 32`) and large-M (`> 32`) ranges use the + // validated `dpas_w8a16_policy_m_32` geometry. The packed-nibble B copy + // fragment + `int8_t` staging + `reorder(tBrB_i8, tCrB)` is only correct + // for the `(tile_k = 32, SG_N = 16)` fragment shape the `m_16` / `m_32` + // policies use (`tile_n = 64`, `SGLayout<1,4,1>` -> `ATOM_N = 4` -> + // `SG_N = tile_n / ATOM_N = 16`). A bespoke large-M policy that widened the + // WG tile to `<128, 128, 32>` (`SGLayout<4,2,1>`) kept `tile_k = 32` but + // changed `SG_N` to `128 / 2 = 64`, which re-broke the nibble<->lane mapping + // for a fraction of B lanes and produced gross outliers (max abs diff ~70 on + // `medium E=8`, K=14336). Pinning `tile_k = 32` alone was insufficient -- + // `SG_N` must also stay 16 -- so large-M reuses the validated `m_32` + // geometry (more M-tiles, identical per-tile decode / reorder) rather than a + // wider, unvalidated tile. See the `Large-M policy` note above. + ARK_DPAS_S4_PG_LAUNCH_SYM(dpas_w8a16_policy_m_32); } #undef ARK_DPAS_S4_PG_LAUNCH_SYM @@ -662,24 +769,33 @@ void moe_prefill_s4_dpas_per_group_dispatch( } // --------------------------------------------------------------------------- -// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default ON, semantics -// identical to `moe_prefill_dpas_int_enabled` / `moe_prefill_dpas_fp8 -// _enabled`). Decoupled from `ARK_MOE_PREFILL_DPAS_INT8` so this new -// single-pass path can be disabled in isolation if it regresses -- -// switching S4 off falls back to the S4->S8 upcast + INT8 DPAS path -// which is itself gated by `ARK_MOE_PREFILL_DPAS_INT8`. +// Env-flag helper -- `ARK_MOE_PREFILL_DPAS_S4` (default OFF). Decoupled +// from `ARK_MOE_PREFILL_DPAS_INT8` so this single-pass path can be +// enabled in isolation for A/B validation -- with S4 off (the default) +// int4-sym falls back to the bit-exact generic dequant path (the two-pass +// S4->S8 upcast + INT8 DPAS path is itself default OFF behind +// `ARK_MOE_PREFILL_DPAS_LOWBIT`, since it exhibits the same numeric defect). +// +// STATUS: default OFF pending on-hardware validation. The large-M +// gross-outlier defect (observed: max abs diff ~70 on the `medium E=8`, +// K=14336 int4-sym + fp16 accuracy case) was root-caused to the bespoke +// large-M policy breaking the packed-nibble decode's `SG_N = 16` invariant +// (it used `SG_N = 64`); large-M now reuses the validated `m_32` geometry. +// The gate stays default OFF until the fix is confirmed on hardware; callers +// get the bit-exact generic dequant path by default, and the single-pass +// path stays available behind `ARK_MOE_PREFILL_DPAS_S4=1`. // -// Truthy values (case-insensitive): "1", "true", "on", "yes". -// Explicit "0" / "false" / "off" / "no" disable. Re-read on every -// call so benchmarks / tests can toggle the path in-process. +// Truthy values (case-insensitive): "1", "true", "on", "yes" enable. +// Anything else (including unset) leaves the path disabled. Re-read on +// every call so benchmarks / tests can toggle the path in-process. // --------------------------------------------------------------------------- inline bool moe_prefill_dpas_s4_enabled() { const char* env = std::getenv("ARK_MOE_PREFILL_DPAS_S4"); - if (env == nullptr) return true; // default ON + if (env == nullptr) return false; // default OFF std::string s(env); for (auto& c : s) c = static_cast(std::tolower(static_cast(c))); - if (s == "0" || s == "false" || s == "off" || s == "no") return false; - return true; + if (s == "1" || s == "true" || s == "on" || s == "yes") return true; + return false; } // --------------------------------------------------------------------------- diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa.hpp index aabe8bff0..5fd7477cc 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa.hpp @@ -62,6 +62,17 @@ namespace detail { using namespace cute; +// Query the multiprocessor (Xe-core / EU) count from the *default queue's own +// device*. The caller binds the default queue to the tensor's queue via +// `compat::set_default_queue(*q)` and the kernel is launched on that same queue, +// so this reports the count for the device the kernel actually runs on. Using a +// hardcoded ordinal 0 instead can name a different device on multi-card systems, +// mis-sizing the persistent-scheduler grid and producing wrong results that only +// manifest when more than one device is visible. +inline int query_default_queue_sm_count() { + return compat::get_default_queue().get_device().get_info(); +} + // Command line options parsing struct Options { const void *q = nullptr, *k = nullptr, *v = nullptr; @@ -371,9 +382,10 @@ struct FMHAConfig { // // The KernelHardwareInfo struct holds the number of EUs on the GPU with a given device ID. This - // information is used by the underlying kernel. + // information is used by the underlying kernel. Derive it from the default queue's own device so + // the grid is sized for the device the kernel actually runs on (see query_default_queue_sm_count). cutlass::KernelHardwareInfo hw_info; - hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + hw_info.sm_count = query_default_queue_sm_count(); using ProblemShapeType = cutlass::fmha::kernel::FMHAProblemShape; @@ -733,9 +745,10 @@ struct SageConfig { // // The KernelHardwareInfo struct holds the number of EUs on the GPU with a given device ID. This - // information is used by the underlying kernel. + // information is used by the underlying kernel. Derive it from the default queue's own device so + // the grid is sized for the device the kernel actually runs on (see query_default_queue_sm_count). cutlass::KernelHardwareInfo hw_info; - hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + hw_info.sm_count = query_default_queue_sm_count(); using ProblemShapeType = cutlass::fmha::kernel::SageProblemShape; 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..779ee0061 100644 --- a/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md +++ b/auto_round_extension/ark/test/README_MOE_PREFILL_PERF.md @@ -276,9 +276,9 @@ falls through to the dequant path. | Precedence | Env flag | Kernel | | ----------- | --------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| 1 (highest) | `ARK_MOE_PREFILL_DPAS_S4` unset or truthy (**default ON**) | **S4-sym single-pass DPAS mixed-input mainloop.** Reads packed `[E, N, K/2]` `uint8_t` nibbles directly and folds the S4→`act_dtype` upcast into the DPAS mainloop via CuTe's `reorder(tBrB, tCrB)` (which relies on `NumericArrayConverter`). B-side global traffic is exactly half of the S8 path. Per-K-group scale is applied through the same deferred group-boundary fold as INT8. Implemented in `sycl_tla_moe_prefill_s4_dpas.hpp`. **Status: NEEDS-HARDWARE-VALIDATION** (untested port). | -| 2 (fallback)| `ARK_MOE_PREFILL_DPAS_S4=0` and `ARK_MOE_PREFILL_DPAS_INT8` truthy (**default ON**) | **S4→S8 upcast + shared INT8 DPAS mainloop.** Two-pass: `launch_upcast_int4_sym_to_int8` writes an `[E, N, K]` `int8_t` view of the dequant workspace, then the standard INT8 per-group DPAS mainloop consumes it. Robust but pays the ~E·N·K byte round-trip vs. path 1. Implemented in `sycl_tla_moe_mixed.hpp` + `sycl_tla_moe_prefill_int_dpas.hpp`. | -| 3 (default) | `ARK_MOE_PREFILL_DPAS_S4=0` and `ARK_MOE_PREFILL_DPAS_INT8=0` | v1 dequant kernel (`sycl_tla_moe_mixed.hpp::launch_dequant_int4`) followed by the stock bf16/fp16 grouped GEMM. Handles both sym and asym. | +| 1 (highest) | `ARK_MOE_PREFILL_DPAS_S4=1` (**default OFF**) | **S4-sym single-pass DPAS mixed-input mainloop.** Reads packed `[E, N, K/2]` `uint8_t` nibbles directly and folds the S4→`act_dtype` upcast into the DPAS mainloop. Each `int4b_t` B fragment is decoded to an `int8_t` staging fragment in registers and then handed to the validated `int8_t → act_dtype` `reorder(...)` shared with the INT8 path (avoiding the interleaved `NumericArrayConverter` fast path, which assumed a pre-shuffled B layout auto-round does not produce). B-side global traffic is exactly half of the S8 path. Per-K-group scale is applied through the same deferred group-boundary fold as INT8. Implemented in `sycl_tla_moe_prefill_s4_dpas.hpp`. **Status: NEEDS-HARDWARE-VALIDATION — default OFF pending on-hardware confirmation. The large-M gross-outlier defect (max abs diff ~70 on `medium E=8`, K=14336 int4-sym + fp16) was root-caused to the bespoke large-M policy breaking the packed-nibble decode's `SG_N = 16` invariant (it used `SG_N = 64`); large-M now reuses the validated `m_32` geometry.** | +| 2 | `ARK_MOE_PREFILL_DPAS_LOWBIT=1` (**default OFF**) with `ARK_MOE_PREFILL_DPAS_S4` unset/`0` and `ARK_MOE_PREFILL_DPAS_INT8` truthy | **S4→S8 upcast + shared INT8 DPAS mainloop.** Two-pass: `launch_upcast_int4_sym_to_int8` writes an `[E, N, K]` `int8_t` view of the dequant workspace, then the standard INT8 per-group DPAS mainloop consumes it. Robust but pays the ~E·N·K byte round-trip vs. path 1. Implemented in `sycl_tla_moe_mixed.hpp` + `sycl_tla_moe_prefill_int_dpas.hpp`. **Status: NEEDS-HARDWARE-VALIDATION — defaulted OFF (`ARK_MOE_PREFILL_DPAS_LOWBIT`) because it exhibits the same ~70 max-abs-diff defect as path 1 on `medium E=8`, K=14336 int4-sym + fp16; opt in only for debugging.** | +| 3 (default) | `ARK_MOE_PREFILL_DPAS_S4` unset/`0` and `ARK_MOE_PREFILL_DPAS_LOWBIT` unset/`0` (or `ARK_MOE_PREFILL_DPAS_INT8=0`) | v1 dequant kernel (`sycl_tla_moe_mixed.hpp::launch_dequant_int4`) followed by the stock bf16/fp16 grouped GEMM. Bit-exact reference-grade path; handles both sym and asym. **This is the default for int4-sym / int2-sym prefill.** | **S4 DPAS path shape preconditions** — the `moe_gemm_prefill` dispatcher silently falls back to precedence 2 (then 3) whenever any of 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..2102aa1a6 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 @@ -208,9 +208,9 @@ S4-sym 有两条独立的 DPAS 路径;asym S4 始终回退到 dequant 路径。 | 优先级 | Env 开关 | Kernel | | ------------ | ------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| 1 (最高) | `ARK_MOE_PREFILL_DPAS_S4` 未设置或为真值(**默认开启**) | **S4-sym 单遍 DPAS 混合输入 mainloop**。直接读取 packed `[E, N, K/2]` `uint8_t` nibble,通过 CuTe `reorder(tBrB, tCrB)`(依赖 `NumericArrayConverter`)在寄存器中把 S4 上转到 `act_dtype`。B 侧 global 带宽正好是 S8 路径的一半。Per-K-group scale 使用与 INT8 相同的组边界延迟折叠。实现于 `sycl_tla_moe_prefill_s4_dpas.hpp`。**状态:NEEDS-HARDWARE-VALIDATION**(未经测试的移植)。 | -| 2 (回退) | `ARK_MOE_PREFILL_DPAS_S4=0` 且 `ARK_MOE_PREFILL_DPAS_INT8` 为真值(**默认开启**) | **S4→S8 上转 + 共享 INT8 DPAS mainloop**。两遍:`launch_upcast_int4_sym_to_int8` 把权重写成 `[E, N, K]` `int8_t`(复用 dequant workspace),再由标准 INT8 per-group DPAS mainloop 消费。相较路径 1 需要付出 ~E·N·K 字节的 workspace 往返。实现于 `sycl_tla_moe_mixed.hpp` + `sycl_tla_moe_prefill_int_dpas.hpp`。 | -| 3 (默认回退) | `ARK_MOE_PREFILL_DPAS_S4=0` 且 `ARK_MOE_PREFILL_DPAS_INT8=0` | v1 dequant kernel(`sycl_tla_moe_mixed.hpp::launch_dequant_int4`)后接标准 bf16/fp16 grouped GEMM。同时支持 sym 与 asym。 | +| 1 (最高) | `ARK_MOE_PREFILL_DPAS_S4=1`(**默认关闭**) | **S4-sym 单遍 DPAS 混合输入 mainloop**。直接读取 packed `[E, N, K/2]` `uint8_t` nibble,在寄存器中把 S4 上转到 `act_dtype`。每个 `int4b_t` B fragment 先解码为 `int8_t` 暂存 fragment,再交给与 INT8 路径共享的、已验证的 `int8_t → act_dtype` `reorder(...)`(从而绕开会假设 B 已预置换布局的交织式 `NumericArrayConverter` 快路径,而 auto-round 并不产生这种置换)。B 侧 global 带宽正好是 S8 路径的一半。Per-K-group scale 使用与 INT8 相同的组边界延迟折叠。实现于 `sycl_tla_moe_prefill_s4_dpas.hpp`。**状态:NEEDS-HARDWARE-VALIDATION —— 默认关闭,待硬件确认。此前在大 M 的 `medium E=8`、K=14336 int4-sym + fp16 上观测到的 ~70 最大绝对误差缺陷,已定位为 bespoke 大 M policy 破坏了 packed-nibble 解码所依赖的 `SG_N = 16` 不变量(它用了 `SG_N = 64`);大 M 现改为复用已验证的 `m_32` 几何。** | +| 2 | `ARK_MOE_PREFILL_DPAS_LOWBIT=1`(**默认关闭**),且 `ARK_MOE_PREFILL_DPAS_S4` 未设置或为 `0`,且 `ARK_MOE_PREFILL_DPAS_INT8` 为真值 | **S4→S8 上转 + 共享 INT8 DPAS mainloop**。两遍:`launch_upcast_int4_sym_to_int8` 把权重写成 `[E, N, K]` `int8_t`(复用 dequant workspace),再由标准 INT8 per-group DPAS mainloop 消费。相较路径 1 需要付出 ~E·N·K 字节的 workspace 往返。实现于 `sycl_tla_moe_mixed.hpp` + `sycl_tla_moe_prefill_int_dpas.hpp`。**状态:NEEDS-HARDWARE-VALIDATION —— 因在 `medium E=8`、K=14336 int4-sym + fp16 上表现出与路径 1 相同的 ~70 最大绝对误差缺陷,已通过 `ARK_MOE_PREFILL_DPAS_LOWBIT` 默认关闭,仅供调试时手动开启。** | +| 3 (默认) | `ARK_MOE_PREFILL_DPAS_S4` 未设置或为 `0`,且 `ARK_MOE_PREFILL_DPAS_LOWBIT` 未设置或为 `0`(或 `ARK_MOE_PREFILL_DPAS_INT8=0`） | v1 dequant kernel(`sycl_tla_moe_mixed.hpp::launch_dequant_int4`)后接标准 bf16/fp16 grouped GEMM。位精确的参考级路径;同时支持 sym 与 asym。**这是 int4-sym / int2-sym prefill 的默认路径。** | **S4 DPAS 路径形状前置条件** — 任何条件不满足时,`moe_gemm_prefill` 分发器会静默回退到优先级 2(再回退到 3): diff --git a/auto_round_extension/ark/test/test_moe_prefill_accuracy.py b/auto_round_extension/ark/test/test_moe_prefill_accuracy.py index c8c1ee168..23019375d 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_accuracy.py +++ b/auto_round_extension/ark/test/test_moe_prefill_accuracy.py @@ -186,19 +186,65 @@ def _tol_for_dtype(base, dtype): bf16 has 7 mantissa bits vs fp16's 10, so K-reductions of ~14K elements (the ``medium E=8`` prefill shape uses K=14336) can produce a handful of - outliers per ~2M outputs that exceed the fp16-calibrated ``7e-2`` bound - by up to ~1.3x (observed: max abs diff ~0.090). This shows up - stochastically across seeds on every quantized path -- int8-sym / fp8 / - DPAS mainloops most consistently, but int4/int2 as well when a rare - outlier lands past the bound -- so we widen for all quant callers on - bf16 and leave fp16 (which absorbs the noise in its wider mantissa) at - the tight bound. + outliers per ~2M outputs that exceed the fp16-calibrated bound by up to + ~1.3x (observed: max abs diff ~0.090). This is genuine accumulator-order + noise, so we widen only for bf16 quant callers and leave fp16 (whose wider + mantissa absorbs the noise) at the tight, quant-appropriate bound. + + fp16 is deliberately NOT loosened here. A previous change widened fp16 to a + ``1e-1`` floor to silence an "intermittent" int4-sym + fp16 failure, but the + observed failure (max abs diff ~70 at a single index) is orders of magnitude + beyond any accumulator noise and beyond the ``1e-1`` floor itself -- i.e. the + loosened tolerance never even made the assertion pass. Such a single wildly + wrong element is a kernel bug (boundary / uninitialized-memory / predication + in the S4 DPAS prefill path), not RNG luck, and must be fixed in the kernel + rather than masked here. Keeping fp16 strict ensures that bug stays visible. """ if dtype is torch.bfloat16: return dict(rtol=max(base["rtol"], 1e-1), atol=max(base["atol"], 1e-1)) return base +def _assert_close_report(out, ref, tol, label): + """``torch.testing.assert_close`` with a localized outlier report on failure. + + When the assertion fails this augments the error with a breakdown that makes + the failure *attributable* to a root cause instead of just a single "greatest + absolute difference" number: + + * the greatest abs diff and its index, + * how many elements exceed a coarse ``1.0`` threshold and their (first few) + indices -- a *single* / handful of huge outliers points at a kernel + boundary or uninitialized-memory bug, whereas a large *fraction* of + elements being off points at a global layout / dequant / scale bug. + + This directly supports triaging the MoE prefill kernels: accumulator-order + noise is many tiny sub-``0.1`` diffs, so any element past ``1.0`` is a real + correctness defect regardless of the configured tolerance. + """ + try: + torch.testing.assert_close(out, ref, msg=lambda m: f"[{label}] {m}", **tol) + except AssertionError as err: + diff = (out.detach().float() - ref.detach().float()).abs() + total = diff.numel() + gross_mask = diff > 1.0 + n_gross = int(gross_mask.sum().item()) + max_diff, max_flat = diff.reshape(-1).max(dim=0) + max_idx = tuple(int(i) for i in torch.unravel_index(max_flat, diff.shape)) + gross_idx = [tuple(int(i) for i in idx) for idx in torch.nonzero(gross_mask)[:8].tolist()] if n_gross else [] + report = ( + f"\n[{label}] outlier report:" + f"\n shape={tuple(diff.shape)} total={total}" + f"\n max abs diff={float(max_diff):.6g} at index {max_idx}" + f"\n elements with |diff|>1.0: {n_gross} ({100.0 * n_gross / total:.4f}%)" + f"\n first gross indices: {gross_idx}" + f"\n -> a single / handful of huge outliers indicates a kernel" + f" boundary or uninitialized-memory bug (fix the kernel, not the tol);" + f"\n a large fraction indicates a global layout / dequant / scale bug." + ) + raise AssertionError(str(err) + report) from None + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -243,7 +289,7 @@ def test_accuracy_fp(self, dtype): assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" assert out.dtype == dtype, f"{label}: bad dtype {out.dtype}" - torch.testing.assert_close(out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_TOL_FP) + _assert_close_report(out, ref, _TOL_FP, label) @pytest.mark.skipif(bool(_QUANT_PREFILL_SKIP), reason=_QUANT_PREFILL_SKIP or "ok") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -280,9 +326,7 @@ def test_accuracy_int4(self, dtype, asym): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT4, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_INT4, dtype), label) @pytest.mark.skipif(bool(_QUANT_PREFILL_SKIP), reason=_QUANT_PREFILL_SKIP or "ok") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -319,9 +363,7 @@ def test_accuracy_int8(self, dtype, asym): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT8, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_INT8, dtype), label) @pytest.mark.skipif(bool(_QUANT_PREFILL_SKIP), reason=_QUANT_PREFILL_SKIP or "ok") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -358,9 +400,7 @@ def test_accuracy_int2(self, dtype, asym): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT2, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_INT2, dtype), label) @pytest.mark.skipif(bool(_QUANT_PREFILL_SKIP), reason=_QUANT_PREFILL_SKIP or "ok") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -407,9 +447,7 @@ def test_accuracy_fp8(self, dtype, fp8_dtype, native): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_FP8, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_FP8, dtype), label) finally: if prev is None: os.environ.pop("ARK_MOE_PREFILL_NATIVE_FP8", None) @@ -468,9 +506,7 @@ def test_accuracy_fp8_per_tensor_dpas(self, dtype, fp8_dtype): ref = _reference_moe_prefill(activations, dequant_NK, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_FP8, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_FP8, dtype), label) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_accuracy_int8_per_tensor_dpas(self, dtype): @@ -522,9 +558,7 @@ def test_accuracy_int8_per_tensor_dpas(self, dtype): ref = _reference_moe_prefill(activations, dequant_NK, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT8, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_INT8, dtype), label) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) @@ -566,9 +600,7 @@ def test_accuracy_fp8_dpas_per_group(self, dtype, fp8_dtype): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_FP8, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_FP8, dtype), label) finally: if prev is None: os.environ.pop("ARK_MOE_PREFILL_DPAS_FP8", None) @@ -618,9 +650,7 @@ def test_accuracy_int8_dpas_per_group(self, dtype): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT8, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_INT8, dtype), label) finally: if prev is None: os.environ.pop("ARK_MOE_PREFILL_DPAS_INT8", None) @@ -634,11 +664,11 @@ def test_accuracy_int4_dpas_per_group(self, dtype): Sibling of :meth:`test_accuracy_int8_dpas_per_group` -- uses the standard ``moe_gemm_prefill`` call path with - ``ARK_MOE_PREFILL_DPAS_S4=1`` (default ON) and - ``ARK_MOE_PREFILL_DPAS_INT8=0`` so the two-pass S4->S8 upcast - fallback is disabled and we exercise the single-pass mainloop - exclusively. The C++ dispatcher should pick the S4 DPAS branch - for shapes that satisfy ``N%64==0 && K%32==0 && K%group_size==0 + ``ARK_MOE_PREFILL_DPAS_S4=1`` (explicitly opting in; the path is + default OFF) and ``ARK_MOE_PREFILL_DPAS_INT8=0`` so the two-pass + S4->S8 upcast fallback is disabled and we exercise the single-pass + mainloop exclusively. The C++ dispatcher should pick the S4 DPAS + branch for shapes that satisfy ``N%64==0 && K%32==0 && K%group_size==0 && group_size%2==0 && group_size in {32,64,128,256}`` and silently fall back to the dequant path otherwise, so this test is checking parity, not that the DPAS branch is exercised. @@ -674,9 +704,7 @@ def test_accuracy_int4_dpas_per_group(self, dtype): ref = _reference_moe_prefill(activations, dequant, ntpe) assert out.shape == (total_tokens, N), f"{label}: bad shape {out.shape}" - torch.testing.assert_close( - out, ref, msg=lambda m, lbl=label: f"[{lbl}] {m}", **_tol_for_dtype(_TOL_INT4, dtype) - ) + _assert_close_report(out, ref, _tol_for_dtype(_TOL_INT4, dtype), label) finally: if prev_s4 is None: os.environ.pop("ARK_MOE_PREFILL_DPAS_S4", None) diff --git a/auto_round_extension/ark/test/test_moe_prefill_perf.py b/auto_round_extension/ark/test/test_moe_prefill_perf.py index c322cb66b..4f0416941 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_perf.py +++ b/auto_round_extension/ark/test/test_moe_prefill_perf.py @@ -668,15 +668,16 @@ def test_perf_int4(self, dtype, asym): deq_ms = _xpu_time_ms(lambda: _dequant_int4_sym(packed, scales, group_size).to(dtype)) base_ms = _xpu_time_ms(lambda: _default_moe_prefill(act_padded, dequant)) - # Default ARK path (dequant + GEMM). INT4-sym is DPAS-accelerated - # via TWO independent branches inside `moe_gemm_prefill`: - # 1. `ARK_MOE_PREFILL_DPAS_S4=1` (default ON) -- single-pass - # mainloop reading packed nibbles directly via CuTe's - # `NumericArrayConverter` in - # `reorder(tBrB, tCrB)`. Preferred; the new hot path. - # 2. `ARK_MOE_PREFILL_DPAS_INT8=1` (default ON) -- two-pass - # S4->S8 upcast into workspace + shared INT8 DPAS - # mainloop. Fallback for when (1) is disabled. + # Default ARK path (dequant + GEMM). INT4-sym has TWO opt-in DPAS + # branches inside `moe_gemm_prefill`, both default OFF (int4-sym + # prefill defaults to the bit-exact dequant + GEMM path): + # 1. `ARK_MOE_PREFILL_DPAS_S4=1` (default OFF) -- single-pass + # mainloop reading packed nibbles directly, decoding each + # `int4b_t` fragment to `int8_t` in registers and reusing + # the validated `int8_t -> act` reorder. Opt-in for debugging. + # 2. `ARK_MOE_PREFILL_DPAS_LOWBIT=1` (default OFF) -- two-pass + # S4->S8 upcast into workspace + shared INT8 DPAS mainloop. + # Opt-in for debugging. # Force BOTH off for this measurement so the `ark(ms)` column # reflects the legacy dequant + BF16 GEMM path independently # of the `dpas(ms)` column below. diff --git a/requirements.txt b/requirements.txt index 1bb587b64..2c1e75186 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -# 1.5.1