[BugFix][MoE] Make xe2 grouped-GEMM tile scheduler cudagraph-capture-safe

[AleksVidmantas]

Purpose

On Intel Arc (Xe2/BMG) the xe2 grouped GEMM (MoE::MoEGEMM,
csrc/xpu/grouped_gemm/xe_2/grouped_gemm_xe2.hpp) cannot be captured into a SYCL graph
(cudagraph) at batch > 1: replaying a captured decode graph with more than one sequence raises
UR_RESULT_ERROR_DEVICE_LOST. Servers using these kernels must therefore pin
cudagraph_capture_sizes to [1] and fall back to eager execution for any concurrent batch,
losing the graph-capture speedup exactly where batching helps most.

Root cause. The kernel uses a persistent work-steal scheduler with a single global atomic
tile counter that it resets from inside the kernel (the atomic_buffer[0] store), guarded only
by "group 0, lane 0" with no device-wide barrier. Workgroups then steal the next tile via
atomicAdd(atomic_buffer, 1) (workgroup-local fence only).

• In eager submits the in-kernel reset is separated from the steals by launch timing, so it
  happens to be correct (this is why the eager batched path works).
• Under SYCL-graph capture/replay that temporal separation is not reproduced: the counter
  carries the previous replay's final value, and the un-ordered reset-vs-steal ordering yields
  duplicated / inconsistent tile indices. A workgroup then computes an m_coord outside an
  expert's real row range → out-of-bounds global load/store → GPU page fault → DEVICE_LOST.
• The hazard scales with the number of tiles, i.e. with batch: at batch 1 the tile count is often
  within the initial grid, so the steal loop barely runs and capture at [1] is safe; at
  batch > 1 the steal loop is heavily exercised and replay corrupts. (Even at batch 1 the
  unsynchronized reset-vs-steal is a latent race rather than a guaranteed-safe design; the
  deterministic scheduler removes the hazard at every batch size.)

Fix. Replace the atomic work-steal with a deterministic grid-stride partition: workgroup
g processes tiles g, g + group_range, g + 2*group_range, .... There is no shared mutable global
state and no in-kernel reset, so the kernel is safe to capture and replay at any batch size.
Coverage is identical to the previous scheduler (every tile computed exactly once), so
single-stream output and performance are unchanged. The launch grid was already occupancy-based and
batch-independent, so no launch change is needed. atomic_buffer / the SLM scratch are left in the
signature (unused, (void)-cast) to keep the change minimal; removing them from the op/launcher
signature is a mechanical follow-up.

Test Plan

Build (Intel Arc Pro B70, Xe2/BMG-G31, oneAPI 2025.3):

source /opt/intel/oneapi/setvars.sh
CMAKE_BUILD_TYPE=Release MAX_JOBS=4 pip install --no-build-isolation -e .

1. Kernel unit test — grouped GEMM vs per-expert reference, M ∈ {1, 4, 16, 8192}:

   python -m pytest tests/fused_moe/test_grouped_gemm.py
   

2. End-to-end on a served MoE model (Qwen3.6-35B-A3B W4A16, int4 grouped MoE) under vLLM-XPU
   PIECEWISE capture, comparing the pinned-[1] baseline against cudagraph_capture_sizes=[1,2,4],
   max_num_seqs=4. (E2E uses a model-specific serving build; the unit test above needs only the
   standard build shown here.)

Test Result

Unit test (fixed kernel): 144 passed, 16 skipped (the 16 skips are the pre-existing
test_grouped_gemm non-xe2 case, skipif reason "Need Fix API" — unrelated to this change). The
path this PR touches, test_xe_grouped_gemm_int4, is 16/16 passed across
M ∈ {1, 4, 16, 8192} × {bf16, fp16} × {bias, no-bias}, within rtol 1e-2. Sibling xe2 paths
(_fp8, _block_fp8_w8a8, _mxfp4) also pass, confirming the scheduler rewrite did not disturb
the shared compute paths.

End-to-end (Qwen3.6-35B-A3B W4A16, max-model-len 131072, max-num-seqs 4,
cudagraph_capture_sizes=[1,2,4], non-spec):

• No regression (batch 1): capture [1] boots clean, output coherent; warm single-stream decode
  is unchanged vs. the pre-fix kernel (~89 t/s, short context). The rewrite has identical tile
  coverage, so batch-1 output and rate are untouched.

• Capture unblocked (the fix): capture of all three batch sizes completes with zero
  DEVICE_LOST; 4 concurrent prompts return coherent, correct answers (previously: DEVICE_LOST at
  batch > 1).

• Throughput scaling (captured batch-4 vs single-stream), across context. Aggregate decode
  measured from the vLLM generation_tokens_total counter delta over the window where all 4 streams
  are co-resident and decoding (median of 3 reps, ± s.d.):

  combined ctx	batch-4 agg t/s	batch-4 / single-stream	GB/s	MBU (of 608)
  8k	269.1 ±0.4	3.13×	312	51.3%
  20k	261.1 ±2.7	3.01×	319	52.5%
  40k	255.4 ±2.4	3.05×	338	55.6%
  80k	230.5 ±2.2	2.87×	353	58.0%
  120k	213.5 ±1.6	2.77×	370	60.9%

  Combined ctx = total tokens across the 4 co-resident streams (per-stream ≈ ÷4); single-stream
  (batch-1) is 77–86 t/s across the same per-stream context range. The ~3× captured-batch-4 speedup
  holds across the full 8k→120k range — capture removes the per-step launch overhead that
  otherwise dominates batched decode — lifting memory-bandwidth utilization from ~50% (single-stream)
  toward ~61% at batch-4.

The change is a scheduling rewrite with identical tile coverage; it does not alter numerics.
Batch-1 greedy output is unchanged (bit-for-bit tile coverage), and batch-4 concurrent greedy
prompts produce correct answers. No accuracy/eval regression expected or observed.

(Optional) Documentation Update

None — kernel-internal scheduling fix; no user-facing API, configuration, or documentation change.

Risk / limitations

• Grid-stride is a static partition; it drops the previous scheduler's dynamic
  load-balancing across experts with very uneven token counts. For highly skewed expert loads a
  few % of tail imbalance is possible; in the decode regime measured here it was not material, and
  the capture win dominates. A hybrid (static assignment + a small bounded steal reset out-of-kernel
  with a memset graph node) could recover it in a follow-up if needed.
• atomic_buffer and the SLM scratch are now unused in this kernel; removing them from the
  op/launcher signature is a mechanical cleanup left out to keep the diff focused.

Related

Independent of open PR #436 (block-FP8 W8A8 grouped-GEMM), which edits this same file: this fix was
developed alongside #436 but does not depend on it and applies cleanly with or without it — only a
textual (not functional) conflict is possible if #436 lands first, so happy to rebase in whichever
order suits. Prior art: #320 (same PIECEWISE capture-padding family) and #276 (torch.compile
umbrella).