Skip to content

Optimize INT4 MoE prefill/decode (sym) XPU kernels with dedicated w4a16 tiling - #2111

Draft
a32543254 with Copilot wants to merge 49 commits into
mainfrom
copilot/optimize-int4-moe-performance
Draft

Optimize INT4 MoE prefill/decode (sym) XPU kernels with dedicated w4a16 tiling#2111
a32543254 with Copilot wants to merge 49 commits into
mainfrom
copilot/optimize-int4-moe-performance

Conversation

Copilot AI commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Description

The S4-sym prefill DPAS path reused the INT8 dpas_w8a16_policy* tile shapes, leaving 4-bit-specific throughput on the table vs. the reference w4a16 grouped GEMM in vllm-project/vllm-xpu-kernels. This tunes the Intel XPU (BMG/Xe2) INT4-sym MoE kernels in auto_round_extension/ark. All changes are runtime-gated (ARK_MOE_PREFILL_DPAS_S4) and math-preserving — tile shapes only affect on-device blocking, not results.

Prefill (sym)sycl_tla_moe_prefill_s4_dpas.hpp, sycl_tla_moe_prefill_fp8_dpas.hpp

  • Added dedicated 4-bit tile policies: dpas_w4a16_policy (WG 128×256×32, SG 4×8) for large M — the halved packed-nibble B stream makes the wider N tile pay off; dpas_w4a16_policy_m_8 (WG 8×64×32) for tiny M. m_16/m_32 alias the shared 64-wide-N INT8 shapes.
  • Switched dispatch to the reference 4-tier A_avg_M bucketing (≤4→m_8, ≤8→m_16, ≤128→m_32, else default). The 32×64 tile now covers M up to 128 instead of jumping to the wide tile at 33, cutting padding waste on chunked-prefill batches.

Decode (sym)sycl_tla_moe_decode.hpp

  • Refactored the launch_int4 GEMV inner loop into a width-templated int4_decode_chunk helper with a 32→16→scalar ladder. Widens the fast path for shipped group sizes (32/64/128/256) while never regressing group_size==16.

DocsREADME_MOE_PREFILL_PERF.md + _CN counterpart updated with the S4 tile-policy table and bucketing.

These are SYCL/CuTe header kernels requiring an on-hardware build (oneAPI + cutlass-sycl, Intel BMG) to benchmark speedups and confirm accuracy parity (test_perf_int4, test_moe_decode_perf.py, test_accuracy_int4_dpas_per_group).

Type of Change

Performance

Related Issues

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.

Note: these are Intel XPU (BMG/Xe2) SYCL kernels; local build/benchmark requires oneAPI + cutlass-sycl on XPU hardware, which is unavailable in this environment. Please validate on target hardware.

a32543254 and others added 3 commits July 31, 2026 11:16
Signed-off-by: Dong, Bo1 <bo1.dong@intel.com>
Signed-off-by: Dong, Bo1 <bo1.dong@intel.com>
… GEMV

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Copilot AI and others added 24 commits August 5, 2026 08:59
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…mangled names

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…atch regression

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Replace per-call sycl::malloc_device/free of the 1-int32 work-group
counter in the prefill and decode S4 DPAS grouped-GEMM dispatchers with
a persistent per-queue buffer. Each malloc/free forces a queue sync,
which is pure overhead on the decode hot path. The kernel self-zeroes
the counter at launch and every launch is synchronous, so a single
reused buffer per queue is safe.

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Two waste removals on the int4-sym decode path, motivated by the
device-event perf measurement (host-side malloc/free reuse showed no
change because it never lands on the device timeline):

- Skip the fill_expert_id_per_token kernel when the S4-sym DPAS fast
  path is taken; that path uses num_tokens_per_expert directly and
  never reads expert_id_per_token_buf, so the extra device kernel
  launch was pure overhead. The predicate is shared with the S4 block
  so the two decisions cannot diverge.
- Cache the per-launch query_device_multiprocessor_count host driver
  query in MoEGEMMLauncher_s4 (device SM count is fixed), removing a
  Level Zero round-trip from every decode launch.

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…m-xpu-kernels)

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…x88 sign-flip)

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…ool the repack buffer

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Decode packed int4 nibbles from 32-bit words instead of 8-bit-typed
sycl::vec ops, and restore true signed-nibble decode for sym so the
zero-point/activation-sum fold (and its extra per-call device kernel
launch) is skipped entirely on the sym path.

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Rewrite the sym-decode, coalesced-repack and activation-sum sections of
README_MOE_PREFILL_PERF.md and its CN counterpart to describe the 32-bit
word decode path and the asym-only activation-sum pre-pass.

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…ection (EN+CN)

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…e allocation-free

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
…ng sync

Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants