test: stabilize flaky int4-sym fp16 MoE prefill accuracy test

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
 
 
 # ---------------------------------------------------------------------------

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<int*>(
@@ -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,

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<max_compute_units>()`, so this is numerically
+  // identical for the correct device.)
+  int sm_count = q->get_device().get_info<sycl::info::device::max_compute_units>();
   cutlass::KernelHardwareInfo hw_info{0, sm_count};
 
   auto dummy_problem_shape = cute::Shape<int, int, int>{1, gemm_k, gemm_n};

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<ElementA, int4b_t, N>` 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
-  // <ElementA, cutlass::int4b_t, N>` 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<const uint8_t*>(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<ScalarT>(
@@ -1083,12 +1107,26 @@ inline void moe_gemm_prefill(sycl::queue* q, void* activations, void* weights, v
     auto* w_kn = static_cast<sycl::half*>(dequant_workspace);
     moe_mixed_detail::dequant_to_KN<sycl::half>(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<BF*>(dequant_workspace);
     moe_mixed_detail::dequant_to_KN<BF>(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");