"""Shared capability contract for the external :mod:`ascend_moonep` path."""
from __future__ import annotations
import operator
from typing import Optional, Type
EXTERNAL_MOONEP_MAX_ROUTER_TOPK = 16
EXTERNAL_MOONEP_MAX_EP_RANKS = 1024
EXTERNAL_MOONEP_COMBINE_UB_BYTES = 192 * 1024
def validate_external_moonep_router_topk(
router_topk,
num_experts,
*,
error_type: Type[Exception] = ValueError,
) -> int:
"""Return a validated router top-k supported by the external provider.
``operator.index`` deliberately rejects floats, strings and booleans. A
permissive ``int(value)`` conversion would otherwise turn invalid values
such as ``1.5`` into a different routing contract on only part of the
stack. The upper bound is the smaller of the model's expert count and the
current native combine-kernel limit.
"""
if isinstance(num_experts, bool):
raise error_type("external MoonEP requires a positive integer expert count")
try:
experts = operator.index(num_experts)
except TypeError as error:
raise error_type(
"external MoonEP requires a positive integer expert count"
) from error
if experts <= 0:
raise error_type("external MoonEP requires a positive integer expert count")
if isinstance(router_topk, bool):
raise error_type("external MoonEP router top-k must be a positive integer")
try:
topk = operator.index(router_topk)
except TypeError as error:
raise error_type(
"external MoonEP router top-k must be a positive integer"
) from error
supported_max = min(experts, EXTERNAL_MOONEP_MAX_ROUTER_TOPK)
if topk < 1 or topk > supported_max:
raise error_type(
"external MoonEP router top-k must satisfy "
f"1 <= top-k <= min(num_experts, {EXTERNAL_MOONEP_MAX_ROUTER_TOPK}); "
f"got top-k={topk}, num_experts={experts}"
)
return topk
def validate_external_moonep_group_routing(
router_topk,
num_experts,
router_num_groups,
router_group_topk,
*,
error_type: Type[Exception] = ValueError,
) -> tuple[Optional[int], Optional[int]]:
"""Return the normalized group-limited routing contract.
Megatron enables group-limited routing only when ``group_topk`` is
truthy. A standalone ``num_groups`` value is therefore inert and is
normalized out of the external runtime contract. When enabled, validate
every bound used by ``group_limited_topk`` before planning or VMM
allocation; in particular, its internal ``topk // group_topk`` must be at
least one and the selected groups must contain enough experts.
"""
topk = validate_external_moonep_router_topk(
router_topk, num_experts, error_type=error_type
)
experts = operator.index(num_experts)
if not router_group_topk:
return None, None
if isinstance(router_num_groups, bool) or isinstance(router_group_topk, bool):
raise error_type(
"external MoonEP group routing requires positive integer "
"moe_router_num_groups and moe_router_group_topk"
)
try:
num_groups = operator.index(router_num_groups)
group_topk = operator.index(router_group_topk)
except TypeError as error:
raise error_type(
"external MoonEP group routing requires positive integer "
"moe_router_num_groups and moe_router_group_topk"
) from error
if (
num_groups <= 0
or group_topk <= 0
or group_topk > num_groups
or experts % num_groups
):
raise error_type(
"external MoonEP group routing requires num_experts divisible by "
"moe_router_num_groups and 1 <= moe_router_group_topk <= "
"moe_router_num_groups"
)
if topk < group_topk:
raise error_type(
"external MoonEP group routing requires router top-k greater than "
f"or equal to moe_router_group_topk; got top-k={topk}, "
f"group_topk={group_topk}"
)
selectable_experts = experts // num_groups * group_topk
if topk > selectable_experts:
raise error_type(
"external MoonEP router top-k exceeds the experts available from "
f"the selected groups: top-k={topk}, "
f"selectable_experts={selectable_experts}"
)
return num_groups, group_topk
def validate_external_moonep_router_semantics(
router_topk,
*,
score_function="softmax",
pre_softmax=False,
load_balancing_type="aux_loss",
error_type: Type[Exception] = ValueError,
) -> None:
"""Mirror MCore's K=1 differentiability contract.
This is intentionally not a blanket pre-softmax requirement: sinkhorn and
non-softmax scoring are valid upstream alternatives. Keeping the same
condition here makes the early external-runtime gate agree with the
TransformerConfig gate that follows it.
"""
if (
operator.index(router_topk) == 1
and score_function == "softmax"
and not bool(pre_softmax)
and load_balancing_type != "sinkhorn"
):
raise error_type(
"external MoonEP router top-k 1 requires pre-softmax routing, "
"sinkhorn load balancing, or a non-softmax score function"
)
def validate_external_moonep_native_dimensions(
*,
hidden_size,
num_experts,
num_ep_ranks,
error_type: Type[Exception] = ValueError,
) -> None:
"""Reject dimensions unsupported by the frozen planning/combine ABI."""
hidden = operator.index(hidden_size)
experts = operator.index(num_experts)
ranks = operator.index(num_ep_ranks)
if ranks > EXTERNAL_MOONEP_MAX_EP_RANKS:
raise error_type(
"external MoonEP planning supports at most "
f"{EXTERNAL_MOONEP_MAX_EP_RANKS} EP ranks; got {ranks}"
)
if ranks * experts > 2**31 - 1:
raise error_type(
"external MoonEP EP * num_experts must fit signed int32 planning indices"
)
if hidden % 16:
raise error_type("external MoonEP combine requires hidden size divisible by 16")
if hidden * 10 + 256 > EXTERNAL_MOONEP_COMBINE_UB_BYTES:
raise error_type(
"external MoonEP hidden size exceeds the single-pass combine UB limit"
)