Skip to content

feat(metax): add moe_topk_softmax operator for MoE top-k gating - #993

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

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

Conversation

@rubik-hua

Copy link
Copy Markdown

Implement top-k gating softmax on Metax backend for Mixture-of-Experts.
Given gate logits and optional expert bias, computes top-k expert
probabilities via softmax and returns normalized routing weights along
with selected expert indices for token dispatch.

单元测试:
image

三个moe算子合一起,跑模型验证
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 top-k gating softmax on Metax backend for Mixture-of-Experts.
  Given gate logits and optional expert bias, computes top-k expert
  probabilities via softmax and returns normalized routing weights along
  with selected expert indices for token dispatch.
@rubik-hua
rubik-hua requested a review from a team September 18, 2026 08:16
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