onnxruntime
f340e685 - [CPU/MLAS] Add QMoE shared layout INT2 AVX-512/AVX2 kernels and direct CPU QMoE (#33785)

Commit
2 days ago
[CPU/MLAS] Add QMoE shared layout INT2 AVX-512/AVX2 kernels and direct CPU QMoE (#33785) ## Description Add direct CPU QMoE execution over immutable SM80-prepacked INT2 expert banks, allowing CPU and hybrid CPU/GPU workers to consume the same weights and scale banks without unpacking or transposing them. ## Changes - Add scalar, AVX2, and AVX-512 INT2 kernels for block sizes 32, 64, and 128, with dedicated INT2 translation units and runtime dispatch. - Support both accepted prepacked weight-shape conventions, FP32/FP16/BF16 CPU inputs, optional biases, and explicit routing weights used as supplied for either normalization setting. - Align INT2 FC1 projection boundaries across CPU, CUDA decode, and CUDA packed prefill. Handle negative-infinite SiLU gates. A range containing a BF16 scale of magnitude >= 2^127, which would overflow before multiplication by a small activation, runs on an exact half-scale scalar kernel chosen by a vectorized pre-check; the AVX2 and AVX-512 kernels are unchanged. - Validate complete gate/up-pair ranges and add independent-reference coverage for both weight layouts, both scale layouts, partial ranges, and explicit AVX2 selection on AVX-512 hosts. - Add CPU/CUDA decode/prefill comparisons with identical SM80 banks, CPU-EP fallback disabled, and dense-dequant fallback prevented. Fix the cached-scale profiling test to start profiling once after warmup. - Fix narrowing and unused-variable errors reported by Windows, Clang, ARM, and WASM CI. Merge current main to resolve the PR conflict. ## Validation - CUDA 13 Release build on H200/SM90, with `onnxruntime_USE_FPA_INTB_GEMM=ON`, CUDA internal tests, and CUDA profiling enabled. - **58/58 focused C++ tests passed:** all shared-layout tests, prepacked INT2 CPU operator tests, CPU/CUDA decode/prefill parity, packed decode and cached-scale regressions, and packed-prefill tests. - `Int2LoneExtremeBf16ScaleIsFoundAtEveryPosition` places one extreme or boundary BF16 scale at every column and block position for both scale layouts, both weight layouts, and all backends. It was mutation-checked against a skipped tail, a skipped last block, an off-by-one count, and a wrong threshold. BF16 block-32, M=1 decode GEMM (K=4096) takes 34.3 us versus 33.8 us without overflow handling. - `test_qmoe_shared_direct_edges.py`: **14 tests and 78 subtests passed**, including the routing-weight regression reported by CI. - Both complete modified test files pass Clang with `-Werror -Wshorten-64-to-32`. Standalone INT2 tests pass with SIMD enabled and with both AVX backends disabled. - `lintrunner` and `git diff --check` pass. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading