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

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)) {

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<int, sycl::memory_order::relaxed,
-                                sycl::memory_scope::device,
-                                sycl::access::address_space::global_space>(
-        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;
@@ -886,7 +885,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<int32_t, 1> local_mem(sycl::range<1>(1), cgh);
     cgh.parallel_for<DpasGemmName<ElementA, ElementB, ElementS, ElementD,
                                   layoutA, layoutB, policy, Mode>>(

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<int, sycl::memory_order::relaxed,
-                                sycl::memory_scope::device,
-                                sycl::access::address_space::global_space>(
-        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;
@@ -737,7 +736,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<int32_t, 1> local_mem(sycl::range<1>(1), cgh);
     cgh.parallel_for<DpasGemmIntName<ElementA, ElementB, ElementS, ElementD,
                                      layoutA, layoutB, policy, Mode>>(
@@ -884,6 +890,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<char>(std::tolower(static_cast<unsigned char>(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