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.

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

W4A8 perf [prefill] (models=qwen3+minimax, group_size=32, act=bfloat16, rescale_group_size=-1) -- ark.moe_gemm_w4a8 vs W4A16 vs torch
shape             E      N      K  tokens  rows/E   torch(ms)   w4a16(ms)    w4a8(ms)    TFLOPS    W GB/s  DRAM GB/s   BW@100T   vs torch   vs w4a16  prepack(ms)
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
qwen3 up        128   1536   2048   65536   512.0          --       3.760       2.856    144.36     141.0      399.4       277         --      1.32x         2.50
qwen3 down      128   2048    768   65536   512.0          --       2.561       2.098     98.25      95.9      319.8       326         --      1.22x         1.14
minimax up      192   1536   3072   65536   341.3          --       6.649       4.565    135.48     198.5      419.0       309         --      1.46x         4.75
minimax down    192   3072   1536   65536   341.3          --       6.529       4.263    145.07     212.5      401.4       277         --      1.53x         4.85

[a32543254]

W4A8 perf [prefill] (models=qwen3+minimax, group_size=32, act=bfloat16, rescale_group_size=-1) -- ark.moe_gemm_w4a8 vs W4A16 vs torch
shape             E      N      K  tokens  rows/E   torch(ms)   w4a16(ms)    w4a8(ms)    TFLOPS    W GB/s  DRAM GB/s   BW@100T   vs torch   vs w4a16  prepack(ms)
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
qwen3 up        128   1536   2048   65536   512.0          --       3.760       2.856    144.36     141.0      399.4       277         --      1.32x         2.50
qwen3 down      128   2048    768   65536   512.0          --       2.561       2.098     98.25      95.9      319.8       326         --      1.22x         1.14
minimax up      192   1536   3072   65536   341.3          --       6.649       4.565    135.48     198.5      419.0       309         --      1.46x         4.75
minimax down    192   3072   1536   65536   341.3          --       6.529       4.263    145.07     212.5      401.4       277         --      1.53x         4.85

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[a32543254]

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.