Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/distributed/expert

3# Licensed under the Apache License, Version 2.0 (the "License");

4# you may not use this file except in compliance with the License.

5# You may obtain a copy of the License at

7# http://www.apache.org/licenses/LICENSE-2.0

9# Unless required by applicable law or agreed to in writing, software

10# distributed under the License is distributed on an "AS IS" BASIS,

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

12# See the License for the specific language governing permissions and

13# limitations under the License.

14# ============================================================================

16"""expert_parallel.routing: MoE router adapters (05 §6.4.7, D-09).

18Routing semantics are model-specific (softmax/sigmoid, top-k normalization,

19scaling), while the expert MLP structure is uniform (SwiGLU). A factory picks

20an adapter from ``MOE_ROUTER_ADAPTERS`` BY NAME in its own code — the choice

21is explicit, never inferred. Each adapter maps

22``(module, hidden) -> (topk_idx [T,K] int64, topk_w [T,K] float)``.

24Benchmark-only balanced variants are derived from a real adapter with

25:func:`make_balanced_router` and registered as ``<adapter>_fixed``, so an

26archetype opts in with one flag (``fix_router``, see :func:`resolve_moe_router`)

27instead of carrying its own copy of the wrapper.

29Split out of components/distributed/ep_utils.py in stage 4e.

30"""

32import logging

33import math

34import os

35from typing import Optional

37import torch

39logger = logging.getLogger(__name__)

41# float32 represents every integer up to 2**24 exactly, so a smaller expert count

42# can use the AICore float sort (see ``experts._SORT_KEY_FP32_LIMIT``).

43_SORT_KEY_FP32_LIMIT = 1 << 24

45# Which slots inside an over-capacity expert survive (``HP_EP_CAPACITY_ORDER``):

46#

47# * ``score`` (default): the heaviest routing weights -- the contributions whose removal

48# perturbs the routed sum the least. A cap exists to bound memory, and spending the bound

49# on the cheapest slots is what makes "which slots die" a defined policy instead of an

50# accident of the sort order;

51# * ``first``: the first slots of the expert-major order, i.e. what the archived

52# ``HP_EP_CAPACITY_FACTOR`` sweeps (drop rate / peak / step time) were measured with.

53#

54# Both policies keep exactly the same number of slots per expert, which is the invariant the

55# memory bound rests on, so the choice cannot move a peak, a row count or a step time -- it

56# only decides whose contribution is lost.

57_CAPACITY_ORDER_SCORE = "score"

58_CAPACITY_ORDER_FIRST = "first"

59_CAPACITY_ORDERS = (_CAPACITY_ORDER_SCORE, _CAPACITY_ORDER_FIRST)

60_FP32_KEY_WARNED = False

63def _softmax_topk_router(module, hidden_states):

64 """default adapter: softmax -> topk -> normalize by sum (Mixtral/Qwen3 semantics).

66 top_k source: config.num_experts_per_tok / config.top_k / module.top_k (default 2);

67 normalization switch: config.norm_topk_prob (default True).

68 """

69 gate = getattr(module, "gate", None)

70 if gate is None:

71 gate = getattr(module, "router", None)

72 if gate is None:

73 raise AttributeError(

74 f"{type(module).__name__}: router not found (neither gate nor router "

75 "attribute exists); please register a custom MOE_ROUTER_ADAPTERS entry"

76 )

77 cfg = getattr(module, "config", None)

78 logits = gate(hidden_states)

79 logits = logits.view(-1, logits.shape[-1])

80 top_k = (getattr(cfg, "num_experts_per_tok", None)

81 or getattr(cfg, "top_k", None)

82 or getattr(module, "top_k", 2))

83 weights = logits.softmax(-1)

84 topk_w, topk_idx = weights.topk(int(top_k), dim=-1)

85 if getattr(cfg, "norm_topk_prob", True):

86 topk_w = topk_w / topk_w.sum(-1, keepdim=True).clamp_min(1e-20)

87 return topk_idx, topk_w

90def _topk_router_module(module, hidden_states):

91 """Qwen2/Qwen3/Mixtral adapter: gate is a TopKRouter module (after the HF 2025

92 refactor); forward directly returns (logits, scores [T,K], indices [T,K])

93 -- take the latter two."""

94 gate = getattr(module, "gate", None)

95 if gate is None:

96 gate = getattr(module, "router", None)

97 out = gate(hidden_states)

98 if isinstance(out, (tuple, list)) and len(out) == 3:

99 _, scores, indices = out

100 return indices, scores

101 raise TypeError(

102 f"{type(module).__name__}: TopKRouter should return (logits, scores, indices), "

103 f"got {type(out).__name__} -- use the default adapter or register a custom adapter"

104 )

105

106

107def _mask_scores_by_group(scores, n_group, topk_group):

108 """Keep scores only in the highest-scoring expert groups."""

109 expert_count = scores.shape[-1]

110 group_scores = (

111 scores.view(-1, n_group, expert_count // n_group)

112 .topk(2, dim=-1)[0]

113 .sum(dim=-1)

114 )

115 group_indices = group_scores.topk(topk_group, dim=-1, sorted=False)[1]

116 group_mask = torch.zeros_like(group_scores)

117 group_mask.scatter_(1, group_indices, 1)

118 score_mask = (

119 group_mask.unsqueeze(-1)

120 .expand(-1, n_group, expert_count // n_group)

121 .reshape(-1, expert_count)

122 )

123 return scores.masked_fill(~score_mask.bool(), float("-inf"))

124

125

126def _sigmoid_group_router(module, hidden_states):

127 """deepseekv3/glm4moe adapter: sigmoid + e_score_correction_bias +

128 group-limited topk + (optional) normalization + routed_scaling_factor

129 (step-by-step consistent with HF DeepseekV3MoE.route_tokens_to_experts /

130 Glm4MoeMoE).

131

132 Parameter source: the module's own attributes take precedence

133 (n_group/topk_group/top_k/norm_topk_prob/routed_scaling_factor), falling

134 back to module.config; when n_group is missing or <=1, the group-limited

135 filter is skipped.

136 """

137 gate = getattr(module, "gate", None)

138 if gate is None:

139 gate = getattr(module, "router", None)

140 cfg = getattr(module, "config", None)

141

142 def _attr(name, default=None):

143 v = getattr(module, name, None)

144 if v is None and cfg is not None:

145 v = getattr(cfg, name, None)

146 return default if v is None else v

147

148 logits = gate(hidden_states)

149 if isinstance(logits, (tuple, list)):

150 logits = logits[0]

151 logits = logits.view(-1, logits.shape[-1]).float()

152 scores = logits.sigmoid()

153 bias = getattr(gate, "e_score_correction_bias", None)

154 scores_for_choice = scores + bias if bias is not None else scores

155

156 n_group = int(_attr("n_group", 0) or 0)

157 topk_group = int(_attr("topk_group", 0) or 0)

158 top_k = int(_attr("top_k", None) or _attr("num_experts_per_tok", 2))

159 if n_group > 1 and topk_group > 0:

160 scores_for_choice = _mask_scores_by_group(

161 scores_for_choice, n_group, topk_group

162 )

163

164 topk_idx = scores_for_choice.topk(top_k, dim=-1, sorted=False)[1]

165 topk_w = scores.gather(1, topk_idx)

166 if _attr("norm_topk_prob", False):

167 topk_w = topk_w / (topk_w.sum(dim=-1, keepdim=True) + 1e-20)

168 topk_w = topk_w * float(_attr("routed_scaling_factor", 1.0))

169 return topk_idx, topk_w

170

171

172def apply_capacity_limit(

173 topk_idx: "torch.Tensor",

174 capacity_factor: Optional[float],

175 num_experts: int,

176 weights: Optional["torch.Tensor"] = None,

177) -> tuple["torch.Tensor", "torch.Tensor"]:

178 """Capacity mask for a routed assignment (MoE token dropping).

179

180 Real routing is skewed: on the 2-SN benchmark the busiest rank has been measured at

181 **6x** the mean token count. A rank's expert buffers are sized by the rows it

182 receives, so that skew -- not the balanced case -- decides whether the step fits in

183 memory: at GBS 256 / seq 8192 the balanced peak is 41.5 GB while real routing reached

184 55.3 GiB and failed to allocate a 5.23 GiB expert output. Every *switch* that does not

185 change the routed row count (GBS, prefetch depth, ``swap_inputs``, ``overlap_shared_expert``,

186 even a 25% shorter sequence) bought at most 1.8 GB, because none of them bound the

187 busiest rank.

188

189 Capping each expert at ``capacity_factor x (T*K / E)`` slots bounds it directly, which is

190 the standard GShard/Switch/Megatron answer to routing skew. Tokens beyond capacity are

191 *dropped* (their contribution is removed), so this changes the model, exactly as it does

192 in those implementations; the returned drop rate is what makes the trade measurable.

193

194 The cap is what bounds memory, and it is policy-independent: every expert keeps exactly

195 ``min(count_e, capacity)`` slots. *Which* slots those are is decided by

196 ``HP_EP_CAPACITY_ORDER``:

197

198 * ``score`` (default): the ``capacity`` heaviest ``weights`` of each expert, so the

199 contributions dropped are the cheapest ones. Cheap and deterministic -- the expert index

200 and the inverted weight share a single fp32 sort key, keeping the AICore float sort (see

201 :func:`_expert_score_order`);

202 * ``first``: the first slots of the expert-major order (see :func:`_expert_major_order`).

203 Deterministic for a given input and backend as well, but not a policy with a meaning:

204 within one expert it is "not necessarily first in token order".

205

206 Args:

207 topk_idx: Routed expert per slot, shape ``[T, K]``.

208 capacity_factor: Multiplier on the mean load, or ``None`` to disable the cap.

209 num_experts: Total routed experts ``E``.

210 weights: Routing weight of every slot, shaped like ``topk_idx``. Only the ``score``

211 policy reads it, so a caller that passes ``None`` gets the ``first`` policy.

212

213 Returns:

214 ``(keep, dropped)`` where ``keep`` is a bool tensor shaped like ``topk_idx`` and

215 ``dropped`` is the number of dropped slots as a 0-dim tensor on the plan's device.

216 Draining it costs a host sync, so callers must not read it on the hot path -- the one

217 caller logs it once and then sparsely (see ``ep_routed_dispatch``).

218

219 Raises:

220 ValueError: If ``num_experts`` is not positive, or ``HP_EP_CAPACITY_ORDER`` is neither

221 ``score`` nor ``first``.

222 """

223 if capacity_factor is None or float(capacity_factor) <= 0:

224 return (

225 torch.ones_like(topk_idx, dtype=torch.bool),

226 torch.zeros((), dtype=torch.int64, device=topk_idx.device),

227 )

228 if num_experts <= 0:

229 raise ValueError(f"num_experts must be positive, got {num_experts}")

230

231 flat_expert = topk_idx.reshape(-1).to(torch.int64)

232 slots = flat_expert.numel()

233 capacity = int(math.ceil(float(capacity_factor) * slots / float(num_experts)))

234 keep_flat = torch.zeros_like(flat_expert, dtype=torch.bool)

235 if capacity > 0:

236 order = _slot_order(flat_expert, weights, num_experts)

237 counts = torch.bincount(flat_expert, minlength=num_experts)

238 starts = torch.cumsum(counts, dim=0) - counts

239 ranks = torch.arange(slots, device=flat_expert.device) - starts[flat_expert[order]]

240 keep_flat[order] = ranks < capacity

241 # ``dropped`` stays a device tensor on purpose: draining it costs a host sync, and this

242 # runs once per MoE layer, so callers that want the number must ask for it explicitly

243 # (they log it once, see ``ep_routed_dispatch``) instead of paying it every layer.

244 return keep_flat.view_as(topk_idx), (~keep_flat).sum()

245

246

247def _resolve_capacity_order() -> str:

248 """Resolve ``HP_EP_CAPACITY_ORDER`` for the current call.

249

250 Read per call rather than at import time, like the capacity factor itself, so the policy

251 can be switched without reloading the module.

252

253 Returns:

254 ``"score"`` or ``"first"``.

255

256 Raises:

257 ValueError: If the knob is neither; a typo must not silently keep a policy the

258 operator did not ask for.

259 """

260 raw = os.environ.get("HP_EP_CAPACITY_ORDER", _CAPACITY_ORDER_SCORE).strip().lower()

261 if raw not in _CAPACITY_ORDERS:

262 raise ValueError(

263 f"HP_EP_CAPACITY_ORDER must be one of {list(_CAPACITY_ORDERS)}, but got {raw!r}"

264 )

265 return raw

266

267

268def _slot_order(

269 flat_expert: "torch.Tensor",

270 flat_weights: Optional["torch.Tensor"],

271 num_experts: int,

272) -> "torch.Tensor":

273 """Order slots by expert, and under the ``score`` policy by descending weight as well."""

274 if flat_weights is None or _resolve_capacity_order() == _CAPACITY_ORDER_FIRST:

275 return _expert_major_order(flat_expert, num_experts)

276 return _expert_score_order(flat_expert, flat_weights, num_experts)

277

278

279def _expert_score_order(

280 flat_expert: "torch.Tensor",

281 flat_weights: "torch.Tensor",

282 num_experts: int,

283) -> "torch.Tensor":

284 """Expert-major order that puts the heaviest slot of every expert first.

285

286 Both keys travel in ONE fp32 sort key: the integer part is the expert index and the

287 fractional part carries the inverted, normalized weight, so a heavier slot always sorts

288 earlier and the fraction can never carry into the next expert. One key also means one sort

289 -- the AICore float sort the expert-major order already uses. A lexicographic or int64 key

290 would lose that kernel (the same reason ``HP_EP_SORT_FP32`` exists for the dispatch sort),

291 and a second pass is not needed either.

292

293 Resolution limit: fp32 spends the integer part's exponent on the expert index, so the

294 fractional step is about ``2**-15`` around ``num_experts = 2**8..2**9`` (and ``2**-11``

295 around ``2**12``). Weights closer together than that, relative to the plan's spread, tie

296 and their relative order is arbitrary -- the *count* kept per expert, which is what bounds

297 the busiest rank, is exact either way. ``HP_EP_CAPACITY_ORDER=first`` stays available for a

298 bit-exact reproduction of the archived sweeps.

299

300 Args:

301 flat_expert: Expert index of every routed slot, shape ``[T*K]``.

302 flat_weights: Routing weight of every slot, same shape.

303 num_experts: Size of the expert index space.

304

305 Returns:

306 An int64 permutation of ``range(flat_expert.numel())``.

307

308 Raises:

309 ValueError: If ``flat_weights`` does not carry exactly one weight per slot.

310 """

311 if flat_weights.numel() != flat_expert.numel():

312 raise ValueError(

313 f"weights must have one entry per slot: got {flat_weights.numel()} for "

314 f"{flat_expert.numel()} slots"

315 )

316 if not 0 < num_experts <= _SORT_KEY_FP32_LIMIT:

317 global _FP32_KEY_WARNED # pylint: disable=global-statement

318 if not _FP32_KEY_WARNED:

319 _FP32_KEY_WARNED = True

320 logger.warning(

321 "HP_EP_CAPACITY_ORDER=score needs the expert index to be exact in a float32 "

322 "key, but num_experts=%d is outside (0, 2**24]; keeping the expert-major order",

323 num_experts,

324 )

325 return _expert_major_order(flat_expert, num_experts)

326 weight = flat_weights.reshape(-1).to(torch.float32)

327 inverted = weight.max() - weight

328 # ``clamp_min`` rather than a host-side ``if span == 0``: reading the span would drain the

329 # device, and this runs once per MoE layer. A zero span leaves every ``inverted`` at zero,

330 # so the fraction stays zero and the expert-major order stands. ``0 * inf`` cannot happen:

331 # a zero span means ``inverted`` is exactly zero.

332 span = inverted.max().clamp_min(torch.finfo(weight.dtype).tiny)

333 # One multiply-add instead of three elementwise passes: ``expert + inverted * 0.5 / span``.

334 key = torch.addcmul(flat_expert.to(torch.float32), inverted, 0.5 / span)

335 return key.argsort()

336

337

338def _expert_major_order(flat_expert: "torch.Tensor", num_experts: int) -> "torch.Tensor":

339 """Return the permutation that groups slots by expert.

340

341 Integer ``argsort`` has no AICore kernel on this stack (the same reason

342 ``HP_EP_SORT_FP32`` exists for the dispatch sort), so keys are cast to float32 when the

343 expert count is exactly representable. Ties inside one expert may then be ordered

344 arbitrarily rather than by token: the *count* kept per expert -- which is what bounds the

345 busiest rank's memory -- is unaffected, and the result is still deterministic for a given

346 input and backend.

347

348 Args:

349 flat_expert: Expert index of every routed slot, shape ``[T*K]``.

350 num_experts: Size of the expert index space.

351

352 Returns:

353 An int64 permutation of ``range(flat_expert.numel())`` ordering slots by expert.

354 """

355 if 0 < num_experts <= _SORT_KEY_FP32_LIMIT:

356 return flat_expert.to(torch.float32).argsort()

357 return torch.argsort(flat_expert, stable=True)

358

359

360def _global_expert_count(module):

361 """Model-level routed expert count (mirrors ``experts._get_global_expert_count``).

362

363 Kept local so this module stays a torch-only leaf (it is imported by

364 model adapters that do not want the expert compute machinery).

365 """

366 for owner in (getattr(module, "experts", None), module):

367 count = getattr(owner, "num_experts", None)

368 if count is not None:

369 return int(count)

370 cfg = getattr(module, "config", None)

371 for name in ("num_experts", "n_routed_experts"):

372 count = getattr(cfg, name, None)

373 if count is not None:

374 return int(count)

375 raise ValueError(

376 f"{type(module).__name__}: cannot determine the global routed expert count"

377 )

378

379

380def _balanced_expert_indices(

381 token_count: int,

382 experts_per_token: int,

383 expert_count: int,

384 local_count: int,

385 *,

386 device,

387) -> "torch.Tensor":

388 """Round-robin the ``T*K`` slots over the DESTINATION ranks.

389

390 Local slot ``i`` goes to one of the ``L`` experts owned by destination

391 ``i mod Q`` (``Q = E/L``), so every rank receives exactly ``T*K/Q`` slots

392 instead of a data-dependent count.

393

394 Spreading by *rank* rather than by expert is what makes this exact: the

395 per-expert token count cannot be made uniform at all when ``E`` does not

396 divide ``T*K``, and a per-expert round robin (``i mod E``) replays the same

397 residue pattern on every rank, so the leftover slots pile onto the same

398 fixed group of ranks (measured: 0.8% above the mean on the busiest rank,

399 2.4% spread between rank groups). Going through the destination first

400 absorbs that leftover per rank instead.

401

402 Args:

403 token_count: Local tokens ``T``.

404 experts_per_token: Routed experts per token ``K``.

405 expert_count: Global routed experts ``E``.

406 local_count: Experts this rank owns ``L``.

407 device: Device for the returned index tensor.

408

409 Returns:

410 ``[T, K]`` int64 indices naming the expert each slot is routed to.

411

412 Raises:

413 ValueError: If ``E`` is not divisible by ``L``.

414 """

415 if local_count <= 0 or expert_count % local_count != 0:

416 raise ValueError(

417 f"num_experts ({expert_count}) must be divisible by the local "

418 f"expert count ({local_count})"

419 )

420 destinations = expert_count // local_count

421 slots = torch.arange(token_count * experts_per_token, device=device)

422 indices = (slots % destinations) * local_count + (slots // destinations) % local_count

423 return indices.view(token_count, experts_per_token)

424

425

426def make_balanced_router(real_router):

427 """Wrap ``real_router`` in a benchmark-only variant with a uniform expert load.

428

429 The real adapter still runs -- so its weights stay differentiable and the

430 router GEMM stays in the compute profile -- but the chosen experts are

431 replaced by a round-robin assignment that spreads the local slots over the

432 **destination ranks** (see :func:`_balanced_expert_indices`). Every rank

433 therefore receives exactly ``T*K/Q`` tokens instead of a data-dependent

434 count, whatever routing semantics ``real_router`` implements.

435

436 Why: unbalanced routing is the dominant source of step-time jitter on MoE

437 training runs -- the slowest rank in a step is the one whose experts drew

438 the most tokens, so identical configurations can differ by tens of percent

439 step to step. Measured on the 293B/18-layer config, the real router put

440 **6.3x** the mean token count on the busiest rank (leaving some experts

441 empty) while the balanced one stayed at 1.00x. Flattening the load makes a

442 step-to-step comparison converge in 2-3 steps instead of needing a dozen,

443 which is what makes small (1-3%) kernel-level effects measurable at all.

444

445 **Benchmark-only.** The assignment ignores the gate scores, so:

446

447 - ``loss`` / ``grad_norm`` / any convergence signal is meaningless here;

448 - expert hit distributions and "effective MFU" are optimistic (the load is

449 artificially flat), so numbers from this mode must not be reported as an

450 efficiency result.

451

452 Correctness must be validated with the switch OFF (``fix_router: False``),

453 where ``real_router`` runs unchanged. The true ``topk_w`` are kept so the

454 gate still receives gradient.

455

456 Args:

457 real_router: A ``(module, hidden) -> (topk_idx, topk_w)`` adapter, e.g.

458 any entry of :data:`MOE_ROUTER_ADAPTERS`.

459

460 Returns:

461 A ``(module, hidden) -> (topk_idx, topk_w)`` adapter with a uniform

462 destination load; its ``__name__`` records the wrapped adapter.

463 """

464

465 def balanced_router(module, hidden_states):

466 """Run ``real_router`` and replace its destinations round-robin."""

467 topk_idx, topk_w = real_router(module, hidden_states)

468 token_count, experts_per_token = topk_idx.shape

469 balanced_idx = _balanced_expert_indices(

470 int(token_count),

471 int(experts_per_token),

472 _global_expert_count(module),

473 # Set by the EP binder (bind_local_expert_forward) before any forward; a

474 # missing value degrades to a plain per-expert round robin.

475 int(getattr(getattr(module, "experts", None),

476 "local_expert_count", None) or 1),

477 device=topk_idx.device,

478 )

479 return balanced_idx.to(topk_idx.dtype), topk_w

480

481 # A callable object or a functools.partial has no ``__name__``; its class name is

482 # then the most informative thing the log line can carry.

483 real_name = getattr(real_router, "__name__", None) or type(real_router).__name__

484 balanced_router.__name__ = f"balanced_{real_name}"

485 balanced_router.__doc__ = (

486 f"Benchmark-only variant of ``{real_name}`` with a uniform destination "

487 "load (see :func:`make_balanced_router`)."

488 )

489 return balanced_router

490

491

492MOE_ROUTER_ADAPTERS = {

493 "softmax_topk": _softmax_topk_router,

494 "topk_router_module": _topk_router_module,

495 "sigmoid_group": _sigmoid_group_router,

496 "default": _softmax_topk_router,

497 "qwen2moe": _topk_router_module,

498 "qwen2_moe": _topk_router_module,

499 "qwen3moe": _topk_router_module,

500 "qwen3_moe": _topk_router_module,

501 "mixtral": _topk_router_module,

502 "deepseekv3": _sigmoid_group_router,

503 "deepseek_v3": _sigmoid_group_router,

504 "glm4moe": _sigmoid_group_router,

505 "glm4_moe": _sigmoid_group_router,

506}

507

508# One balanced wrapper per distinct real adapter, published as ``<name>_fixed`` for

509# every alias: the convention ``fix_router`` resolves through.

510_BALANCED_VARIANTS = {

511 adapter: make_balanced_router(adapter)

512 for adapter in dict.fromkeys(MOE_ROUTER_ADAPTERS.values())

513}

514# The archived benchmark router: the balanced wrapper over the sigmoid-group adapter,

515# kept as a module-level name because the recipes' docs and the earlier 2-SN sweeps

516# refer to it.

517_balanced_router = _BALANCED_VARIANTS[_sigmoid_group_router]

518for _alias, _adapter in list(MOE_ROUTER_ADAPTERS.items()):

519 MOE_ROUTER_ADAPTERS[f"{_alias}_fixed"] = _BALANCED_VARIANTS[_adapter]

520del _alias, _adapter

521

522

523def resolve_moe_router(router, fix_router: bool = False):

524 """Resolve a router adapter, optionally to its benchmark-only balanced variant.

525

526 ``fix_router`` is the switch every archetype factory exposes; it has to mean the

527 same thing everywhere, so it is resolved here rather than re-implemented per

528 factory (see :func:`make_balanced_router` for what it does and why its numbers are

529 not efficiency results).

530

531 Args:

532 router: An adapter callable, or the name of an entry in

533 :data:`MOE_ROUTER_ADAPTERS`.

534 fix_router: When true, return the balanced variant: the ``<name>_fixed``

535 registry entry for a name, or ``make_balanced_router(router)`` for a

536 callable.

537

538 Returns:

539 A ``(module, hidden) -> (topk_idx, topk_w)`` adapter.

540

541 Raises:

542 KeyError: If ``router`` is a name that is not registered.

543 TypeError: If ``router`` is neither a registry name nor callable.

544 """

545 if isinstance(router, str):

546 key = f"{router}_fixed" if fix_router else router

547 return MOE_ROUTER_ADAPTERS[key]

548 if not callable(router):

549 raise TypeError(

550 "router must be a registry name or a callable, got "

551 f"{type(router).__name__}"

552 )

553 return make_balanced_router(router) if fix_router else router

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / routing.py: 88%

154 statements