Skip to content

feat(metax): add moe_fused_dense operator for fused MoE expert forward - #992

Open
rubik-hua wants to merge 1 commit into
InfiniTensor:masterfrom
rubik-hua:metax_moe_fused_dense
Open

rubik-hua wants to merge 1 commit into
InfiniTensor:masterfrom
rubik-hua:metax_moe_fused_dense

Conversation

@rubik-hua

Copy link
Copy Markdown

Implement fused dense forward on Metax backend, combining gate-up
projection (w1/w3), SiLU/GELU activation, down projection (w2), and
token shuffle into a single dispatch. Supports both fp16 and bf16
precision for end-to-end MoE layer inference.

image

单元测试:
image

模型验证
python examples/test_infer.py --device metax --model=/data/rubik/models/Qwen3-30B-A3B-Instruct-2507/ --enable-paged-attn --tp 2 --disable-prefix-caching --prompt "who are you"
image

  Implement fused dense forward on Metax backend, combining gate-up
  projection (w1/w3), SiLU/GELU activation, down projection (w2), and
  token shuffle into a single dispatch. Supports both fp16 and bf16
  precision for end-to-end MoE layer inference.
@rubik-hua
rubik-hua requested a review from a team September 18, 2026 08:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant