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

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

[a32543254]

@copilot resolve the merge conflicts in this pull request

[Copilot]

@copilot resolve the merge conflicts in this pull request

Merged in 1f7c408. Branch is now 0 commits behind main, and GitHub moved the PR from dirty to blocked (i.e. conflicts gone, just awaiting review).

Only three files actually overlapped with main, and git flagged just one — but that one was a real collision, not a textual one.

The conflict: two kernels claimed the same scratch slot. main added the XMX Hadamard kernel (ecc5b00), which took DeviceMemoryPool slot 11 for its staging buffer. This branch had already taken slots 11 and 12 for the W4A8 activation-quantization and expert-map buffers. Since get_scratch_ptr keys slabs by (buf_loc, device_key), leaving both at 11 would have made the two kernels share one slab — and because the sizes differ, each call would free and reallocate the other's buffer out from under it. Silent corruption, not a compile error.

Resolved by leaving main's slot 11 alone and moving W4A8 to 12/13, with MaxLocNum 12/13 → 14 so the std::array slot maps and the three buf_loc >= MaxLocNum guards actually cover them. Final map: 8 DPAS counter, 9/10 int4 decode, 11 Hadamard staging, 12/13 W4A8.

The two auto-merges were verified, not assumed — ark.cpp and __init__.py merged without markers, which is exactly how the semantic breakage in the previous merge slipped through. Checked: no m.def name registered twice (55 total, both sides' registrations present), brace/paren/bracket balance clean, __init__.py parses and keeps both sides' APIs, and main's from .mxfp4_hadamard import ... still sits after the lib binding and before __main__, as its "imported last" comment requires. A line-level diff against both parents shows zero added lines dropped from either side.

CI on the merge commit is green so far — 18 checks passed (DCO, CodeQL, all 12 compatibility matrix jobs, both ark-change detectors), 0 failures. The ARK wheel build and the unit-test/XPU jobs are still running; the wheel build is the one that would catch a compile regression from this resolution, since there's no SYCL compiler in my environment.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

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

[Copilot]

🟡 Changes recommended

One or more issues must be addressed before approval.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds an XPU ARK W4A8 MoE path with int4→int8 prepacking, dynamic activation quantization, prefill DPAS GEMM, decode GEMV, caching, and fused reduction.

Changes:

• Added W4A8 kernels, scratch management, and generated translation units.
• Added C++/Python APIs and prepack caching.
• Added performance/accuracy tests and bilingual documentation.

File summaries

File	Description
auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_w4a8.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_w4a8_scratch.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_w4a8_kernels.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_w4a8_helpers.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_moe_prefill_fp8_dpas.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/sycl_tla_moe.cpp.in	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/sycl_tla_moe_w4a8_scratch.cpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/sycl_tla_generation.cmake	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/ark.cpp	Updated as part of this pull request.
auto_round_extension/ark/auto_round_kernel/init.py	Updated as part of this pull request.

Review details

Suppressed comments (3)

auto_round_extension/ark/auto_round_kernel/init.py:3918

• This cache key uses the storage addresses and shape but not tensor versions or strides (and not the effective block after ARK_MOE_W4A8_AUTO_S8 is applied). If packed weights/scales are updated in place, or a different view reuses the same address/shape, the cache silently returns an old int8 prepack and produces results for stale weights. Track mutation/layout/effective-block identity or invalidate the entry when inputs change.

        key = (
            device.type,
            device.index,
            weights.data_ptr(),
            scales.data_ptr(),
            tuple(weights.shape),
            int(group_size),
            int(rescale_group_size),
            str(scales.dtype),
        )

auto_round_extension/ark/auto_round_kernel/init.py:3580

• The public API returns three values here—weights_s8, wscales, and the resolved block size—but the PR's usage example unpacks only two (w_s8, w_scale = ...). That example will raise ValueError: too many values to unpack; update the example to include block (which is required by moe_gemm_w4a8) or change the API consistently.

    Returns:
        ``(weights_s8, wscales, rescale_block_size)`` where ``weights_s8`` is
        ``[E, N, K]`` ``torch.int8``, ``wscales`` is
        ``[E, N, K // rescale_block_size]`` ``torch.float32``, and
        ``rescale_block_size`` is the resolved block size (pass it to
        :func:`moe_gemm_w4a8`).

auto_round_extension/ark/auto_round_kernel/sycl_tla_moe_w4a8_scratch.cpp:55

• DeviceMemoryPool::get_scratch_ptr explicitly avoids q->wait() during XPU graph capture because that host wait is illegal (utils.hpp:186-200), but this new growth path performs the wait itself. A captured W4A8 call that grows the activation or expert-map slab will therefore fail; use the pool's nonblocking growth/lifetime mechanism or explicitly reject capture before reaching this path.

void* acquire_w4a8_slab(sycl::queue* q, size_t bytes, size_t buf_loc) {
  auto* pool = DeviceMemoryPool::Instance();
  const size_t held = pool->get_scratch_size(buf_loc, q);
  if (held != 0 && held < bytes) {
    q->wait();
  }
  void* ptr = pool->get_scratch_mem(bytes, buf_loc, q);

• Files reviewed: 12/15 changed files
• Comments generated: 5
• Review effort level: Lite

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[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.

[a32543254]

@copilot resolve the merge conflicts in this pull request

[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.