feat: W4A8 ARK XPU MoE kernel (int4 weight / int8 compute) with prefill + decode

[Copilot]

Description

Adds a W4A8 MoE kernel to the ARK XPU backend covering both prefill and decode: int4 symmetric weights, int8 compute dtype, activations dynamically quantized per token to int8. Follows the W4A8 weight-only GEMM in zhenzhong/woqgemm_s8_update, and uses ARK's AUTO_S8 trick — re-scaling int4 group=32 weights into int8 group=-1 so the K loop needs a single full-width int32 accumulation instead of per-group folding.

Built on top of copilot/optimize-int4-moe-performance.

Numerics

AUTO_S8 re-scale, ported from packscale/unpackq in xpu_wrapper.hpp:

sxt[e][n][j] = max_{g in block j} |s[e][n][g]| * 8 / 127
w8           = round(w4 * s / sxt)

Default block = K (group=-1) → blks == 1. Since _pack_int4_sym divides by 7.0, max|w8| = 7 * 127/8 ≈ 111 ≤ 127, so the conversion never clips. Epilogue is out = acc_s32 * scale_b[col] * scale_a[row].

Kernel — wrapper/include/sycl_tla_moe_w4a8.hpp (new)

• One-shot moe_w4a8_prepack: [E,N,K/2] int4 → [E,N,K] int8 + [E,N,blks] fp32 scales. Costs E*N*K bytes (2× the packed int4).
• Per-token activation quant producing x_s8 and scale_a.
• Prefill: grouped DPAS GEMM on XE_DPAS_TT<8, int32_t, int8_t, int8_t>, tile ladder mirroring the reference (m<16 → 8x128, m<128 → 64x128, m<=1024 → 128x128, else 256x128). num_tokens_per_expert is a device tensor, so it reuses the existing persistent work-stealing scheduler.
• Decode: GEMV, SG_SIZE=16 / N_TILE=16, one output column per lane.
• ARK_MOE_W4A8_AUTO_S8 overrides the rescale block size; invalid values fall back to K. Shape gate: N % 16 == 0, K % 64 == 0, group_size % 8 == 0, K % group_size == 0.

Plumbing

• sycl_tla_common.hpp: 4 public declarations.
• ark.cpp: include, 2 wrappers, 4 m.def registrations.
• auto_round_kernel/__init__.py: Python API plus a prepack cache. The cache key includes device type/index and the entry pins the source tensors, so a freed-and-reallocated weight buffer can't alias another layer's int8 weights.

from auto_round_extension.ark import auto_round_kernel as ark

out = ark.moe_w4a8(x, weights, scales, num_tokens_per_expert, group_size=32, phase="prefill")

# or manage the prepack lifetime explicitly
w_s8, w_scale = ark.moe_w4a8_prepack(weights, scales, group_size=32)
out = ark.moe_gemm_w4a8(x, w_s8, w_scale, num_tokens_per_expert, phase="decode")

Benchmark — test/test_moe_w4a8_perf.py (new)

Standalone perf + accuracy harness, runnable under pytest or directly. Qwen3-MoE shapes (E=128, hidden=2048, inter=768, top_k=8, group_size=32). Reports SNR/cosine/max-rel-err against a torch bf16 baseline and latency/TFLOPS/speedup against the W4A16 kernel, for prefill across batch×seq and decode across batch.

pytest auto_round_extension/ark/test/test_moe_w4a8_perf.py -v
python auto_round_extension/ark/test/test_moe_w4a8_perf.py --warmup 10 --iters 50

Type of Change

New feature (Performance)

Checklist Before Submitting

• My code has been tested locally.
• Documentation has been updated as needed.
• New or updated tests are included where applicable.
• The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

Important

Needs hardware validation. No XPU or SYCL compiler was available, so the kernel is unbuilt and has not run on device; the header carries the same STATUS: NEEDS-HARDWARE-VALIDATION marker as the sibling MoE DPAS headers. What was checked offline: a numpy replay of the full AUTO_S8 + int8-activation path (no clipping, ~38.5 dB end-to-end output SNR vs. the script's 20 dB gate), Python↔C++ parity of moe_w4a8_rescale_block_size across 11 cases, pybind argument order against the Python call sites, and offset overflow at E=128/N=2048/K=2048.

Worth a close look on hardware: CuTe tile/policy tuning, and the 128-token cutoff for decode auto-dispatch.

Docs: test/README_MOE_W4A8.md + README_MOE_W4A8_CN.md.