[KMCompiler]Fix fused moe vendor gating (metax fast paths)

[LittleShun1214]

PR Category

[ Operator ]

Type of Change

[ Bug Fix ]

Description

Fix and refine the generic fused MoE operator for vendor-specific tuning.

Bug fixes

• ops/fused_moe.py: remove the hard-coded metax-only activation switch (use_separate_activation_for_metax_qwen36_large) that was evaluated on every platform and changed the fused-MoE path for non-metax users running Qwen3.6 shapes (E=256, K=2048, M>=8192). The decision now comes from the vendor descriptor (moe_separate_activation), preserving the original path everywhere else.
• Restore the vendor-aware MoE smem/NS clamping (_get_max_smem_bytes, _clamp_num_stages_for_vendor, _is_h20) and the H20 SWAP_AB heuristic that were dropped from the generic config heuristics; NVIDIA default behavior is unchanged (VendorDescriptor defaults reproduce previous values exactly).

New vendor-gated kernel paths (all default-off, other platforms unaffected)

• USE_INT32_OFFSETS: int32 token/offset arithmetic when every offset product stays below 2^31 (saves registers); GEMM2 topk-merged rows are normalized back to model-level M before the vendor M threshold applies.
• FAST_BF16_OUTPUT: bf16 accumulator conversion with rtne_no_nan.
• Both are exposed as VendorDescriptor fields (moe_use_int32_offsets, moe_fast_bf16_output, moe_separate_activation, moe_fast_paths_min_m) with safe defaults (False / 0), consistent with the existing MUSA vendor parameterization (max_smem_bytes, moe_direct_sum_enabled, ...).

MetaX backend (runtime/backend/_metax/)

• Enable the fast paths above (bf16, M >= 8192) via the vendor descriptor; drop the per-shape config keys (USE_INT32_OFFSETS, FAST_BF16_OUTPUT, SEPARATE_ACTIVATION) in favor of the descriptor fields.

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT. (Qwen3.6 real-shape correctness vs vllm_metax reference: 2 passed, all token counts)

Performance

vs upstream vLLM metax baseline on MetaX C550, Qwen3.6 real shapes, bf16 fixed-seed CUDA Graph benchmark, SpeedUp = baseline / FlagGems (334 cases, 3 prefill scenarios x intermediate 128/512 x tokens 1..16384 incl. non-aligned):

scenario	i512 mean	i128 mean	ALL mean
p1024d1024	0.995	1.236	1.115
p4096d1024	0.987	1.244	1.115
p32768d1024	0.996	1.247	1.122

ALL: min 0.715, max 1.708. i128 shapes peak at 1.70 (tokens 40-64) and 1.27 (tokens 16384); i512 shapes average ~0.99. Kernel parameters verified identical to the previous config-key implementation (no behavior/perf change, only the parameterization mechanism).

[CLAassistant]

All committers have signed the CLA.