[4/15] Gate true-on-policy MoE execution

megatron/core/transformer/moe/moe_layer.py

@@ -28,11 +28,20 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.typed_torch import apply_module
 from megatron.core.utils import internal_api
+from miles_megatron_plugins.true_on_policy.moe_layer_ext import (
+    forward_compacted_true_on_policy_padding,
+    run_direct_sglang_ep_forward,
+    uses_true_on_policy_moe_kernel,
+)
 
 try:
     import transformer_engine as te  # pylint: disable=unused-import
 
-    from megatron.core.extensions.transformer_engine import TELinear, te_checkpoint
+    from megatron.core.extensions.transformer_engine import (
+        TELinear,
+        set_save_original_input,
+        te_checkpoint,
+    )
 
     HAVE_TE = True
 except ImportError:
@@ -389,10 +398,18 @@ def forward(
         if padding_mask is not None:
             padding_mask = padding_mask.transpose(0, 1).bool()
 
+        use_true_on_policy_moe_kernel = uses_true_on_policy_moe_kernel(self)
+
         # MoE forward: route -> dispatch -> compute -> combine
         def custom_forward(hidden_states, intermediate_tensors, padding_mask=None):
+            shared_expert_output = None
             try:
                 if "route" in self.fwd_execution_map:
+                    if use_true_on_policy_moe_kernel:
+                        return run_direct_sglang_ep_forward(
+                            self, hidden_states, padding_mask, intermediate_tensors
+                        )
+
                     shared_expert_output = self.shared_experts_compute(hidden_states)
                     probs, routing_map = self.route(hidden_states, padding_mask)
                     hidden_states, probs = self.preprocess(hidden_states, probs, routing_map)
@@ -433,7 +450,21 @@ def custom_forward(hidden_states, intermediate_tensors, padding_mask=None):
 
             return output, mlp_bias
 
-        if self.moe_layer_recompute:
+        use_compact = (
+            use_true_on_policy_moe_kernel
+            and padding_mask is not None
+            and intermediate_tensors is None
+            and padding_mask.dtype == torch.bool
+            and bool(padding_mask.any().item())
+            and not self.use_shared_expert
+            and not self.config.moe_latent_size
+        )
+
+        if use_compact:
+            outputs = forward_compacted_true_on_policy_padding(
+                hidden_states, padding_mask, custom_forward
+            )
+        elif self.moe_layer_recompute:
             if self.config.fp8 or self.config.fp4:
                 outputs = te_checkpoint(
                     custom_forward,
@@ -464,6 +495,4 @@ def set_for_recompute_pre_mlp_layernorm(self):
         # If shared_experts_recompute is used, nothing needs to be done because the checkpoint
         # function will save the original input tensors.
         if self.shared_experts is not None and not self.shared_experts_recompute:
-            from megatron.core.extensions.transformer_engine import set_save_original_input
-
             set_save_original_input(self.shared_experts.linear_fc1)

miles_megatron_plugins/true_on_policy/moe.py

@@ -0,0 +1,18 @@
+"""True-on-policy MoE re-exports.
+
+The implementation is split across:
+  * :mod:`moe_reduce`    -- ``sglang_moe_ep_tree_all_reduce`` deterministic EP combine.
+  * :mod:`moe_experts`   -- ``SGLangGroupedMLP`` SGLang fused-expert helpers.
+  * :mod:`moe_layer_ext` -- EP-invariant local-masked forward and padding compaction.
+
+Importers should prefer the dedicated submodules; this module is kept as a
+backward-compatible re-export point for older call sites.
+"""
+
+from .moe_experts import SGLangGroupedMLP
+from .moe_reduce import sglang_moe_ep_tree_all_reduce
+
+__all__ = [
+    "SGLangGroupedMLP",
+    "sglang_moe_ep_tree_all_reduce",
+]

miles_megatron_plugins/true_on_policy/moe_reduce.py

@@ -0,0 +1,21 @@
+"""Deterministic EP combine for true-on-policy MoE.
+
+Wraps SGLang's ``tree_all_reduce_sum`` so SGLang rollout and Megatron training
+both perform the post-EP sum in the same fixed-tree order. Using this in place
+of NCCL's ring all-reduce is what keeps the EP combine bit-exact across the
+two engines.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def sglang_moe_ep_tree_all_reduce(
+    input_: torch.Tensor,
+    ep_group,
+) -> torch.Tensor:
+    """Mirror SGLang's fixed-order EP sum for true-on-policy MoE combine."""
+    from sglang.srt.tp_invariant_ops import tree_all_reduce_sum
+
+    return tree_all_reduce_sum(input_, device_group=ep_group)