Skip to content

perf(moe): wire qwen3_vl_moe to the fused decode-MoE kernel - #310

Merged
inureyes merged 1 commit into
mainfrom
perf/issue-306-qwen3-vl-moe-fused-kernel
Jun 16, 2026
Merged

inureyes merged 1 commit into
mainfrom
perf/issue-306-qwen3-vl-moe-fused-kernel

Conversation

@inureyes

Copy link
Copy Markdown
Member

Summary

Wires qwen3_vl_moe's SparseMoeBlock::forward to the fused single-token decode-MoE kernel, matching the pattern already used by qwen3_moe, qwen3_next, dots1, gemma4, qwen2_moe, and lfm2.

What changed

  • src/models/qwen3_vl_moe.rs: replaced the unconditional experts.forward + moe_weighted_sum block with the gated fused-kernel branch. When n_tokens == 1 and fused_moe_enabled() returns true, the block calls SwitchGLU::forward_fused_kernel(&x_flat, &topk_indices, &scores) (signature: fn forward_fused_kernel(&self, x: &MlxArray, indices: &MlxArray, scores: &MlxArray) -> Option<UniquePtr<MlxArray>>). The kernel returns None for any unsupported config (non-4/8-bit gate/up, non-4/6/8-bit down, group_size mismatch, or multi-token decode), and the match arm falls through to the original forward + moe_weighted_sum path. Prefill is unaffected. Activation is SwiGLU (act=0); no GeGLU or relu2 template is needed.
  • docs/benchmark_results/fused-moe-decode-kernel-design.md: updated the "Models covered" section, moving qwen3_vl_moe from the unwired list to the wired list and updating the count from six to seven.

Test plan

  • cargo check --lib --tests --features metal,accelerate: clean
  • cargo clippy --lib --tests --features metal,accelerate -- -D warnings: clean
  • Real-checkpoint validation on mlx-community/Qwen3-VL-30B-A3B-4bit (text-only decode, greedy temp-0 parity fused-on vs fused-off, tok/s measurement) is pending orchestrator.

Closes #306

SparseMoeBlock::forward in qwen3_vl_moe now gates on single-token
decode (n_tokens == 1) and fused_moe_enabled(), then delegates to
SwitchGLU::forward_fused_kernel (imported from qwen3_moe). The
fallback path (experts.forward + moe_weighted_sum) is unchanged for
prefill and any unsupported bit-width config. Activation is SwiGLU
(act=0); no GeGLU/relu2 template is needed.

Updates the 'Models covered' list in the fused-kernel design doc
to reflect seven wired model paths, moving qwen3_vl_moe out of the
unwired list.
@inureyes inureyes added type:performance Performance improvements priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) platform:macos macOS (Apple Silicon) specific status:review Under review status:done Completed and removed status:review Under review labels Jun 16, 2026
@inureyes
inureyes merged commit b8f99db into main Jun 16, 2026
5 checks passed
@inureyes
inureyes deleted the perf/issue-306-qwen3-vl-moe-fused-kernel branch June 16, 2026 08:54
@inureyes inureyes added this to the 0.3 milestone Jun 21, 2026
@inureyes inureyes self-assigned this Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata platform:macos macOS (Apple Silicon) specific priority:medium Medium priority status:done Completed type:performance Performance improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf(moe): wire qwen3_vl_moe to the fused decode-MoE kernel

1 participant