Skip to content

Make MoE router-weight placement explicit in experts - #79

Draft
zianglih wants to merge 3 commits into
radixark:miles-mainfrom
zianglih:agent/moe-apply-probs-on-output
Draft

Make MoE router-weight placement explicit in experts#79
zianglih wants to merge 3 commits into
radixark:miles-mainfrom
zianglih:agent/moe-apply-probs-on-output

Conversation

@zianglih

@zianglih zianglih commented Aug 6, 2026

Copy link
Copy Markdown

Summary

@HumansAnd

  • Add an explicit MoERouterWeightPlacement expert API with FC2_INPUT and FC2_OUTPUT modes.
  • Keep the historical FC2-input path as the default; MEGATRON_MOE_APPLY_PROBS_ON_OUTPUT=1 selects FC2-output placement.
  • Implement the contract in GroupedMLP, TEGroupedMLP, and SequentialMLP, while preserving the legacy three-argument call for disabled/custom-expert paths.
  • Keep the one-time rank-zero marker so end-to-end runs can confirm that the opt-in branch was taken.

Numerical contract

The FC2-output mode no longer passes unit probabilities through the weighted activation path. It runs SwiGLU and FC2 unweighted, promotes the materialized FC2 output, expert bias, and original router probability to at least FP32, computes (fc2_output + bias) * router_prob, and casts the result once to the expert output dtype.

This removes the unit-probability allocation and activation-sized multiply/reduction, and it avoids narrowing the router probability or creating a probability-weighted BF16 FC2 input.

This is not a grouped-GEMM epilogue: TE has already materialized and rounded the FC2 output to the expert output dtype. Applying the router probability to the GEMM accumulator would require a future TE/kernel API extension and is intentionally outside this PR.

Scope

  • FC2-output placement currently supports SwiGLU only; other activations are rejected.
  • Expert tensor parallelism greater than 1 is rejected because weighting rank-local FC2 partials would not provide the intended post-FC2 numerical contract.
  • moe_apply_probs_on_input and moe_combine_in_fp32 remain incompatible with FC2-output placement.
  • Quantized expert padding keeps router probabilities unpadded and applies them after output unpadding. Quantized experts with bias are rejected because bias-generated padded activations can contaminate quantization metadata.
  • The default FC2-input path remains unchanged.

Validation

  • git diff --check
  • python3 -m py_compile on all changed Python files
  • Black and isort hooks on the changed files
  • B200 focused tests: explicit fixed-weight SwiGLU algebra and router-probability gradients for both placements, plus Sequential-vs-TE forward/backward parity for FC2-input/output with fused and unfused SwiGLU (5 passed)
  • B200 MXFP8 forward/backward smoke with uneven per-expert token counts, a zero-token expert, and quantization padding
  • B200 MXFP8 activation-recompute forward/backward smoke
  • B200 BF16 expert-bias forward/backward smoke

GPU validation used radixark/miles@sha256:d9e01378d8820afd88824c798ea628b3b6cb87a6c6911db5165ef9d98187db55 without reinstalling packages.

@zianglih zianglih changed the title Add post-expert MoE probability weighting Make MoE router-weight placement explicit in experts Aug 8, 2026
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