Hidden size and expert sizes 128 for fused MoE on OpenVINO GPU

#2
by rkazants - opened
Optimum Intel Internal Testing org

Increases the hidden size and the expert sizes from 16 to 128 (hidden_size, moe_intermediate_size, shared_expert_intermediate_size, and the vision out_hidden_size that follows the hidden size). Nothing else in the architecture changes.

Why: with expert input dimensions of 16, a 4-bit export with the default group size (128) leaves the experts uncompressed, so OpenVINO GPU does not fuse them into its MoE operation (moe_3gemm_fused_compressed). Smaller group sizes do not help either: the GPU MoE kernel needs a compression group of at least 64. With 128, the default int4 export compresses all experts and they are fused on both CPU (GatherMatmul) and GPU.

  • Same generator script, same initialization scheme; the layer std is 0.09 (about 1 / sqrt(hidden size)) instead of 0.25.
  • Same file layout: fused 3D experts, n-gram table shards, MTP head under mtp..
  • Health check on text, image and video prompts: logits std about 11, median greedy top-2 margins above 1.3, all experts used on 70+ token prompts, QSA drops 23–93% of visible tokens.
  • The model card is updated with the new sizes and checks.
  • The model grows from 8M to 71M parameters (33 MB β†’ 289 MB checkpoint), mostly the two 248320 Γ— 128 embedding tables.
  • An fp32 OpenVINO export of this checkpoint matches transformers token for token on CPU (text and image).

πŸ€– Generated with Claude Code

rkazants changed pull request status to merged

Sign up or log in to comment