[CPU/MLAS] Add QMoE shared 4-bit expert layouts, AVX-512 kernels, and FileExpertStore by tianleiwu · Pull Request #33450 · microsoft/onnxruntime

Tianlei Wu (tianleiwu) · 2026-10-19T08:45:11Z

Description

Part 1 of 3 in the Portable QMoE Offload series adds CPU shared INT4, MXFP4 and NVFP4 expert layouts, scalar/AVX-512
expert kernels, and FileExpertStore with bounded, pin-counted host-slot caching and asynchronous prefetch.
The MLAS build owns the kernels by default, with a standalone/provider-owned compilation path when disabled.

Key changes

Shared layouts and reference INT4 packers preserve the prepacked expert bytes without introducing a second weight file.
Runtime-dispatched AVX-512 kernels support multi-row reuse, FP16/BF16 boundaries and optional bias handling.
File-backed banks are read into bounded cache slots; eligible small heap-deserialized banks receive owned resident backing.
INT4 partial ranges fall back to scalar execution when endpoints are not four-aligned.
SIMD scale rounding and BF16 narrowing preserve NaNs; scalar FP16 conversion preserves signed zero.
Standalone Clang compilation uses scoped target attributes instead of GCC-only target pragmas.
Generic POSIX initializer mappings retain writable private copy-on-write behavior. This PR does not change global
initializer mapping policy; caller-created immutable read-only mappings can still be used as expert banks.

Validation

Commands below were run from the repository root unless noted:

.venv/bin/python onnxruntime/test/python/transformers/test_qmoe_shared_layout.py:
passed, 12 tests passed and 4 optional tests skipped.
/usr/local/bin/cmake --build build/pr33505_cpu/Release --target onnxruntime_common --parallel 8:
passed; recompiles the actual POSIX mapping implementation.
/usr/local/bin/cmake --build build/pr33505_cpu/Release --target onnxruntime_mlas --parallel 8:
passed, including the separate AVX-512 translation unit.
Focused C++ executables built from the committed tests:
GCC and Clang each passed all 7 tests under Intel SDE -skx; the production MLAS-linked executable also
passed all 7 tests. ORT_QMOE_FP4_GROUP8_MASK=15 passed all 7 tests.
Focused platform tests run from the build output directory with
--gtest_filter='FileIoTest.*:RandomAccessFileTest.*': passed all 14 tests, including writable COW mappings.
Native ASan/UBSan focused C++ run: passed, 4 tests passed and 3 SIMD tests explicitly skipped.
Touched-file lintrunner: passed. git diff --check: passed.

The local AMD host lacks AVX-512, so SIMD correctness was verified by actual instruction emulation, not native
performance testing. CUDA golden/artifact tests and macOS execution were not run.

Motivation and Context

Cold experts need a portable CPU execution path and bounded host caching while retaining the same quantized storage
used by accelerator execution. Keeping checked geometry, reference execution and SIMD dispatch behind a shared
layout contract allows later offload integration without making immutable bank policy apply to mutable training tensors.

The review follow-up addresses all six applicable GitHub comments, adds regressions for two additional correctness
gaps, and confirms the spin-pool comment is obsolete because that file is no longer part of this PR.

Copilot

🟡 Changes recommended

Memory-safety, mutable-initializer compatibility, numerical correctness, and portability defects remain unresolved.

7 open issues

What changed in this PR

Adds CPU infrastructure for the first stage of portable QMoE offloading.

Changes:

Adds shared INT4/FP4 layouts, scalar and AVX-512 kernels, and bounded expert caching.
Introduces topology-aware workers and shared read-only POSIX mappings.
Integrates MLAS compilation and adds Python reference packers and tests.

File	Description
`onnxruntime/test/python/transformers/test_qmoe_shared_layout.py`	Adds layout, kernel parity, and cache tests.
`onnxruntime/test/python/transformers/qmoe_shared_layout_ref.py`	Adds INT4 reference packing and computation.
`onnxruntime/core/platform/posix/env.cc`	Changes little-endian mappings to shared read-only storage.
`onnxruntime/core/mlas/lib/platform.cpp`	Implements AVX-512 capability query.
`onnxruntime/core/mlas/inc/mlas.h`	Declares the capability query.
`onnxruntime/contrib_ops/cuda/moe/qmoe_spin_pool.h`	Adds topology-aware workers and spin controls.
`onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout.h`	Declares layouts, stores, and kernel interfaces.
`onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout.cc`	Implements layouts, caching, scalar kernels, and dispatch.
`onnxruntime/contrib_ops/cpu/moe/qmoe_shared_layout_avx512.cc`	Adds vectorized expert kernels and conversions.
`cmake/onnxruntime_providers_cpu.cmake`	Adjusts ownership of QMoE source compilation.
`cmake/onnxruntime_mlas.cmake`	Adds offload option and separate AVX-512 compilation.

_{Review effort: Balanced}

Copilot

🟡 Changes recommended

Unresolved storage-validation and mapping defects can permit out-of-bounds reads or silently change expert weights.

6 open issues

7 resolved since last review

_{Review effort: Balanced}

Copilot

🟡 Changes recommended

Incorrect legacy-weight prepacking, an allocation-failure data race, and CUDA tests that permit CPU fallback remain unresolved.

3 open issues

6 resolved since last review

_{Review effort: Balanced}

+  alignas(64) static __m512 fallback_lut[256];
+  PopulateScaleFp4Lut2D(base_lut, global_scale, round_scale_bf16, fallback_lut);
+  return fallback_lut;


+
+        outs = []
+        for model in (m_raw, m_pk):
+            sess = ort.InferenceSession(model, providers=["CUDAExecutionProvider"])


… FileExpertStore Add SharedExpertBlockV1/V2/V3 4-bit expert layouts (MXFP4/INT4/NVFP4), AVX-512/AVX-512-FP16/AVX2 CPU expert GEMV kernels, MLAS FP16 dispatch, FileExpertStore with bounded host-slot caching over MAP_SHARED/MAP_PRIVATE/heap-scale banks, and read-only MAP_SHARED external initializer mapping on little-endian POSIX. Address review feedback: - Fix scalar HalfToFloat subnormal exponent (127 - 14 - e) and add exhaustive scalar/SIMD FP16 bit-pattern tests. - Track regular-file vs heap backing in FileExpertStore, retain independent resident backing for accepted heap banks, and add IsFileBackedReadOnlyMapping / EvictFileBackedHostPages helpers with heap eviction safety tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Guard partial INT4 SIMD ranges, preserve NaNs and signed zero, fix Clang target scoping, and validate inferred storage offsets. Restore writable private POSIX mappings for mutable initializers and add focused kernel/cache/mapping regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…oE support - Reject null or zero-byte buffers in TensorExpertStore::Initialize before range validation. - Exclude private mappings (r--p) from IsFileBackedReadOnlyMapping and automatic /proc/self/maps file reloading so modified COW mappings are snapshotted or rejected rather than discarded. - Remove the 4096-byte cap in ComputePointerAlignment so alignments > 4096 work in MakeExpertBlockLease, TensorExpertStore, and FileExpertStore. - Retain an open Win32 HANDLE (FILE_SHARE_READ|WRITE|DELETE) across FileExpertStore lifetime and read with positional OVERLAPPED ReadFile. - Coalesce pending Prefetch requests by expert ID, refresh generations, and bound the pending queue to pre-reserved cache_slots. - Move the 16 KiB FP4 2D LUT cache entries out of inline TLS into a small per-thread handle with lazily allocated ways. - Gate onnxruntime_ENABLE_MLAS_QMOE_OFFLOAD off when onnxruntime_DISABLE_CONTRIB_OPS or minimal builds are enabled. - Add canonical INT4 QMoE view/kernels (row-major and block-major scales, SM80 prepacked and raw row-major weights, FP16/BF16) and support weights_prepacked=1 INT4 QMoE in QMoECPU. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

github-advanced-security Bot found potential problems Oct 19, 2026

View reviewed changes

Comment thread onnxruntime/test/python/transformers/qmoe_shared_layout_ref.py Fixed

Comment thread onnxruntime/test/python/transformers/qmoe_shared_layout_ref.py Fixed

Comment thread onnxruntime/test/python/transformers/test_qmoe_shared_layout.py Fixed

Tianlei Wu (tianleiwu) force-pushed the tlwu/20261019/qmoe-offload-p1-cpu-shared-layout branch from a0d08bf to 7a92038 Compare October 19, 2026 09:00

mirounga self-requested a review October 19, 2026 13:41

Tianlei Wu (tianleiwu) force-pushed the tlwu/20261019/qmoe-offload-p1-cpu-shared-layout branch from 7a92038 to 7e3016c Compare October 20, 2026 09:19

This was referenced Oct 20, 2026

[CUDA] Add session-scoped MoE/QMoE offload planner, lifecycle controller, and memory budget ledger #33451

Open

[CUDA] Add portable QMoE CPU/GPU offload backend with doorbell decode, adaptive promotion, and DMA prefill #33452

Open

Tianlei Wu (tianleiwu) force-pushed the tlwu/20261019/qmoe-offload-p1-cpu-shared-layout branch from 7e3016c to 5f1dc94 Compare October 20, 2026 10:09

Tianlei Wu (tianleiwu) added this pull request to stack #33495 October 20, 2026 18:52

Tianlei Wu (tianleiwu) force-pushed the tlwu/20261019/qmoe-offload-p1-cpu-shared-layout branch 2 times, most recently from c4535f3 to 3c007ec Compare October 21, 2026 06:32

Tianlei Wu (tianleiwu) marked this pull request as ready for review October 21, 2026 18:12

Copilot AI balanced review requested due to automatic review settings October 21, 2026 18:12

Copilot started reviewing on behalf of Tianlei Wu (tianleiwu) October 21, 2026 18:14

View session

Copilot AI reviewed Oct 21, 2026

View reviewed changes

Tianlei Wu (tianleiwu) force-pushed the tlwu/20261019/qmoe-offload-p1-cpu-shared-layout branch from 3c007ec to 0cc4f64 Compare October 22, 2026 01:25

Tianlei Wu (tianleiwu) requested a balanced review from Copilot October 23, 2026 00:07

Copilot started reviewing on behalf of Tianlei Wu (tianleiwu) October 23, 2026 00:08

View session

Copilot AI reviewed Oct 23, 2026

View reviewed changes

github-advanced-security Bot found potential problems Oct 23, 2026

View reviewed changes

Comment thread onnxruntime/test/python/transformers/test_qmoe_shared_layout.py Fixed

Tianlei Wu (tianleiwu) requested review from Akshay Sonawane (apsonawane), kunal-vaishnavi and Ti-Tai Wang (titaiwangms) and a balanced review from Copilot October 23, 2026 23:38

Copilot started reviewing on behalf of Tianlei Wu (tianleiwu) October 23, 2026 23:44

View session

Copilot AI reviewed Oct 23, 2026

View reviewed changes

Tianlei Wu (tianleiwu) and others added 5 commits October 23, 2026 20:59

test(qmoe): initialize ort on import failure to satisfy CodeQL

8c7e40c

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

fix(qmoe): tolerate stat/maps device mismatch on stacked filesystems

ea94e4b

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Tianlei Wu (tianleiwu) force-pushed the tlwu/20261019/qmoe-offload-p1-cpu-shared-layout branch from dd0da0f to ea94e4b Compare October 24, 2026 03:59

Navigation Menu

Repository navigation

[CPU/MLAS] Add QMoE shared 4-bit expert layouts, AVX-512 kernels, and FileExpertStore - #33450

[CPU/MLAS] Add QMoE shared 4-bit expert layouts, AVX-512 kernels, and FileExpertStore#33450
Tianlei Wu (tianleiwu) wants to merge 5 commits into
mainfrom
tlwu/20261019/qmoe-offload-p1-cpu-shared-layout

Tianlei Wu (tianleiwu) commented Oct 19, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Tianlei Wu (tianleiwu) commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Tianlei Wu (tianleiwu) commented Oct 19, 2026 •

edited

Loading