Problem

In hyper_parallel/platform/torch/common/moe.py, function _compute_load_balance_loss computed:

• expert_fraction with shape [num_experts]
• mean_score with shape [top_k]
• then multiplied via broadcasting and summed all elements

This collapses to a near-constant value and does not reflect routing imbalance severity, which is inconsistent with the expected behavior of MoE load-balance auxiliary loss.

Reproduction (simplified)

Use the same top_scores but different selected_experts distributions (balanced vs. heavily imbalanced).
The old implementation returns effectively constant behavior; loss does not track imbalance correctly.

Expected

Loss should vary with routing distribution and penalize imbalanced expert usage.

Fix

1. Compute expert_fraction from normalized assignment counts.
2. Compute expert_prob by accumulating per-token top-k probabilities onto selected experts.
3. Use num_experts * sum(expert_fraction * expert_prob) as scalar auxiliary loss.
4. Add unit test to verify imbalanced routing gets larger loss than balanced routing.

Changed files

• hyper_parallel/platform/torch/common/moe.py
• tests/ut/platform/torch/common/test_moe.py