Skip to content

feat(metax): add MoE optimizations and contiguous prefill for MetaX backend - #320

Open
yswang777 wants to merge 1 commit into
flagos-ai:mainfrom
yswang777:feat/metax-qwen36-a3b-optimizations
Open

feat(metax): add MoE optimizations and contiguous prefill for MetaX backend#320
yswang777 wants to merge 1 commit into
flagos-ai:mainfrom
yswang777:feat/metax-qwen36-a3b-optimizations

Conversation

@yswang777

Copy link
Copy Markdown

Summary

MoE and attention performance optimizations for MetaX backend, tested on MetaX C550 with Qwen3.6-35B-A3B.

Changes

MetaX backend (vendor/metax)

  • activation: fallback from SiluAndMul/GeluAndMul to torch.ops._C or native PyTorch ops
  • flash_attn: VLLM020_CONTIGUOUS_SINGLE_PREFILL - single-request prefill with contiguous KV cache
  • fused_moe: moe_align_block_size, moe_sum, SGL MoE, BF16 noscale, even K fastpath, naive decode, stage trace
  • register_ops: register moe_align_block_size, moe_sum
  • metax config: vendor:metax for topk_softmax

MoE utilities (ops/fused_moe)

  • activation: try torch.ops._C.silu_and_mul first, fallback to dispatch
  • fused_moe_utils: split staged configs, runtime profile config (VLLM020_MOE_RUNTIME_PROFILE), stage2 realign, mcoplib stage1/stage2 token fast paths, small M block override, v626 config

Environment Variables

Env Var Purpose
VLLM020_CONTIGUOUS_SINGLE_PREFILL Contiguous single-request prefill
VLLM020_MOE_BF16_NOSCALE BF16 noscale MoE kernel
VLLM020_MOE_EVEN_K_FASTPATH Even K dimension fastpath
VLLM020_SGL_MOE / VLLM020_SGL_ZERO_BIAS SGL MoE kernel
VLLM020_MOE_RUNTIME_PROFILE Runtime profile config selection
VLLM020_MOE_STAGE2_TOKEN_SUM_* Stage2 token-sum fast path
VLLM020_MOE_SMALLM_BLOCK_M Small M block size override
VLLM_FL_ENABLE_MOE_METHOD_OOT Enable OOT MoE without FlagGems

…ackend

- activation: fallback from SiluAndMul/GeluAndMul to torch.ops._C or native PyTorch ops
- flash_attn: VLLM020_CONTIGUOUS_SINGLE_PREFILL for single-request prefill with contiguous KV cache
- fused_moe: add moe_align_block_size, moe_sum, SGL MoE invoke, BF16 noscale fast path,
  even K fastpath, naive decode mode, stage trace
- fused_moe_utils: split staged configs, runtime profile config selection (VLLM020_MOE_RUNTIME_PROFILE),
  stage2 realign, mcoplib stage1/stage2 token fast paths, small M block override
- metax config: add vendor:metax backend for topk_softmax

Co-Authored-By: Claude <noreply@anthropic.com>
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@github-actions github-actions Bot added the core label Jul 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants