diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 512fab6a55..a688ad4efa 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -101,7 +101,17 @@ def get_stream(A: torch.Tensor) -> int: if A.device.type == "cpu": return 0 if A.device.type == "xpu": - return torch.xpu.current_stream().sycl_queue + # Query the stream for *A's own device*, not the global current device. + # `torch.xpu.current_stream()` with no argument resolves the device via + # `torch.xpu.current_device()`, a process-global that another operation + # (or, in a test suite, a preceding test) may have left pointing at a + # different card. Passing `A.device` guarantees the returned SYCL queue + # runs on the same device the tensor's memory lives on; otherwise the + # native kernel would launch on one card while dereferencing pointers + # into another card's memory, silently corrupting results. On a + # single-visible-device system the two always coincide, which is why + # the mismatch only surfaces with multiple cards visible. + return torch.xpu.current_stream(A.device).sycl_queue def _normalize_tensor_layout(tensor_layout: str) -> str: @@ -2044,6 +2054,19 @@ def _moe_gemm_prefill_int_pertensor( return outputs +def _alloc_moe_prefill_workspace(num_experts: int, K: int, N: int, device, dtype) -> torch.Tensor: + """Allocate the ``[E, K, N]`` dequant workspace, fully zero-initialised. + + The generic dequant kernels write only the first ``N`` columns of every row + and the downstream Grouped-GEMM reads exactly that ``[E, K, N]`` region -- + it never over-reads past the nominal extent (confirmed by NaN-poisoning the + trailing capacity and widening it: the poison never leaked into the output). + Zero-initialising (rather than ``torch.empty``) keeps zero-token experts and + unwritten slices reading back as deterministic zeros. + """ + return torch.zeros((num_experts, K, N), device=device, dtype=dtype) + + def moe_gemm_prefill( activations: torch.Tensor, weights: torch.Tensor, @@ -2174,13 +2197,20 @@ def moe_gemm_prefill( weights_ptr = dequant_workspace.data_ptr() workspace_ptr = dequant_workspace.data_ptr() else: - # Reuse a persistent `[E, K, N]` workspace across calls with the same - # (device, dtype, E, K, N). For real MoE prefill workloads the same - # shape is dispatched on every iteration; allocating a fresh - # `E*K*N*sizeof(act)` tensor each call adds non-trivial caching- - # allocator overhead (and, on the small shapes, dominates the - # quantized GEMM cost). The workspace is kept alive by the cache so - # we hand the data_ptr() to the kernel without taking a new ref. + # Allocate a fresh `[E, K, N]` workspace for every call rather than + # sharing a persistent, module-level cached buffer. + # + # A cached buffer is only safe when calls that share the same + # (device, dtype, E, K, N) shape never overlap. That assumption breaks + # under multi-card / multi-stream execution (e.g. tensor-parallel MoE + # prefill): two concurrently-launched kernels with the same shape would + # be handed the *same* `data_ptr()`, so one launch's dequant write races + # the other launch's Grouped-GEMM read of the same scratch. On INT4 + # prefill this reproduces deterministically on multi-card runs and + # corrupts the output. Allocating a distinct tensor per call removes the + # aliasing entirely; PyTorch's caching allocator makes the repeated + # allocation cheap (reused device blocks, no fresh device malloc on the + # steady-state path). # # We allocate the workspace unconditionally for all quantized paths, # including native FP8. The native FP8 launcher fuses GEMM+scale and @@ -2193,10 +2223,25 @@ def moe_gemm_prefill( # (`N % 16`, `K % 32`, `K % group_size`, `group_size % 32`) may not # hold, or the act dtype may not be F16/BF16. Without a workspace the # fall-through would hit the generic null-pointer check in - # `sycl_tla_moe_mixed.hpp` and raise. Since the workspace lives in - # the module-level cache, allocation happens once per shape and adds - # no per-call overhead when the native path is taken. - dequant_workspace = _get_moe_prefill_workspace(activations.device, activations.dtype, num_experts, K, N) + # `sycl_tla_moe_mixed.hpp` and raise. + # + # Zero-initialise (rather than `torch.empty`) so experts that receive + # no tokens in this prefill batch have deterministic, zeroed rows. The + # generic `[E, K, N]` dequant kernels in `sycl_tla_moe_mixed.hpp` skip + # every expert with `num_tokens_per_expert[e] == 0` and never write its + # slice of the workspace; leaving that slice uninitialised exposes stale + # allocator memory to any consumer that still reads those rows. + # + # Allocate exactly the ``[E, K, N]`` footprint. The native kernels + # address the buffer as a flat ``E * K * N`` region via ``workspace_ptr`` + # (stride N) and the downstream Grouped-GEMM reads only that region -- + # it does not over-read past the nominal extent. (A previous 2× headroom + # plus NaN-poison diagnostic confirmed this: poisoning and widening the + # trailing capacity never leaked NaN into the output, so the extra + # padding was unnecessary and has been removed.) + dequant_workspace = _alloc_moe_prefill_workspace( + num_experts, K, N, device=activations.device, dtype=activations.dtype + ) weights_ptr = weights.data_ptr() workspace_ptr = dequant_workspace.data_ptr() @@ -2225,58 +2270,37 @@ def moe_gemm_prefill( # (see `moe_detail::moe_gemm_launcher` in `sycl_tla_moe.hpp`), so by the # time `lib.moe_gemm_prefill` returns the device has already consumed the # workspace. For the unquantized fast path the workspace is a per-call - # transposed copy of `weights` -- drop it now. For the quantized paths - # the workspace lives in the module-level cache (`_get_moe_prefill_workspace`) - # and is intentionally retained for reuse on the next call. The native - # fp8 path allocates no workspace at all, so there is nothing to drop. - if is_unquantized: - del dequant_workspace + # transposed copy of `weights`; for the quantized paths it is the per-call + # `[E, K, N]` dequant scratch allocated above. Either way it is a local, + # non-shared buffer that is safe to drop now (the native fp8 path allocates + # no workspace at all, so `dequant_workspace` is simply unused there). + del dequant_workspace return outputs # --------------------------------------------------------------------------- -# `moe_gemm_prefill` dequant-workspace cache. +# `moe_gemm_prefill` dequant-workspace. # # The Stage-1 quantized prefill kernel dequantises weights into an # `[E, K, N]` act-dtype scratch buffer before dispatching to the existing -# CUTLASS-SYCL grouped GEMM. In real model usage the same `(E, K, N, dtype)` -# tuple is hit on every prefill step, so allocating a fresh -# `E * K * N * sizeof(act_dtype)` tensor per call adds caching-allocator -# overhead that is significant on the small/medium shapes. +# CUTLASS-SYCL grouped GEMM. This scratch is now allocated fresh per call +# (see `moe_gemm_prefill`) rather than shared through a module-level cache: +# a shared buffer aliases across concurrent same-shape launches on multi-card +# / multi-stream setups and races the dequant write against the GEMM read. # -# We cache one tensor per `(device, dtype, E, K, N)` key. The cache holds -# references that keep the tensors alive across calls; callers can clear it -# explicitly via `clear_moe_prefill_workspace_cache()` if they need to -# release the memory (e.g., before allocating large buffers for a different -# subsystem). +# `clear_moe_prefill_workspace_cache()` is retained as a backwards-compatible +# no-op for callers that used to drop the cache explicitly. # --------------------------------------------------------------------------- -_MOE_PREFILL_WORKSPACE_CACHE: "dict[tuple, torch.Tensor]" = {} +def clear_moe_prefill_workspace_cache() -> None: + """Deprecated no-op. -def _get_moe_prefill_workspace(device: torch.device, dtype: torch.dtype, E: int, K: int, N: int) -> torch.Tensor: - """Return a persistent `[E, K, N]` workspace tensor for the prefill kernel. - - The tensor is allocated lazily on first use and retained in a module-level - cache so subsequent calls with the same `(device, dtype, E, K, N)` reuse - the same memory. Returned tensors are contiguous and uninitialised; the - kernel writes every element before reading. + The `moe_gemm_prefill` dequant workspace is no longer cached across calls; + each call allocates and releases its own buffer, so there is nothing to + clear. Kept for backwards compatibility with earlier callers. """ - # `device` may be a `torch.device` or a string; normalise so the cache key - # is hashable and identifies the exact device (including ordinal). - if not isinstance(device, torch.device): - device = torch.device(device) - key = (device.type, device.index, dtype, int(E), int(K), int(N)) - ws = _MOE_PREFILL_WORKSPACE_CACHE.get(key) - if ws is None: - ws = torch.empty((E, K, N), device=device, dtype=dtype) - _MOE_PREFILL_WORKSPACE_CACHE[key] = ws - return ws - - -def clear_moe_prefill_workspace_cache() -> None: - """Release all cached `moe_gemm_prefill` dequant-workspace tensors.""" - _MOE_PREFILL_WORKSPACE_CACHE.clear() + return None # --------------------------------------------------------------------------- diff --git a/auto_round_extension/ark/auto_round_kernel/sdpa.cpp b/auto_round_extension/ark/auto_round_kernel/sdpa.cpp index 10afcf6469..48525a06a8 100644 --- a/auto_round_extension/ark/auto_round_kernel/sdpa.cpp +++ b/auto_round_extension/ark/auto_round_kernel/sdpa.cpp @@ -244,6 +244,7 @@ void sage_prefill(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O options.vscale = vscale; options.lse = lse; compat::set_default_queue(*q); + options.queue = q; KernelLauncher launcher = select_sage_prefill_launcher(q_dtype, pv_dtype, head_dim, use_int8_pv); if (launcher == nullptr) { @@ -268,6 +269,7 @@ void flash_attn_prefill(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, v num_heads_kv, seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal); options.lse = lse; compat::set_default_queue(*q); + options.queue = q; KernelLauncher launcher = select_prefill_launcher(q_dtype, head_dim); if (launcher == nullptr) { @@ -292,6 +294,7 @@ void flash_attn_decode(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, vo num_heads_kv, 1, seq_len_kv, head_dim, softmax_scale, is_causal); options.lse = lse; compat::set_default_queue(*q); + options.queue = q; KernelLauncher launcher = select_decode_launcher(q_dtype, head_dim); if (launcher == nullptr) { @@ -448,6 +451,7 @@ void sage_prefill_varlen(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, options.lse = lse; compat::set_default_queue(*q); + options.queue = q; // Zero-filled workspace for cu_seqlens_kv_cache via DnnlContext scratch pool. int* zero_cu_buf = static_cast( @@ -509,6 +513,7 @@ void sdpa_varlen_impl(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, voi options.lse = lse; compat::set_default_queue(*q); + options.queue = q; // When isVarLen=true, the kernel's apply_variable_length accesses // cumulative_length for ALL three fields. Even with max_seqlen_kv_cache=0, diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe.hpp index cff5626de8..5e28a0ecc1 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 e6a3ececcd..c5ee6cf7c3 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)) { @@ -947,6 +963,14 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v q, static_cast(weights), upcast_i8, num_experts, N, K, num_tokens_per_expert); } + // The upcast kernel above is submitted asynchronously and captures no + // event, while the DPAS dispatch below reads the same `upcast_i8` + // workspace. On an out-of-order stream (PyTorch XPU queues are + // out-of-order) submission order alone does not serialise the two, so the + // GEMM could read the workspace before the upcast finishes writing it -- + // a read-before-write race that yields localized garbage. Block until the + // upcast writes are visible before the GEMM consumes them. + q->wait(); if (act_dtype == BTLA_DTYPE::F16) { using ScalarT = sycl::half; moe_dpas_int::moe_prefill_int_dpas_per_group_dispatch( @@ -1083,12 +1107,26 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v auto* w_kn = static_cast(dequant_workspace); moe_mixed_detail::dequant_to_KN(q, weights, scales, zeros, w_kn, weight_dtype, num_experts, N, K, group_size, asym, num_tokens_per_expert); + // The dequant kernels are submitted asynchronously and capture no event, + // whereas `moe_gemm` reads the `[E, K, N]` workspace they just wrote. On an + // out-of-order stream (PyTorch XPU queues are out-of-order) submission order + // does not serialise the two launches, so the grouped GEMM can read the + // workspace before the dequant pass has finished writing it -- a + // read-before-write race that surfaces as a handful of localized, wildly + // wrong outputs whose values depend on the workspace's prior contents. + // Wait for the dequant writes to become visible before the GEMM consumes + // them. + q->wait(); moe_gemm(q, activations, w_kn, /*scales=*/nullptr, outputs, act_dtype, N, K, num_tokens_per_expert, num_experts); } else if (act_dtype == BTLA_DTYPE::BF16) { using BF = sycl::ext::oneapi::bfloat16; auto* w_kn = static_cast(dequant_workspace); moe_mixed_detail::dequant_to_KN(q, weights, scales, zeros, w_kn, weight_dtype, num_experts, N, K, group_size, asym, num_tokens_per_expert); + // See the F16 branch above: serialise the async dequant writes before the + // grouped GEMM reads the workspace to avoid a read-before-write race on + // out-of-order streams. + q->wait(); moe_gemm(q, activations, w_kn, /*scales=*/nullptr, outputs, act_dtype, N, K, num_tokens_per_expert, num_experts); } else { throw std::invalid_argument("moe_gemm_prefill: act_dtype must be F16 or 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 1638bf7af0..0e9308c1ba 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 2e05dfa442..b1394b6091 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 fad8d87eb7..8788a4369a 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 aabe8bff05..919541d429 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 directly from the *launch +// queue's own device*. The kernel is launched on this same queue, so this +// reports the count for the device the kernel actually runs on. Reading it from +// a hardcoded ordinal 0 or from syclcompat's implicit current device 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_queue_sm_count(const sycl::queue& q) { + return q.get_device().get_info(); +} + // Command line options parsing struct Options { const void *q = nullptr, *k = nullptr, *v = nullptr; @@ -89,6 +100,11 @@ struct Options { float softmax_scale = 0.0f; float* lse = nullptr; // LSE output buffer (null = skip) bool persistent = false; + /// SYCL queue used to launch the kernel. Transient device allocations + /// (kernel workspace, zero-filled kv_cache scratch) are made on this queue so + /// they bind to its device instead of syclcompat's implicit current device. + /// Must be non-null whenever a launcher is invoked (see sdpa.cpp). + sycl::queue* queue = nullptr; void print(std::ostream& os = std::cout) const { os << std::boolalpha << "Options {\n" @@ -174,13 +190,44 @@ struct KernelRunner { StrideV stride_V_cache; StrideO stride_O; uint64_t seed = 0; - /// Scratch buffer for zero-filled kv_cache cumulative_length. - cutlass::device_memory::allocation zero_cu_cache_; + /// Scratch buffer for zero-filled kv_cache cumulative_length, allocated on the + /// launch queue's device (see make_zero_cu_cache). + int* zero_cu_cache_ = nullptr; + sycl::queue* zero_cu_cache_queue_ = nullptr; // // Methods // + KernelRunner() = default; + KernelRunner(const KernelRunner&) = delete; + KernelRunner& operator=(const KernelRunner&) = delete; + ~KernelRunner() { + if (zero_cu_cache_ != nullptr) { + sycl::free(zero_cu_cache_, *zero_cu_cache_queue_); + } + } + + // Allocate (or grow) the zero-filled kv_cache cumulative_length scratch on the + // launch queue's device and return a device pointer usable by the kernel. Using + // sycl::malloc_device on *options.queue avoids syclcompat's implicit current + // device, which can point at a different device under a full test run. + int* make_zero_cu_cache(const Options& options, int count) { + sycl::queue* q = options.queue; + if (zero_cu_cache_ != nullptr) { + sycl::free(zero_cu_cache_, *zero_cu_cache_queue_); + zero_cu_cache_ = nullptr; + } + zero_cu_cache_queue_ = q; + zero_cu_cache_ = sycl::malloc_device(count, *q); + if (zero_cu_cache_ == nullptr) { + throw std::runtime_error("KernelRunner: failed to allocate zero_cu_cache scratch (" + + std::to_string(count * sizeof(int)) + " bytes)"); + } + q->memset(zero_cu_cache_, 0, count * sizeof(int)).wait(); + return zero_cu_cache_; + } + // Note that the GemmUniversalAdapter currently doesn't support flash attention, which is why this // secondary `run` function is required to launch the kernel. static void run(typename FMHAKernel::Params params) { @@ -228,9 +275,7 @@ struct KernelRunner { if (options.cu_seqlens_kv_cache) { problem_size_for_launch.seq_len_kv_cache.cumulative_length = const_cast(options.cu_seqlens_kv_cache); } else { - zero_cu_cache_.reset(options.batch + 1); - std::fill_n(zero_cu_cache_.get(), options.batch + 1, 0); - problem_size_for_launch.seq_len_kv_cache.cumulative_length = zero_cu_cache_.get(); + problem_size_for_launch.seq_len_kv_cache.cumulative_length = make_zero_cu_cache(options, options.batch + 1); } problem_size_for_launch.head_size_qk = get<6>(problem_size); problem_size_for_launch.head_size_vo = get<7>(problem_size); @@ -288,9 +333,7 @@ struct KernelRunner { if (options.cu_seqlens_kv_cache) { shape.seq_len_kv_cache.cumulative_length = const_cast(options.cu_seqlens_kv_cache); } else { - zero_cu_cache_.reset(options.batch + 1); - std::fill_n(zero_cu_cache_.get(), options.batch + 1, 0); - shape.seq_len_kv_cache.cumulative_length = zero_cu_cache_.get(); + shape.seq_len_kv_cache.cumulative_length = make_zero_cu_cache(options, options.batch + 1); } } return shape; @@ -321,9 +364,7 @@ struct KernelRunner { options.use_paged_kv ? options.num_pages_per_seq : nullptr}, {}, hw_info}; - // Define device-global scratch memory - size_t workspace_size = FMHAKernel::get_workspace_size(arguments); - cutlass::device_memory::allocation workspace(workspace_size); + sycl::queue* q = options.queue; if (!FMHAKernel::can_implement(arguments)) { std::cout << "Invalid Problem Size: " << options.batch << 'x' << options.num_heads_q << 'x' << options.seq_len_qo << 'x' << options.seq_len_kv << 'x' << options.head_size_qk << 'x' << options.head_size_vo @@ -331,14 +372,34 @@ struct KernelRunner { return cutlass::Status::kErrorInvalidProblem; } + // Define device-global scratch memory on the same queue that launches the + // kernel so the allocation binds to that queue's device rather than + // syclcompat's implicit current device. + size_t workspace_size = FMHAKernel::get_workspace_size(arguments); + uint8_t* workspace = nullptr; + if (workspace_size > 0) { + workspace = sycl::malloc_device(workspace_size, *q); + if (workspace == nullptr) { + throw std::runtime_error("KernelRunner: failed to allocate kernel workspace (" + + std::to_string(workspace_size) + " bytes)"); + } + } + // Initialize the workspace - CUTLASS_CHECK(FMHAKernel::initialize_workspace(arguments, workspace.get())); + CUTLASS_CHECK(FMHAKernel::initialize_workspace(arguments, workspace)); // Convert host-side arguments to device-side arguments to be passed to the kernel - auto params = FMHAKernel::to_underlying_arguments(arguments, workspace.get()); + auto params = FMHAKernel::to_underlying_arguments(arguments, workspace); run(params); + // The launch above is asynchronous; wait for it on the launch queue before + // releasing the workspace the kernel is still reading from. + q->wait(); + if (workspace != nullptr) { + sycl::free(workspace, *q); + } + return cutlass::Status::kSuccess; } }; @@ -371,9 +432,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 launch queue's own device so + // the grid is sized for the device the kernel actually runs on (see query_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_queue_sm_count(*options.queue); using ProblemShapeType = cutlass::fmha::kernel::FMHAProblemShape; @@ -531,13 +593,44 @@ struct SageKernelRunner { StrideV stride_V_cache; StrideO stride_O; uint64_t seed = 0; - /// Scratch buffer for zero-filled kv_cache cumulative_length. - cutlass::device_memory::allocation zero_cu_cache_; + /// Scratch buffer for zero-filled kv_cache cumulative_length, allocated on the + /// launch queue's device (see make_zero_cu_cache). + int* zero_cu_cache_ = nullptr; + sycl::queue* zero_cu_cache_queue_ = nullptr; // // Methods // + SageKernelRunner() = default; + SageKernelRunner(const SageKernelRunner&) = delete; + SageKernelRunner& operator=(const SageKernelRunner&) = delete; + ~SageKernelRunner() { + if (zero_cu_cache_ != nullptr) { + sycl::free(zero_cu_cache_, *zero_cu_cache_queue_); + } + } + + // Allocate (or grow) the zero-filled kv_cache cumulative_length scratch on the + // launch queue's device and return a device pointer usable by the kernel. Using + // sycl::malloc_device on *options.queue avoids syclcompat's implicit current + // device, which can point at a different device under a full test run. + int* make_zero_cu_cache(const Options& options, int count) { + sycl::queue* q = options.queue; + if (zero_cu_cache_ != nullptr) { + sycl::free(zero_cu_cache_, *zero_cu_cache_queue_); + zero_cu_cache_ = nullptr; + } + zero_cu_cache_queue_ = q; + zero_cu_cache_ = sycl::malloc_device(count, *q); + if (zero_cu_cache_ == nullptr) { + throw std::runtime_error("KernelRunner: failed to allocate zero_cu_cache scratch (" + + std::to_string(count * sizeof(int)) + " bytes)"); + } + q->memset(zero_cu_cache_, 0, count * sizeof(int)).wait(); + return zero_cu_cache_; + } + // Note that the GemmUniversalAdapter currently doesn't support flash attention, which is why this // secondary `run` function is required to launch the kernel. static void run(typename FMHAKernel::Params params) { @@ -585,9 +678,7 @@ struct SageKernelRunner { if (options.cu_seqlens_kv_cache) { problem_size_for_launch.seq_len_kv_cache.cumulative_length = const_cast(options.cu_seqlens_kv_cache); } else { - zero_cu_cache_.reset(options.batch + 1); - std::fill_n(zero_cu_cache_.get(), options.batch + 1, 0); - problem_size_for_launch.seq_len_kv_cache.cumulative_length = zero_cu_cache_.get(); + problem_size_for_launch.seq_len_kv_cache.cumulative_length = make_zero_cu_cache(options, options.batch + 1); } problem_size_for_launch.head_size_qk = get<6>(problem_size); problem_size_for_launch.head_size_vo = get<7>(problem_size); @@ -645,9 +736,7 @@ struct SageKernelRunner { if (options.cu_seqlens_kv_cache) { shape.seq_len_kv_cache.cumulative_length = const_cast(options.cu_seqlens_kv_cache); } else { - zero_cu_cache_.reset(options.batch + 1); - std::fill_n(zero_cu_cache_.get(), options.batch + 1, 0); - shape.seq_len_kv_cache.cumulative_length = zero_cu_cache_.get(); + shape.seq_len_kv_cache.cumulative_length = make_zero_cu_cache(options, options.batch + 1); } } return shape; @@ -680,9 +769,7 @@ struct SageKernelRunner { options.use_paged_kv ? options.num_pages_per_seq : nullptr}, {}, hw_info}; - // Define device-global scratch memory - size_t workspace_size = FMHAKernel::get_workspace_size(arguments); - cutlass::device_memory::allocation workspace(workspace_size); + sycl::queue* q = options.queue; if (!FMHAKernel::can_implement(arguments)) { std::cout << "Invalid Problem Size: " << options.batch << 'x' << options.num_heads_q << 'x' << options.seq_len_qo << 'x' << options.seq_len_kv << 'x' << options.head_size_qk << 'x' << options.head_size_vo @@ -690,14 +777,34 @@ struct SageKernelRunner { return cutlass::Status::kErrorInvalidProblem; } + // Define device-global scratch memory on the same queue that launches the + // kernel so the allocation binds to that queue's device rather than + // syclcompat's implicit current device. + size_t workspace_size = FMHAKernel::get_workspace_size(arguments); + uint8_t* workspace = nullptr; + if (workspace_size > 0) { + workspace = sycl::malloc_device(workspace_size, *q); + if (workspace == nullptr) { + throw std::runtime_error("KernelRunner: failed to allocate kernel workspace (" + + std::to_string(workspace_size) + " bytes)"); + } + } + // Initialize the workspace - CUTLASS_CHECK(FMHAKernel::initialize_workspace(arguments, workspace.get())); + CUTLASS_CHECK(FMHAKernel::initialize_workspace(arguments, workspace)); // Convert host-side arguments to device-side arguments to be passed to the kernel - auto params = FMHAKernel::to_underlying_arguments(arguments, workspace.get()); + auto params = FMHAKernel::to_underlying_arguments(arguments, workspace); run(params); + // The launch above is asynchronous; wait for it on the launch queue before + // releasing the workspace the kernel is still reading from. + q->wait(); + if (workspace != nullptr) { + sycl::free(workspace, *q); + } + return cutlass::Status::kSuccess; } }; @@ -733,9 +840,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 launch queue's own device so + // the grid is sized for the device the kernel actually runs on (see query_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_queue_sm_count(*options.queue); 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 ff97deaa13..779ee00619 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 732cfec70c..2102aa1a67 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.py b/auto_round_extension/ark/test/test_moe.py index 8bad20cf98..ff567d31a6 100644 --- a/auto_round_extension/ark/test/test_moe.py +++ b/auto_round_extension/ark/test/test_moe.py @@ -408,7 +408,10 @@ def _moe_decode_reference(activations, dequant_weights, num_tokens_per_expert): """Reference: each token is matmul'd against its routed expert's weights.""" total_tokens, K = activations.shape E, N, _ = dequant_weights.shape - out = torch.empty(total_tokens, N, dtype=activations.dtype, device=activations.device) + # Zero-initialise (not ``torch.empty``) so any row the per-expert loop does + # not cover (zero-token experts, or a short ``num_tokens_per_expert`` sum) + # reads back as deterministic zeros instead of stale allocator memory. + out = torch.zeros(total_tokens, N, dtype=activations.dtype, device=activations.device) offset = 0 for e in range(E): n_tokens = int(num_tokens_per_expert[e].item()) @@ -739,7 +742,9 @@ def _run_prefill_reference(activations, weights_NK, num_tokens_per_expert): """Reference: per-expert ``A @ W.T`` over ``[E, N, K]`` weights.""" total_tokens, _ = activations.shape E, N, _ = weights_NK.shape - out = torch.empty(total_tokens, N, dtype=activations.dtype, device=activations.device) + # Zero-initialise (not ``torch.empty``) so any uncovered row reads back + # as deterministic zeros rather than stale allocator memory. + out = torch.zeros(total_tokens, N, dtype=activations.dtype, device=activations.device) offset = 0 for e in range(E): n_tokens = int(num_tokens_per_expert[e].item()) 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 c8c1ee168c..7b921ad8ab 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_accuracy.py +++ b/auto_round_extension/ark/test/test_moe_prefill_accuracy.py @@ -151,10 +151,23 @@ def _reference_moe_prefill(activations, dequant_weights_NK, num_tokens_per_exper This is the same reference used by ``TestMoEGemmPrefill`` in ``test_moe.py`` and matches what a model would compute when no fused kernel is available. + + The matmuls are executed on CPU (inputs are moved off the accelerator) so + the reference is independent of the device kernels under test; the result is + moved back to the original device so it stays comparable to the kernel output. """ + orig_device = activations.device + activations = activations.cpu() + dequant_weights_NK = dequant_weights_NK.cpu() + num_tokens_per_expert = num_tokens_per_expert.cpu() total_tokens, _ = activations.shape E, N, _ = dequant_weights_NK.shape - out = torch.empty(total_tokens, N, dtype=activations.dtype, device=activations.device) + # Zero-initialise (not ``torch.empty``): any row not covered by the + # per-expert loop below -- e.g. a zero-token expert, or a + # ``num_tokens_per_expert`` sum that is short of ``total_tokens`` -- would + # otherwise expose stale allocator memory (frequently read back as + # all-zeros), silently corrupting the reference the kernel is compared to. + out = torch.zeros(total_tokens, N, dtype=activations.dtype, device="cpu") offset = 0 for e in range(E): n_tokens = int(num_tokens_per_expert[e].item()) @@ -163,7 +176,7 @@ def _reference_moe_prefill(activations, dequant_weights_NK, num_tokens_per_exper a = activations[offset : offset + n_tokens] out[offset : offset + n_tokens] = a @ dequant_weights_NK[e].T offset += n_tokens - return out + return out.to(orig_device) # --------------------------------------------------------------------------- @@ -186,19 +199,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 # --------------------------------------------------------------------------- @@ -226,6 +285,77 @@ def _fix_seed(self): if hasattr(torch, "xpu") and torch.xpu.is_available(): torch.xpu.manual_seed_all(0) + @pytest.fixture(autouse=True) + def _isolate_ark_moe_env(self): + """Give every test a clean ``ARK_MOE_PREFILL_*`` env baseline. + + Several tests (here and in ``test_moe_prefill_perf.py``) opt into an + experimental dispatch path by exporting an ``ARK_MOE_PREFILL_*`` flag + -- some of which route through kernels with known numeric defects + (e.g. the S4 DPAS large-M outliers). If such a flag leaks -- because a + preceding test aborted before restoring it, or it was exported in the + shell -- the *default-path* parity tests below silently run the wrong + kernel and fail. That is the classic "passes in isolation, fails only + in the full suite" symptom. + + Snapshot every ``ARK_MOE_PREFILL_*`` variable, clear them so each test + starts from the documented default (all opt-in paths OFF), then + restore the original process environment afterwards. Tests that need a + flag set it explicitly within their own body, so clearing the baseline + never changes their behaviour. + """ + prefix = "ARK_MOE_PREFILL_" + saved = {k: v for k, v in os.environ.items() if k.startswith(prefix)} + for k in saved: + del os.environ[k] + try: + yield + finally: + for k in [k for k in os.environ if k.startswith(prefix)]: + del os.environ[k] + os.environ.update(saved) + + @pytest.fixture(autouse=True) + def _xpu_cleanup_between_tests(self): + """Release the XPU allocator cache before and after every test. + + The quantized prefill paths allocate ``torch.empty`` buffers -- the + ``[total_tokens, N]`` output and (for the dequant path) a large + ``[E, K, N]`` scratch workspace -- and the grouped GEMM only writes the + token/expert rows it is actually dispatched for. Any element it does not + overwrite reads back whatever bytes the caching allocator handed out. + + Run in isolation the accumulator starts empty, so those buffers come from + fresh device pages that read back as zeros and the comparison passes. + Run as part of the full suite, the preceding XPU test files + (``test_matmul``, ``test_moe``, ``test_moe_decode_perf``, ...) leave the + allocator holding a large, *non-zero* working set; the next ``torch.empty`` + here reuses one of those dirty blocks and the stale bytes surface as a + handful of wildly-wrong outputs. That is the classic "passes in isolation, + fails only in the full suite" symptom (observed on + ``test_accuracy_int4``). + + Every other XPU test file in this directory already brackets its tests + with ``torch.xpu.empty_cache()`` for exactly this reason (see + ``_xpu_cleanup_between_tests`` in ``test_moe_decode_perf.py`` and the + per-test ``empty_cache`` calls in ``test_matmul.py`` / + ``test_weightonly.py``). Mirror that here so each accuracy test starts + from a clean allocator state regardless of what ran before it. + """ + self._release_xpu_memory() + try: + yield + finally: + self._release_xpu_memory() + + @staticmethod + def _release_xpu_memory(): + """Synchronize and free cached XPU memory (no-op when XPU is absent).""" + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.synchronize() + if hasattr(torch.xpu, "empty_cache"): + torch.xpu.empty_cache() + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_accuracy_fp(self, dtype): for label, E, tpe, N, K in PREFILL_SHAPES: @@ -243,7 +373,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 +410,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 +447,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 +484,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 +531,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 +590,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 +642,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 +684,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 +734,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 +748,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 +788,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 c322cb66bf..2b856190d9 100644 --- a/auto_round_extension/ark/test/test_moe_prefill_perf.py +++ b/auto_round_extension/ark/test/test_moe_prefill_perf.py @@ -489,6 +489,34 @@ def _xpu_cleanup_between_tests(): _release_xpu_memory() +@pytest.fixture(autouse=True) +def _isolate_ark_moe_env(): + """Guarantee ``ARK_MOE_PREFILL_*`` flags never leak out of a perf test. + + The perf sweeps toggle ``ARK_MOE_PREFILL_*`` flags per shape to attribute + each timing column to a specific dispatch path, restoring the prior value + only after the last measurement in the loop body. If a measurement raises + (OOM, kernel error, assertion) the manual restore is skipped and the flag + -- e.g. ``ARK_MOE_PREFILL_DPAS_S4=1``, which routes int4-sym through a + kernel with a known large-M defect -- leaks into subsequent tests. Those + tests then fail only when run after this one (i.e. only in the full + suite), while passing in isolation. + + Snapshot the flags at setup and restore them unconditionally at teardown so + a leak can never escape a single test, regardless of exceptions. + """ + prefix = "ARK_MOE_PREFILL_" + saved = {k: v for k, v in os.environ.items() if k.startswith(prefix)} + for k in saved: + del os.environ[k] + try: + yield + finally: + for k in [k for k in os.environ if k.startswith(prefix)]: + del os.environ[k] + os.environ.update(saved) + + def _print_header(title: str) -> None: """Print a benchmark header. @@ -668,15 +696,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 1bb587b64f..2c1e751862 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -# 1.5.1