[7/15] Complete true-on-policy MoE EP forward orchestration

miles_megatron_plugins/true_on_policy/moe_layer_ext.py

@@ -137,3 +137,165 @@ def backward(ctx, grad_output):
         return grad, None
 
 
+def _forward_sglang_ep(moe_layer, hidden_states: torch.Tensor):
+    """gather -> route -> forward -> reduce -> slice"""
+    hidden_shape = hidden_states.shape
+    flat = hidden_states.reshape(-1, hidden_shape[-1])
+    ep_group = moe_layer.token_dispatcher.ep_group
+    ep_size = moe_layer.token_dispatcher.ep_size
+    ep_rank = torch.distributed.get_rank(group=ep_group)
+
+    max_num_tokens, token_counts = _gather_ep_token_counts(
+        flat.shape[0], flat.device, ep_group, ep_size
+    )
+    if max_num_tokens == 0:
+        return torch.zeros_like(hidden_states), None
+
+    # Step 1: EP all-gather
+    if torch.is_grad_enabled():
+        global_hidden = _PaddedEPAllGather.apply(
+            flat, max_num_tokens, tuple(token_counts), ep_group, ep_rank, ep_size
+        )
+    else:
+        global_hidden = torch.cat(
+            _all_gather_padded(flat, max_num_tokens, ep_group, ep_size), dim=0
+        )
+
+    # Step 2: Router
+    topk_weights, topk_ids = _sglang_topk_route(moe_layer, global_hidden, hidden_shape[-1])
+    if topk_weights is None:
+        return None
+
+    # Step 3: Router grad ownership (training only)
+    if torch.is_grad_enabled():
+        local_start = ep_rank * max_num_tokens
+        local_count = token_counts[ep_rank]
+        topk_weights = _with_local_router_grad_owner(topk_weights, local_start, local_count)
+
+    # Step 4: Expert forward
+    w1 = moe_layer.experts.sglang_w13_weight()
+    w2 = moe_layer.experts.sglang_w2_weight()
+
+    if torch.is_grad_enabled():
+        global_output = sglang_fused_experts_autograd(
+            layer_number=moe_layer.layer_number,
+            hidden_states=global_hidden,
+            w1=w1, w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            num_experts=moe_layer.config.num_moe_experts,
+            num_local_experts=moe_layer.num_local_experts,
+            ep_rank=ep_rank, ep_size=ep_size, ep_group=ep_group,
+            activation="silu",
+            allreduce_grad_hidden=False,
+        )
+    else:
+        global_output = sglang_moe_forward(
+            hidden_states=global_hidden,
+            w1=w1, w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            num_experts=moe_layer.config.num_moe_experts,
+            num_local_experts=moe_layer.num_local_experts,
+            ep_rank=ep_rank, ep_size=ep_size,
+            activation="silu",
+            layer_id=moe_layer.layer_number,
+        )
+
+    # Step 5: EP reduce
+    if torch.is_grad_enabled():
+        global_output = _SGLangEPAllReduceSum.apply(global_output, ep_group)
+    else:
+        global_output = sglang_moe_ep_tree_all_reduce(global_output, ep_group)
+
+    # Step 6: Slice local tokens
+    local_start = ep_rank * max_num_tokens
+    local_count = token_counts[ep_rank]
+    local_output = global_output[local_start : local_start + local_count].contiguous()
+
+    if local_count < flat.shape[0]:
+        padded_output = hidden_states.new_zeros(flat.shape)
+        padded_output[:local_count] = local_output
+        local_output = padded_output
+
+    return local_output.view(hidden_shape), None
+
+
+# ---------------------------------------------------------------------------
+# Router
+# ---------------------------------------------------------------------------
+
+def _with_local_router_grad_owner(topk_weights, local_start, local_count):
+    """Keep global routing values, but route gradients only through local tokens."""
+    if local_count == topk_weights.shape[0]:
+        return topk_weights
+    mask = torch.zeros((topk_weights.shape[0], 1), dtype=torch.bool, device=topk_weights.device)
+    mask[local_start : local_start + local_count] = True
+    return torch.where(mask, topk_weights, topk_weights.detach())
+
+
+def _sglang_topk_route(moe_layer, global_hidden, hidden_size):
+    """Compute SGLang-ordered top-k routing. Returns (weights, ids) or (None, None)."""
+    config = moe_layer.config
+    router = moe_layer.router
+    policy = resolve_true_on_policy_runtime_policy(config)
+
+    if not policy.deterministic_moe_routing or policy.moe_topk_tiebreak != "stable_sort":
+        return None, None
+    if stable_topk_softmax is None:
+        return None, None
+    if getattr(router, "routing_type", None) == "sinkhorn":
+        return None, None
+    if getattr(router, "score_function", None) != "softmax":
+        return None, None
+    if getattr(config, "moe_router_pre_softmax", False):
+        return None, None
+    if getattr(config, "moe_router_num_groups", None) is not None:
+        return None, None
+    if getattr(config, "moe_router_fusion", False):
+        return None, None
+    if getattr(router, "expert_bias", None) is not None:
+        return None, None
+    if getattr(config, "moe_expert_capacity_factor", None) is not None:
+        return None, None
+
+    router._maintain_float32_expert_bias()
+    router_weight = router.weight
+    if router_weight.device.type == "cpu" and global_hidden.is_cuda:
+        router_weight.data = router_weight.data.to(device=torch.cuda.current_device())
+    bias = getattr(router, "bias", None)
+    if bias is not None and bias.device.type == "cpu" and global_hidden.is_cuda:
+        bias.data = bias.data.to(device=torch.cuda.current_device())
+
+    router_input = global_hidden.to(config.params_dtype).view(-1, 1, hidden_size)
+    logits = router_gating_linear(
+        router_input, router_weight, bias, config.params_dtype
+    ).view(-1, config.num_moe_experts)
+    logits = router.apply_z_loss(logits, padding_mask=None)
+
+    topk_weights, topk_ids = stable_topk_softmax(logits, config.moe_router_topk)
+    topk_scaling = getattr(config, "moe_router_topk_scaling_factor", None)
+    if topk_scaling:
+        topk_weights = topk_weights * topk_scaling
+    return topk_weights.type_as(logits), topk_ids
+
+
+# ---------------------------------------------------------------------------
+# EP collective helpers
+# ---------------------------------------------------------------------------
+
+def _gather_ep_token_counts(local_num_tokens, device, ep_group, ep_size):
+    local_count = torch.tensor([local_num_tokens], device=device, dtype=torch.long)
+    gathered = [torch.empty_like(local_count) for _ in range(ep_size)]
+    torch.distributed.all_gather(gathered, local_count, group=ep_group)
+    counts = [int(c.item()) for c in gathered]
+    return max(counts), counts
+
+
+def _all_gather_padded(local_tensor, max_num_tokens, ep_group, ep_size):
+    padded = local_tensor.new_zeros((max_num_tokens, *local_tensor.shape[1:]))
+    if local_tensor.shape[0] != 0:
+        padded[: local_tensor.shape[0]] = local_tensor
+    gathered = [torch.empty_like(padded) for _ in range(ep_size)]
+    torch.distributed.all_gather(gathered, padded, group=ep_group)
+    return gathered