[GPU] Support u2 in the fused MoE 3-GEMM path by Cyberpunk1210 · Pull Request #39062 · openvinotoolkit/openvino

Cyberpunk1210 · 2026-10-21T07:34:46Z

Details:

Part 6 of the split of #37335 (series: #39057 … #39063). Depends on #39060 and #39061; this branch
contains both commits, so until they merge please review only the last commit. The native u2 prefill
GEMM is the follow-up #39063.

Brings moe_3gemm_swiglu to u2 expert weights and enables u2 for the GPU MoE passes.

Decode: u2 GEMV / batched GEMV variants, dispatched per GEMM as dtype code 3, with activation
loads aligned to the contiguous K values in each lane's weight byte. Per-group u2 zero points are
decoded in-kernel; a scalar zp (INT2_SYM) is broadcast, and whether a GEMM's zp is scalar is
answered per GEMM (GATE/UP/DOWN_ZP_SCALAR), since a mixed INT2_SYM/INT2_ASYM layer gives gate and
up different forms. Only the u2 variants have an ELEMS_PER_LANE == 1 branch, so the smaller group
size is admitted only when all three GEMMs are u2.
Prefill: u2 expert weights (and zp) are unpacked once into u4 scratch and served by the existing
oneDNN grouped GEMM. Functional, but the scratch is twice the size of the u2 expert weights; [GPU] Add a native u2 prefill GEMM to the fused MoE 3-GEMM path #39063
removes it.
Shared experts: a u2 shared projection is unpacked to u4 for oneDNN, per projection.
Mixed precision: MoE op fusion now also admits u4 in a mixed layer, as long as gate and up agree
on whether they are u4. This admits the NNCF recipe that keeps gate/up at u2 and gives down u4.
Pipeline: u2 offered to ConvertTiledMoeBlockToGatherMatmuls and
ConvertGatherMatmulToGatherMatmulCompressed only when MoE fusion is enabled — GatherMatmulCompressed
has no u2 kernel, and those passes now also run under OV_GPU_MOE_DISABLE_FUSION — and never to the
oneDNN-only GroupedMatMul pass; FuseGatedMLP and KeepMOE3GemmConstPrecision accept u2.

Tests: moe_3gemm u2 cases for decode GEMV, batched GEMV and prefill, scalar and per-group zero
points, mixed u2/u8 and u2/u4 layers, and the u2 shared-expert unpack; the MoE op fusion dtype table
for the u4 mixes.

Known: moe_accuracy_test_random/10 fails on master itself (see #39061) and is unchanged here.

Tickets:

N/A

AI Assistance:

AI assistance used: yes
Claude (Claude Code) was used to split [GPU][CPU] Support u2 (INT2) weight decompression: optimized FC kernel and fused MoE 3-GEMM #37335 into this series and to draft the change under
close supervision.
Human validation performed: Windows Release build on Arc B390; GPU *moe* unit tests (218/219,
the one failure being the pre-existing case above) and transformation tests (33/33).

🤖 Generated with Claude Code

Prepares MoeOpFusion for INT2 (u2) MoE models, and stops it fusing layers whose expert GEMMs mix dtypes, which the fused kernel decodes wrongly today. u2 itself stays inert until the GPU kernels land in a follow-up PR: the only caller is the GPU pipeline, and it does not yet offer u2 to ConvertTiledMoeBlockToGatherMatmuls / ConvertGatherMatmulToGatherMatmulCompressed, so no u2 GatherMatmulCompressed reaches this pass. Changes to the compressed branch of the 3-GEMM callback: - Weight dtype gate. Previously any combination fused, but moe_3gemm decodes all three GEMMs with the gate GEMM's dtype, so a layer with e.g. u4 gate/up and u8 down (NNCF mixed precision keeps sensitive experts at INT8) produced wrong output with no error. A layer now fuses only when its three GEMMs share one dtype in {u2, u4, i4, u8, i8}, or mix only u2 and u8 -- the one mix the follow-up kernels decode per GEMM. Other mixes stay in the GatherMatmul form, which is correct. - Scalar zero point. A rank-0 zp constant is a folded per-tensor zero point. An all-zero one is normalized to "no zp" (symmetric); a non-zero one is kept only for u2 GEMMs, and op::MOE validation now skips the num_experts check for rank-0 inputs. Mixing symmetric and asymmetric GEMMs in one layer bails out. - Group size. MOECompressed needs one group size for all three GEMMs, so a per-channel scale (and its u8 zp) is broadcast to the group-wise layout of the others; genuinely different group sizes bail out. Sub-byte zps are not expanded -- they are bit-packed and Type::size() rounds up to a byte -- so that case declines to fuse. - The 2-GEMM matcher applies the same scalar-zp normalization and keeps u2 in GatherMatmul form, since moe_gemm's only implementation (oneDNN) has no u2. Tests: compressed GatherMatmulCompressed variants of the MoE block -- until now every model in this file used plain f32 GatherMatmul, so the whole compressed half of the callback was untested -- covering the dtype admission table, the scalar-zp normalization, group-size unification and the per-channel scale/zp broadcast. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

The moe_3gemm decode kernels took one weight dtype from the gate GEMM (WEIGHT_COMPRESSEION_DT) and derived from it the GEMV variant, the per-expert weight stride and the per-expert zp stride for all three GEMMs. A layer whose GEMMs differ -- NNCF mixed precision produces e.g. gate/up u8 with down u4 -- therefore decoded at least one GEMM with the wrong dtype and the wrong stride: wrong output, no error. Emit GATE/UP/DOWN_WEIGHT_DT and use them for the GEMV dispatch, the expert weight stride and the zp stride of each GEMM. The activation interleave stays keyed on one dtype per kernel: gate and up share one SLM copy whose layout is picked from the gate dtype, so they must still agree on whether they are u4; down stages its own copy and is free. The prefill paths already took each GEMM's dtype. Also fix the host-side GEMV check: FAKE_GROUP_SIZE in the kernel is the minimum of the gate/up and the down group size, but the assert only looked at gate/up, so a smaller down group slipped through to an ELEMS_PER_LANE with no kernel branch. Tests: moe_3gemm_compressed_gpu_mixed_dtype, gate/up u8 + down u4 and gate/up u4 + down u8, for single-token decode, batched GEMV and just above the batched-GEMV threshold, against a reference built from the exact dequantized weights. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Brings the fused moe_3gemm_swiglu path to u2 (INT2) expert weights and enables u2 for the GPU MoE passes. Depends on "[Transformations] Support u2 expert weights in MoE op fusion" and on "[GPU] Decode each GEMM of the fused MoE 3-GEMM with its own weight dtype". Prefill runs u2 through the existing oneDNN grouped GEMM by unpacking it to u4; a native u2 prefill GEMM follows separately. Decode (GEMV and batched GEMV). u2 GEMV variants of gate_up and down, dispatched per GEMM as dtype code 3. They read activations from SLM so that lane L's activations line up with the contiguous K values packed in its weight byte (w[4L..4L+3]); the u4 path's strided read targets an interleaved layout and is left unchanged. Per-group u2 zero points are decoded in-kernel; a scalar zero point (INT2_SYM) is broadcast, and whether a GEMM's zp is scalar is answered per GEMM (GATE/UP/DOWN_ZP_SCALAR), since a mixed INT2_SYM/INT2_ASYM layer gives gate and up different forms. Only the u2 variants have an ELEMS_PER_LANE == 1 branch, so the smaller group size is admitted only when all three GEMMs are u2. Prefill. oneDNN and the micro-GEMM have no u2 type, so the u2 expert weights (and zp) are unpacked once into u4 scratch and the u4 grouped GEMM consumes them. The unpack runs once per bound weight buffer rather than per prefill -- the weights are constant, since expert offload is rejected for u2 -- which removed ~430 ms of TTFT per prefill on Qwen3.6-35B-A3B. Shared experts. A u2 shared projection is reported to oneDNN as u4, so it is unpacked the same way; the unpack is per projection because NNCF mixed precision makes any subset of the projections u2. Mixed precision. MoE op fusion now also admits u4 in a mixed layer, as long as gate and up agree on whether they are u4 (they share one SLM activation copy whose layout the gate dtype picks). This admits the NNCF recipe that keeps gate/up at u2 and gives down u4. It lands here rather than with the fusion change because only now does every admitted combination decode per GEMM. Pipeline: u2 is offered to ConvertTiledMoeBlockToGatherMatmuls and ConvertGatherMatmulToGatherMatmulCompressed only when MoE fusion is enabled -- GatherMatmulCompressed has no u2 kernel, and those passes now also run under OV_GPU_MOE_DISABLE_FUSION -- and never to ConvertGroupedMatMulToGroupedMatMulCompressed (oneDNN-only). FuseGatedMLP and KeepMOE3GemmConstPrecision accept u2. Tests: moe_3gemm u2 cases for decode GEMV, batched GEMV and prefill, scalar and per-group zero points, mixed u2/u8 and u2/u4 layers, and the u2 shared-expert unpack; the MoE op fusion dtype table for the u4 mixes. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

A failed u2 weight unpack or grouped GEMM used to be caught and silently replaced by the batched GEMV. No other GPU primitive degrades that way at execution time, so let the error propagate. Also drop measurement figures from comments. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Comment-only cleanup of the u2 fused MoE 3-GEMM changes and tests: keep the why, drop restated code and history. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Comment-only cleanup: keep the why, drop restated code, history and internal labels. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Comment-only cleanup: keep the why, drop restated code and history. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

p-durandin · 2026-10-22T05:36:43Z

build_jenkins

Drops an unneeded lambda capture (clang -Wunused-lambda-capture). No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

p-durandin · 2026-10-22T09:31:28Z

build_jenkins

…ith folded ops - Each compressed expert weight's dtype is now a pattern predicate (type_matches_any in the 3-GEMM matcher, "not u2" in the 2-GEMM matcher); only the cross-GEMM rule (uniform, or all mixable with gate/up agreeing on u4) stays in the callback. - Per-channel scale/zp expansion uses Reshape + Broadcast with make_try_fold instead of a raw byte copy. - The scalar-zp normalization is shared by both matchers and takes whether a non-zero zp may pass; the 2-GEMM matcher never admits one, which the old comment misstated. No change in what fuses. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

…op the broadcast wrapper lambda Same behaviour: non-Constant zp still passes unless rank-0; the wrapper lambda only forwarded group_size. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

…e WEIGHT_COMPRESSEION_DT validate_impl now checks WEIGHT/SCALE/ZP of gate, up and down, plus the cross-input constraints the kernel already relies on. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

moe_3gemm_swiglu_opt.cpp: took upstream's gate GEMM fused with the up output as a binary post-op, and routed its weight and zp through grouped_weight_mem/grouped_zp_mem so u2 gate weights still come from the unpacked u4 scratch. fuse_gated_mlp.cpp: deleted, as upstream removed FuseGatedMLP (openvinotoolkit#39202); our u2 additions to it had no remaining callers. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Prefill runs the shared expert through oneDNN with per-projection dtypes; validate_impl cannot tell decode from prefill, so the check rejected valid models. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

src/common/transformations/src/transformations/common_optimizations/moe_op_fusion.cpp, tests/.../convert_gather_matmuls_moe_block_to_moe_op.cpp, src/core/src/op/moe.cpp: same resolution as int2-moe-fusion (one zp path for u2/u3, u3 in the weight list, upstream moe.cpp), plus this branch's mixed u2/u4/u8 gate/up rule re-applied on top and kept; test table keeps both u3 and our u2 rows. src/plugins/intel_gpu/src/graph/impls/ocl_v2/utils/jitter.cpp: kept u2 and u3. src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp: u3 added to supported_compressed_weights_types; our gather_matmul_weights_types (that list + u2 when fusion is on) is derived from it, so u3 reaches the GatherMatmul passes too. src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe/moe_3gemm_swiglu_opt.hpp: u2 and u3 both supported; the "zp bit width == weight bit width" check (batched GEMV zp stride, non-scalar zps only) exempts u3. src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe/moe_3gemm_swiglu_opt.cpp: ported onto MoeGemmImpl. u2 is batched-GEMV capable (decode_impl batched_gemv); any u2 GEMM forces prefill_impl grouped_gemm. The u2->u4 unpack of routed weights/zp runs in exec_grouped_gemm_core (both mask gen entry points); shared-expert unpack stays in execute() for every prefill path. exec_gemm and its args-cache key use the unpacked buffers for u2 GEMMs. The grouped GEMM zp dtype is now per GEMM (u4 for an unpacked u2 zp). A scalar zp stays on our broadcast-into-u4 unpack for u2 GEMMs: upstream host-scalar zp (u3 only) cannot create a oneDNN matmul with u4 weights, so u2 never uses _grouped_zp_host_scalars. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Lyamin-Roman · 2026-11-02T17:30:45Z

@@ -44,7 +44,11 @@ struct Moe3GemmReference {
    std::vector<uint8_t> pack(const std::vector<uint8_t>& unpacked) {


Please add functional tests to src/plugins/intel_gpu/tests/functional/subgraph_tests/dynamic/moe.cpp for the u2 case, covering both 3-GEMM and 2-GEMM scenarios (to ensure that the pipeline does not crash, but instead falls back to decomposition)

Lyamin-Roman · 2026-11-02T17:34:42Z

+        // The shared-expert oneDNN primitives consume their u4 scratch on every prefill path. The routed experts' unpack
+        // lives in exec_grouped_gemm_core, shared by both mask gen entry points.
+        if (has_u2_shared_weights()) {
+            unpack_u2_shared_weights_for_prefill(instance, scratch, events);


The returned event is not handled in any way

Lyamin-Roman · 2026-11-02T17:50:29Z

+        MoeU2SharedParams{false, false, false, 64, 64, "none_u2_prefill"},
+        MoeU2SharedParams{true, false, false, 64, 64, "gate_u2_prefill"},
+        MoeU2SharedParams{false, false, true, 64, 64, "down_u2_prefill"},
+        MoeU2SharedParams{true, false, true, 64, 64, "gate_down_u2_prefill"}),


And what about the decode?

v-Golubev

A question on the overall direction. oneDNN already supports u2 FC (#37212), and the oneDNN path is preferable: once it lands, prefill needs no custom code on our side, which means much less to maintain. So maybe a better idea is to contribute to u2 grouped GEMM support in OneDNN instead of a custom kernel on our side.

Moreover, the fused MoE path is enabled only for IMMAD + oneDNN (see the gate in transformations_pipeline.cpp), so non-IMMAD devices never reach these kernels. So there is no argument on non-IMMAD devices support in MoE node.

Could you please clarify which code paths this PR closes that oneDNN grouped GEMM wouldn't?

praasz · 2026-11-03T11:29:44Z

+constexpr size_t kE = 2;
+constexpr size_t kHidden = 16;


Suggested change

constexpr size_t kE = 2;

constexpr size_t kHidden = 16;

constexpr size_t k_e = 2;

constexpr size_t k_hidden = 16;

Detail: in common part this naming convention should be used.

Cyberpunk1210 and others added 3 commits October 21, 2026 10:13

Cyberpunk1210 requested review from a team as code owners October 21, 2026 07:34

github-actions Bot added category: Core OpenVINO Core (aka ngraph) category: GPU OpenVINO GPU plugin category: transformations OpenVINO Runtime library - Transformations labels Oct 21, 2026

sys-openvino-ci added the ExternalPR External contributor label Oct 21, 2026

Cyberpunk1210 and others added 4 commits October 21, 2026 16:19

[GPU] Trim u2 fused MoE 3-GEMM comments

8b085b8

Comment-only cleanup of the u2 fused MoE 3-GEMM changes and tests: keep the why, drop restated code and history. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

[Transformations] Trim u2 MoE fusion comments

ddaa10a

Comment-only cleanup: keep the why, drop restated code, history and internal labels. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

[GPU] Trim per-GEMM MoE dtype comments

ac3cd0a

Comment-only cleanup: keep the why, drop restated code and history. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

[GPU] Make the u2 unpack work-group size constexpr

86e5253

Drops an unneeded lambda capture (clang -Wunused-lambda-capture). No behaviour change. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Cyberpunk1210 <genji.1231@outlook.com>

Cyberpunk1210 mentioned this pull request Oct 23, 2026

[GPU][CPU][Transformations] u2 (INT2) weight support on releases/2026/4 #38837

Closed

blairducrayoppat mentioned this pull request Oct 26, 2026

[GPU] Break top-k ties by expert index in the MoE 3-GEMM test reference #39205

Open

Cyberpunk1210 and others added 5 commits October 27, 2026 15:44

github-actions Bot removed the category: Core OpenVINO Core (aka ngraph) label Oct 30, 2026

Lyamin-Roman reviewed Nov 2, 2026

View reviewed changes

v-Golubev requested changes Nov 3, 2026

View reviewed changes

praasz reviewed Nov 3, 2026

View reviewed changes

Navigation Menu

Repository navigation

[GPU] Support u2 in the fused MoE 3-GEMM path - #39062

[GPU] Support u2 in the fused MoE 3-GEMM path#39062
Cyberpunk1210 wants to merge 14 commits into
openvinotoolkit:masterfrom
Cyberpunk1210:int2-moe-gpu

Cyberpunk1210 commented Oct 21, 2026 •

edited

Loading

Uh oh!

p-durandin commented Oct 22, 2026

Uh oh!

p-durandin commented Oct 22, 2026

Uh oh!

Lyamin-Roman Nov 2, 2026

Uh oh!

Lyamin-Roman Nov 2, 2026

Uh oh!

Lyamin-Roman Nov 2, 2026

Uh oh!

v-Golubev left a comment

Uh oh!

praasz Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

		@@ -44,7 +44,11 @@ struct Moe3GemmReference {
		std::vector<uint8_t> pack(const std::vector<uint8_t>& unpacked) {

Conversation

Cyberpunk1210 commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Details:

Tickets:

AI Assistance:

Uh oh!

p-durandin commented Oct 22, 2026

Uh oh!

p-durandin commented Oct 22, 2026

Uh oh!

Lyamin-Roman Nov 2, 2026

Choose a reason for hiding this comment

Uh oh!

Lyamin-Roman Nov 2, 2026

Choose a reason for hiding this comment

Uh oh!

Lyamin-Roman Nov 2, 2026

Choose a reason for hiding this comment

Uh oh!

v-Golubev left a comment

Choose a reason for hiding this comment

Uh oh!

praasz Nov 3, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Cyberpunk1210 commented Oct 21, 2026 •

edited

Loading