Optimize INT4 MoE prefill/decode (sym) XPU kernels with dedicated w4a16 tiling - #2111
Conversation
Signed-off-by: Dong, Bo1 <bo1.dong@intel.com>
Signed-off-by: Dong, Bo1 <bo1.dong@intel.com>
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
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. |
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
Move the int4 repack and activation-sum scratch buffers off the header-local DeviceScratchPool (keyed by raw sycl::queue*) and onto the extension-wide DeviceMemoryPool, which keys on device UUID. The bookkeeping now lives in a single .cpp so the module holds exactly one instance; the header only declares the API. release_decode_scratch() detaches the slabs under its lock and performs queue::wait()/sycl::free() outside it, so no unbounded device sync happens while the lock is held. Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
Update README_MOE_PREFILL_PERF.md and its CN counterpart for the
DeviceMemoryPool migration, and correct the moe_decode_release_scratch
contract: it must not overlap an in-flight decode call, since the acquire
entry points hand out a raw pointer.
Also stop a failed allocation from poisoning the slot. The pool records
{size, nullptr} before the result is checked, so every later request of
that size or smaller would be satisfied from the cached null without ever
retrying; detach the entry before throwing.
Co-authored-by: a32543254 <53296245+a32543254@users.noreply.github.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
luoyu-intel
left a comment
luoyu-intel
left a comment
There was a problem hiding this comment.
The main scratch-memory lifetime issue has been addressed in this PR. MOE no longer maintains a header-local DeviceScratchPool keyed by raw sycl::queue*. Instead, the scratch buffers reuse the extension-wide DeviceMemoryPool, the MOE bookkeeping is defined in a dedicated .cpp file, and release_decode_scratch() detaches the buffers before releasing the lock and then performs queue.wait() / sycl::free().
The remaining context-aware keying issue should be moved to PR #2280 or a follow-up to that PR. MOE slots 9 and 10 are still obtained through DeviceMemoryPool::get_device_key(), which currently keys only by device UUID. Therefore, if multiple SYCL contexts are created for the same physical device, MOE could theoretically reuse a USM scratch buffer allocated from a different context. PR #2280 already introduces (device, context) awareness for oneDNN engines, streams, and oneDNN scratch buffers, but it does not currently cover the MOE scratch slots or the DPAS atomic scratch slot.
Recommended approach:
- Keep the current scratch-lifetime fixes in this PR and merge it.
- In PR #2280, or in a follow-up, extend
DeviceMemoryPoolkeys to include both the device UUID and the SYCL context. - Use the same context-aware keying for oneDNN scratch, MOE repack/activation-sum scratch, and the DPAS atomic scratch buffer.
- Add a regression test covering two different SYCL contexts on the same device and verifying scratch-buffer isolation.
Until the context-aware keying is implemented in #2280, MOE decode scratch should be documented as requiring a single SYCL context per device and must not be driven concurrently from multiple contexts.
Signed-off-by: Dong, Bo1 <bo1.dong@intel.com>
|
/azp run Unit-Test-CUDA-AutoRound |
|
Azure Pipelines could not run because the pipeline triggers exclude this branch/path. |
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 referencew4a16grouped GEMM invllm-project/vllm-xpu-kernels. This tunes the Intel XPU (BMG/Xe2) INT4-sym MoE kernels inauto_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.hppdpas_w4a16_policy(WG128×256×32, SG4×8) for large M — the halved packed-nibble B stream makes the wider N tile pay off;dpas_w4a16_policy_m_8(WG8×64×32) for tiny M. m_16/m_32 alias the shared 64-wide-N INT8 shapes.A_avg_Mbucketing (≤4→m_8,≤8→m_16,≤128→m_32, else default). The32×64tile 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.hpplaunch_int4GEMV inner loop into a width-templatedint4_decode_chunkhelper with a 32→16→scalar ladder. Widens the fast path for shipped group sizes (32/64/128/256) while never regressinggroup_size==16.Docs —
README_MOE_PREFILL_PERF.md+_CNcounterpart 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
/azp run Unit-Test-CUDA-AutoRound.