Skip to content

[CPU] Add MoEPacked contrib op and AVX2 INT8 GGUF kernels for block-packed weights - #33975

Draft
BBaditi wants to merge 2 commits into
mainfrom
qmoe-packed-gguf-cpu-ep
Draft

BBaditi wants to merge 2 commits into
mainfrom
qmoe-packed-gguf-cpu-ep

Conversation

@BBaditi

@BBaditi BBaditi commented Nov 8, 2026 •

Copy link
Copy Markdown
Contributor

Description

Adds a native CPU Execution Provider implementation for com.microsoft::MoEPacked along with a header-only 32-wide AVX2 INT8 (_mm256_maddubs_epi16 / vpmaddubsw) GGUF block-quantized dot-product library (gguf_vec_dot_cpu.h) supporting all 12 quantized GGUF formats in gguf_quant_tables.h (Q2_0, Q4_0, Q5_0, Q8_0, Q5_K, Q6_K, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ4_NL, IQ4_XS) plus F32, F16, and BF16.

Key Optimizations

  • Per-32-Element GgufQ8Block Activation Quantization & 32-Wide AVX2 INT8 (vpmaddubsw):
    • Quantizes input activations (x) once per active token and intermediate SwiGLU tiles (inter) on the fly into 32-element GgufQ8Block (int8_t qs[32], float d, int32_t bsum, matching CUDA block_q8_1).
    • Evaluates 32 weights per AVX2 instruction using _mm256_maddubs_epi16 + _mm256_madd_epi16 across codebook grid formats (IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS via 2-KB L1D sign LUT + vpsignb), non-linear LUT formats (IQ4_NL, IQ4_XS via pshufb + signed×signed vpsignb), and uniform/superblock formats (Q2_0, Q4_0, Q5_0, Q8_0, Q5_K, Q6_K), deferring horizontal reduction (HsumAvx2) and bsum zero-point correction to the end of each row.
  • Multi-Row Register Blocking (2x Gate/Up + 4x Down) & 2D Low-Latency Worker Pool:
    • VecDotGgufGateUpRow2xQ8Cpu processes 2 intermediate rows (4 weight rows: gate_r0, gate_r1, up_r0, up_r1) in lockstep in Phase 1 (kFc1Tile = 32), and VecDotGgufRow4xQ8Cpu processes 4 output rows in lockstep in Phase 2 (kFc2Tile = 64).
  • Expert-Grouped Multi-Token Prefill (M > 1) Reuse (ColdExpertGroup):
    • Groups multi-token routes by expert_id when M > 1 so all B tokens routed to the same expert reuse each GGUF weight row pair/quad hot in L1D/L2 cache (2.20x–3.60x speedup at M = 16), while keeping a zero-overhead direct path for M = 1 decode.

Single-Layer Qwen3.8-Flash-Next (E=512, H=2560, I=640, k=10) CPU MoEPacked Latency & Parity

Layer (Gate/Up + Down Format) Expert Size (MB/exp) Cosine vs FP32 Dequant Ref Relative L2 Error M=1 CPU MoEPacked (us) M=16 Ungrouped (ms) M=16 Expert-Grouped (ms / Speedup)
UD-IQ1_S L00 (IQ2_XXS + Q2_0) 1.31 MB 0.999935 1.14% 370.6 us 3.92 ms 1.09 ms (3.60x)
UD-IQ1_S L04 (IQ2_S + Q2_0) 1.51 MB 0.999929 1.19% 434.4 us 4.18 ms 1.42 ms (2.94x)
UD-IQ2_XS L00 (IQ2_XXS + IQ4_NL) 1.77 MB 0.999933 1.16% 376.3 us 4.05 ms 1.28 ms (3.16x)
UD-IQ2_XS L03 (IQ2_S + IQ4_NL) 1.97 MB 0.999928 1.20% 446.6 us 4.29 ms 1.61 ms (2.66x)
UD-IQ2_XS L04 (IQ2_S + Q5_0) 2.18 MB 0.999930 1.18% 476.9 us 4.41 ms 1.78 ms (2.48x)
UD-IQ3_XXS L14 (IQ3_XXS + IQ4_NL) 2.18 MB 0.999934 1.15% 533.7 us 4.30 ms 1.95 ms (2.20x)
UD-IQ3_XXS L47 (IQ4_XS + IQ4_NL) 2.66 MB 0.999932 1.17% 327.6 us 3.84 ms 1.21 ms (3.17x)

Verification

  • ctest -R GgufQuantPackedTest: Verifies single-row, fused Gate+Up, 2x Gate/Up, and 4x Down dot products across all 12 quantized GGUF formats (plus F32/F16/BF16) and MoEPacked CPU vs reference dequantization (M = 1 and M = 8 with bit-exact grouped vs ungrouped parity).
  • pytest onnxruntime/test/python/transformers/test_gguf_packed_ops.py: 11 passed, 174 subtests passed.

@BBaditi
BBaditi force-pushed the qmoe-packed-gguf-cpu-ep branch 2 times, most recently from d3283ae to 5ac6332 Compare November 9, 2026 06:17
Base automatically changed from tlwu/qwen38-gguf to main November 9, 2026 17:56
@BBaditi
BBaditi force-pushed the qmoe-packed-gguf-cpu-ep branch from 5ac6332 to 85653e0 Compare November 10, 2026 02:20
@tianleiwu

Copy link
Copy Markdown
Contributor

Please fix CI failures

@BBaditi
BBaditi force-pushed the qmoe-packed-gguf-cpu-ep branch from 85653e0 to 3cf8acc Compare November 10, 2026 08:11
…acked weights

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

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants