Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / route.py: 95%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

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 

6# 

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

8# 

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# ============================================================================ 

15"""Prepare trusted Router output for managed Torch MegaMoe execution.""" 

16 

17from __future__ import annotations 

18 

19from array import array 

20from collections.abc import Callable 

21from dataclasses import dataclass, replace 

22from itertools import accumulate 

23from typing import TYPE_CHECKING, Any 

24 

25import torch 

26import torch.distributed as dist 

27 

28from hyper_multicore._loader import load_torch_npu 

29from hyper_multicore.torch import ops as multicore_ops 

30 

31from .spec import MegaMoeSpec, _align_capacity 

32 

33if TYPE_CHECKING: 

34 from .workspace import MegaMoeWorkspace 

35 

36 

37# Avoid large Host transfers and Python boxing when the gathered count table grows. 

38_HOST_ROUTE_MAX_COUNTS = 2048 

39# Preserve the smaller descriptor budget when counts change on every call. 

40_HOST_DESCRIPTOR_MAX_COUNTS = 1024 

41_ROUTE_CACHE_REPROBE_STEPS = 8 

42_METADATA_ALIGNMENT = 512 

43_PULL_STAGING_MAX_BYTES = 64 * 1024 * 1024 

44 

45 

46@dataclass(frozen=True) 

47class RouteMetadata: 

48 """Exact offsets and counts consumed by the low-level native ABI.""" 

49 

50 dispatch_src_off: torch.Tensor 

51 dispatch_target_off: torch.Tensor 

52 dispatch_size: torch.Tensor 

53 combine_src_off: torch.Tensor 

54 combine_target_off: torch.Tensor 

55 combine_size: torch.Tensor 

56 group_list: torch.Tensor 

57 expert_capacity: int 

58 local_num_tokens: int | None = None 

59 plan_tokens: int | None = None 

60 

61 

62@dataclass(frozen=True) 

63class PreparedTopKRoute: 

64 """Permuted payload and exact metadata for one Router decision.""" 

65 

66 routed_tokens: torch.Tensor 

67 unpermute_mapping: torch.Tensor 

68 tokens_per_expert: torch.Tensor 

69 received_counts: torch.Tensor 

70 metadata: RouteMetadata 

71 maximum_received_slots: int = 0 

72 

73 

74@dataclass(frozen=True) 

75class _HostRouteCache: 

76 """One validated wire table and immutable route descriptors for a workspace.""" 

77 

78 signature: tuple[Any, ...] 

79 raw: list[list[int]] 

80 summary: tuple[int, int, int] 

81 metadata: RouteMetadata | None 

82 stream: Any = None 

83 ready: Any = None 

84 

85 def wait(self, device: Any) -> None: 

86 """Order a different consumer stream after the descriptor upload. 

87 

88 Args: 

89 device: Device whose current stream consumes the cached descriptors. 

90 """ 

91 if self.ready is not None and self.metadata is not None: 

92 current = torch.npu.current_stream(device) 

93 if current != self.stream: 

94 current.wait_event(self.ready) 

95 for tensor in vars(self.metadata).values(): 

96 if isinstance(tensor, torch.Tensor): 

97 tensor.record_stream(current) 

98 

99 

100@dataclass(frozen=True) 

101class _RouteHeaderCache: 

102 """One immutable control header, retained across identical local invocations.""" 

103 

104 key: tuple[Any, ...] 

105 tensor: torch.Tensor | None 

106 stream: Any = None 

107 ready: Any = None 

108 

109 def wait(self, device: Any) -> None: 

110 """Preserve upload ordering and allocation lifetime on another stream. 

111 

112 Args: 

113 device: Device whose current stream consumes the cached header. 

114 """ 

115 if self.ready is not None and self.tensor is not None: 

116 current = torch.npu.current_stream(device) 

117 if current != self.stream: 

118 current.wait_event(self.ready) 

119 self.tensor.record_stream(current) 

120 

121 

122def _validate_topk_inputs( 

123 hidden_states: torch.Tensor, 

124 topk_ids: torch.Tensor, 

125 topk_weights: torch.Tensor, 

126) -> torch.Tensor: 

127 """Validate route tensor metadata without synchronizing tensor contents.""" 

128 if topk_ids.dtype not in (torch.int32, torch.int64): 

129 raise TypeError(f"topk_ids must be int32 or int64, got {topk_ids.dtype}.") 

130 if not topk_weights.is_floating_point(): 

131 raise TypeError( 

132 f"topk_weights must be floating point, got {topk_weights.dtype}." 

133 ) 

134 if ( 

135 topk_ids.device != hidden_states.device 

136 or topk_weights.device != hidden_states.device 

137 ): 

138 raise ValueError( 

139 "hidden_states, topk_ids and topk_weights must be on one device." 

140 ) 

141 return topk_ids.reshape(-1) 

142 

143 

144def _resolve_counts( 

145 flat_ids: torch.Tensor, 

146 supplied_counts: torch.Tensor | None, 

147 spec: MegaMoeSpec, 

148) -> torch.Tensor: 

149 """Use trusted Router counts or compute them only when omitted.""" 

150 if supplied_counts is None: 

151 return ( 

152 torch.bincount(flat_ids.to(torch.int64), minlength=spec.num_experts) 

153 .to(torch.int32) 

154 .contiguous() 

155 ) 

156 if tuple(supplied_counts.shape) != (spec.num_experts,): 

157 raise ValueError( 

158 f"tokens_per_expert must have shape ({spec.num_experts},), " 

159 f"got {tuple(supplied_counts.shape)}." 

160 ) 

161 if supplied_counts.device != flat_ids.device: 

162 raise ValueError("tokens_per_expert must be on the route tensor device.") 

163 if supplied_counts.dtype not in (torch.int32, torch.int64): 

164 raise TypeError( 

165 f"tokens_per_expert must be int32 or int64, got {supplied_counts.dtype}." 

166 ) 

167 if supplied_counts.dtype == torch.int32 and supplied_counts.is_contiguous(): 

168 return supplied_counts 

169 return supplied_counts.to(torch.int32).contiguous() 

170 

171 

172def _start_count_gather( 

173 counts: torch.Tensor, 

174 spec: MegaMoeSpec, 

175) -> tuple[torch.Tensor, Any | None]: 

176 """Start the rank-major INT32 count gather used by route construction.""" 

177 if spec.ep_size == 1: 

178 return counts.reshape(1, spec.num_experts), None 

179 gathered_counts = torch.empty( 

180 spec.ep_size * spec.num_experts, 

181 dtype=torch.int32, 

182 device=counts.device, 

183 ) 

184 work = dist.all_gather_into_tensor( 

185 gathered_counts, 

186 counts, 

187 group=spec.ep_group, 

188 async_op=True, 

189 ) 

190 return gathered_counts.reshape(spec.ep_size, spec.num_experts), work 

191 

192 

193def _finish_count_gather( 

194 counts_by_source: torch.Tensor, 

195 work: Any | None, 

196 spec: MegaMoeSpec, 

197) -> torch.Tensor: 

198 """Wait for gathered counts before reading route metadata.""" 

199 if work is not None: 

200 work.wait() 

201 local_start = spec.rank_id * spec.local_experts 

202 local_end = local_start + spec.local_experts 

203 return counts_by_source[:, local_start:local_end].contiguous() 

204 

205 

206def _expert_capacity( 

207 counts_by_source: torch.Tensor, 

208 spec: MegaMoeSpec, 

209) -> tuple[int, int]: 

210 """Validate the global bound and size rank-local computation tensors.""" 

211 destination_loads = counts_by_source.reshape( 

212 spec.ep_size, 

213 spec.ep_size, 

214 spec.local_experts, 

215 ).sum(dim=(0, 2), dtype=torch.int32) 

216 # One host transfer serves both the coordinated overflow check and local 

217 # allocation; the existing gathered counts require no extra collective. 

218 loads = destination_loads.tolist() 

219 # Keep a non-null ABI argument on ranks whose experts receive no tokens. 

220 # Source outputs and symmetric communication buffers retain their own sizes. 

221 return max(1, loads[spec.rank_id]), max(loads) 

222 

223 

224def _permute_topk_input( 

225 hidden_states: torch.Tensor, 

226 topk_ids: torch.Tensor, 

227) -> tuple[torch.Tensor, torch.Tensor]: 

228 """Expand hidden rows with the native NPU MoE permutation operator.""" 

229 if hidden_states.shape[0] == 0: 

230 return hidden_states, torch.empty(0, dtype=torch.int32, device=hidden_states.device) 

231 routed_tokens, unpermute_mapping = load_torch_npu().npu_moe_token_permute( 

232 hidden_states, 

233 topk_ids, 

234 ) 

235 return ( 

236 routed_tokens, 

237 unpermute_mapping.reshape(-1).to(torch.int32).contiguous(), 

238 ) 

239 

240 

241def _permute_topk_input_out( 

242 hidden_states: torch.Tensor, 

243 topk_ids: torch.Tensor, 

244 output: torch.Tensor, 

245) -> tuple[torch.Tensor, torch.Tensor]: 

246 """Keep the inverse mapping owned while writing rows into leased SHMEM.""" 

247 mapping = torch.empty(topk_ids.numel(), dtype=torch.int32, device=topk_ids.device) 

248 multicore_ops.moe_token_permute_out(hidden_states.contiguous(), topk_ids.contiguous(), output, mapping) 

249 return output, mapping 

250 

251 

252def _compute_route_metadata( 

253 counts: torch.Tensor, 

254 counts_by_source: torch.Tensor, 

255 received_counts: torch.Tensor, 

256 spec: MegaMoeSpec, 

257 expert_capacity: int, 

258) -> RouteMetadata: 

259 """Build exact element offsets without padding or route truncation.""" 

260 counts_i32 = counts.to(dtype=torch.int32).contiguous() 

261 source_prefix = counts_by_source.cumsum(dim=1, dtype=torch.int64) - counts_by_source 

262 destination_layout = ( 

263 counts_by_source.reshape( 

264 spec.ep_size, 

265 spec.ep_size, 

266 spec.local_experts, 

267 ) 

268 .permute(1, 2, 0) 

269 .contiguous() 

270 ) 

271 destination_flat = destination_layout.reshape(spec.ep_size, -1) 

272 destination_prefix = ( 

273 destination_flat.cumsum(dim=1, dtype=torch.int64) - destination_flat 

274 ).reshape(spec.ep_size, spec.local_experts, spec.ep_size) 

275 local_destination_prefix = destination_prefix[spec.rank_id] 

276 local_destination_counts = destination_layout[spec.rank_id] 

277 local_start = spec.rank_id * spec.local_experts 

278 local_end = local_start + spec.local_experts 

279 return RouteMetadata( 

280 dispatch_src_off=(source_prefix[:, local_start:local_end].reshape(-1).to(torch.int64) 

281 if spec.dispatch_mode == "pull" else source_prefix[spec.rank_id].to(torch.int64)), 

282 dispatch_target_off=( 

283 local_destination_prefix.T.reshape(-1).to(torch.int64) if spec.dispatch_mode == "pull" 

284 else destination_prefix[:, :, spec.rank_id].reshape(-1).to(torch.int64) 

285 ), 

286 dispatch_size=(received_counts.reshape(-1).to(torch.int32) if spec.dispatch_mode == "pull" else counts_i32), 

287 combine_src_off=local_destination_prefix.T.reshape(-1).to(torch.int64), 

288 combine_target_off=( 

289 source_prefix[:, local_start:local_end].reshape(-1).to(torch.int64) 

290 ), 

291 combine_size=received_counts.reshape(-1).to(torch.int32), 

292 group_list=( 

293 local_destination_prefix[:, -1] + local_destination_counts[:, -1] 

294 ).to(torch.int64), 

295 expert_capacity=expert_capacity, 

296 ) 

297 

298 

299def _upload_route_tensor(tensor: torch.Tensor, device: Any) -> torch.Tensor: 

300 """Keep Host metadata alive through the pinned allocator during asynchronous upload.""" 

301 if torch.device(device).type == "npu": 

302 return tensor.pin_memory().to(device, non_blocking=True) 

303 return tensor.to(device) 

304 

305 

306def _host_route_metadata( 

307 raw: list[list[int]], 

308 spec: MegaMoeSpec, 

309 capacity: int, 

310 device: Any, 

311) -> RouteMetadata: 

312 """Pack the native INT64 offsets and INT32 route sizes in one retained allocation.""" 

313 counts = [row[3:] for row in raw] 

314 local, rank = spec.local_experts, spec.rank_id 

315 start = rank * local 

316 prefix, destinations, destination_prefix = _host_route_prefixes(counts, spec) 

317 source_offsets = [ 

318 prefix[source][expert] for source in range(spec.ep_size) for expert in range(start, start + local) 

319 ] 

320 target_offsets = [ 

321 destination_prefix[rank][expert * spec.ep_size + source] 

322 for source in range(spec.ep_size) 

323 for expert in range(local) 

324 ] 

325 received = [counts[source][expert] for source in range(spec.ep_size) for expert in range(start, start + local)] 

326 if spec.dispatch_mode == "pull": 

327 dispatch_source, dispatch_target, dispatch_sizes = source_offsets, target_offsets, received 

328 else: 

329 dispatch_source = prefix[rank] 

330 dispatch_target = [ 

331 destination_prefix[target][expert * spec.ep_size + rank] 

332 for target in range(spec.ep_size) 

333 for expert in range(local) 

334 ] 

335 dispatch_sizes = counts[rank] 

336 groups = [ 

337 destination_prefix[rank][(expert + 1) * spec.ep_size - 1] + destinations[rank][(expert + 1) * spec.ep_size - 1] 

338 for expert in range(local) 

339 ] 

340 wide = (dispatch_source, dispatch_target, target_offsets, source_offsets, groups) 

341 return _pack_host_route(wide, (dispatch_sizes, received), capacity, device) 

342 

343 

344def _host_route_prefixes( 

345 counts: list[list[int]], spec: MegaMoeSpec, 

346) -> tuple[list[list[int]], list[list[int]], list[list[int]]]: 

347 """Scan each source and each destination independently in native layout order.""" 

348 local = spec.local_experts 

349 prefix = [list(accumulate(row, initial=0))[:-1] for row in counts] 

350 destinations = [ 

351 [ 

352 counts[source][expert] 

353 for expert in range(target * local, (target + 1) * local) 

354 for source in range(spec.ep_size) 

355 ] 

356 for target in range(spec.ep_size) 

357 ] 

358 destination_prefix = [list(accumulate(row, initial=0))[:-1] for row in destinations] 

359 return prefix, destinations, destination_prefix 

360 

361 

362def _pack_host_route( 

363 wide: tuple[list[int], ...], 

364 size_fields: tuple[list[int], ...], 

365 capacity: int, 

366 device: Any, 

367) -> RouteMetadata: 

368 """Transfer aligned descriptor fields together and retain storage through their views.""" 

369 fields = [*wide, *size_fields] 

370 dtypes = [torch.int64] * len(wide) + [torch.int32] * 2 

371 storage = bytearray() 

372 slices = [] 

373 for values, dtype in zip(fields, dtypes): 

374 encoded = array("q" if dtype == torch.int64 else "i", values).tobytes() 

375 start = len(storage) 

376 storage.extend(encoded) 

377 slices.append((start, len(storage), dtype)) 

378 storage.extend(bytes((-len(storage)) % _METADATA_ALIGNMENT)) 

379 buffer = _upload_route_tensor(torch.frombuffer(storage, dtype=torch.uint8), device) 

380 tensors = [buffer[start:end].view(dtype) for start, end, dtype in slices] 

381 offsets, sizes = tensors[:5], tensors[5:] 

382 return RouteMetadata( 

383 dispatch_src_off=offsets[0], 

384 dispatch_target_off=offsets[1], 

385 dispatch_size=sizes[0], 

386 combine_src_off=offsets[2], 

387 combine_target_off=offsets[3], 

388 combine_size=sizes[1], 

389 group_list=offsets[4], 

390 expert_capacity=capacity, 

391 ) 

392 

393 

394def prepare_topk_route( 

395 hidden_states: torch.Tensor, 

396 topk_ids: torch.Tensor, 

397 topk_weights: torch.Tensor, 

398 spec: MegaMoeSpec, 

399 tokens_per_expert: torch.Tensor | None, 

400 workspace: MegaMoeWorkspace | None = None, 

401 *, 

402 autograd_mask: int = 0, 

403 prepare_workspace: Callable[[int, int], None] | None = None, 

404) -> PreparedTopKRoute: 

405 """Overlap Top-K permutation with count exchange and build route metadata. 

406 

407 Args: 

408 hidden_states: Flattened token states to route. 

409 topk_ids: Global experts selected for every token. 

410 topk_weights: Router weights paired with ``topk_ids``. 

411 spec: Bound shape and expert-parallel specification. 

412 tokens_per_expert: Optional trusted Router histogram. 

413 workspace: Optional workspace ordered before the count exchange. Pull 

414 requires a lease acquired by the caller or the automatic-capacity callback. 

415 autograd_mask: Group-consistent gradient participation flags in dynamic mode. 

416 prepare_workspace: Automatic-capacity callback, run after count agreement and before pull permutation. 

417 

418 Returns: 

419 Permuted tokens and exact native route metadata. 

420 """ 

421 flat_ids = _validate_topk_inputs(hidden_states, topk_ids, topk_weights) 

422 if spec.dynamic_tokens: 

423 return _prepare_dynamic_route(hidden_states, topk_ids, flat_ids, spec, tokens_per_expert, 

424 workspace, autograd_mask, prepare_workspace) 

425 counts = _resolve_counts(flat_ids, tokens_per_expert, spec) 

426 permutation_output = None 

427 if workspace is not None: 

428 if spec.dispatch_mode == "pull": 

429 if not workspace.in_use or workspace.source_buffer is None: 

430 raise RuntimeError("Pull permutation requires an initialized, claimed workspace.") 

431 permutation_output = workspace.source_buffer 

432 else: 

433 workspace.wait_for_reuse() 

434 counts_by_source, count_work = _start_count_gather(counts, spec) 

435 if permutation_output is None: 

436 routed_tokens, unpermute_mapping = _permute_topk_input(hidden_states, topk_ids) 

437 else: 

438 routed_tokens, unpermute_mapping = _permute_topk_input_out(hidden_states, topk_ids, permutation_output) 

439 received_counts = _finish_count_gather(counts_by_source, count_work, spec) 

440 expert_capacity, maximum_received_slots = _expert_capacity(counts_by_source, spec) 

441 return PreparedTopKRoute( 

442 routed_tokens=routed_tokens, 

443 unpermute_mapping=unpermute_mapping, 

444 tokens_per_expert=counts, 

445 received_counts=received_counts, 

446 maximum_received_slots=maximum_received_slots, 

447 metadata=_compute_route_metadata( 

448 counts, 

449 counts_by_source, 

450 received_counts, 

451 spec, 

452 expert_capacity, 

453 ), 

454 ) 

455 

456 

457def _valid_route_counts(row: list[int], top_k: int) -> bool: 

458 """Validate one source histogram against its actual distinct Top-K rows.""" 

459 return row[3] == row[1] * top_k and row[4] >= 0 and row[5] <= row[1] 

460 

461 

462def _validate_dynamic_sizes(rows: list[list[int]], spec: MegaMoeSpec) -> None: 

463 """Keep native size and legacy bound rejection ahead of route publication.""" 

464 flexible = spec.auto_tokens 

465 if flexible and any(row[0] or row[1] < 0 for row in rows): 

466 raise RuntimeError("MegaMoe dynamic route exceeds native INT32 size limits") 

467 if not flexible and any( 

468 row[0] or not 0 <= row[1] <= spec.max_local_num_tokens for row in rows): 

469 raise RuntimeError(f"MegaMoe token capacity overflow: limit={spec.max_local_num_tokens}, " 

470 f"actual={[row[1] for row in rows]}") 

471 

472 

473def _validate_dynamic_rows(rows: list[list[int]], spec: MegaMoeSpec) -> None: 

474 """Validate every rank before any automatic allocation or SHMEM write.""" 

475 _validate_dynamic_sizes(rows, spec) 

476 if any(row[2] != rows[0][2] for row in rows): 

477 raise RuntimeError("MegaMoe autograd participation differs across EP ranks") 

478 if any(not _valid_route_counts(row, spec.top_k) for row in rows): 

479 raise RuntimeError("MegaMoe counts must describe the actual distinct Top-K routes") 

480 

481 

482def _dynamic_summary(payload: torch.Tensor, spec: MegaMoeSpec) -> tuple[int, int, int]: 

483 """Use one Host summary for shape, participation, receive limits, and schedule.""" 

484 counts = payload[:, 3:] 

485 loads = counts.reshape(spec.ep_size, spec.ep_size, spec.local_experts).sum(dim=(0, 2)) 

486 summary = torch.cat((payload[:, :3], counts.sum(1, keepdim=True), 

487 counts.amin(1, keepdim=True), counts.amax(1, keepdim=True), loads[:, None]), dim=1) 

488 rows = summary.tolist() 

489 return _summarize_dynamic_rows(rows, spec) 

490 

491 

492def _summarize_dynamic_rows(rows: list[list[int]], spec: MegaMoeSpec) -> tuple[int, int, int]: 

493 """Keep overflow and participation checks common to both metadata policies.""" 

494 _validate_dynamic_rows(rows, spec) 

495 maximum = max(row[6] for row in rows) 

496 maximum_tokens = max(row[1] for row in rows) if spec.auto_tokens else spec.local_num_tokens 

497 maximum_capacity = _align_capacity(spec.ep_size * maximum_tokens * spec.top_k) 

498 if maximum > maximum_capacity: 

499 raise RuntimeError( 

500 f"MegaMoe receive capacity overflow: configured_capacity={maximum_capacity}, actual_maximum={maximum}" 

501 ) 

502 plan_tokens = max(128, _align_capacity(max(row[1] for row in rows))) 

503 if plan_tokens * max(spec.hidden_size, spec.top_k) > 2**31 - 1 or maximum > 2**31 - 1: 

504 raise RuntimeError("MegaMoe dynamic route exceeds native INT32 size limits") 

505 return max(1, rows[spec.rank_id][6]), plan_tokens, maximum 

506 

507 

508def _read_dynamic_route( 

509 payload: torch.Tensor, 

510 spec: MegaMoeSpec, 

511) -> tuple[tuple[int, int, int], list[list[int]] | None]: 

512 """Reuse the required readback for small descriptors; reduce large tables on device.""" 

513 if spec.ep_size * spec.num_experts > _HOST_ROUTE_MAX_COUNTS: 

514 return _dynamic_summary(payload, spec), None 

515 raw = payload.tolist() 

516 return _summarize_host_route(raw, spec), raw 

517 

518 

519def _summarize_host_route(raw: list[list[int]], spec: MegaMoeSpec) -> tuple[int, int, int]: 

520 """Validate a bounded Host table before allocating or publishing source rows.""" 

521 counts = [row[3:] for row in raw] 

522 local = spec.local_experts 

523 loads = [ 

524 sum(sum(values[start : start + local]) for values in counts) for start in range(0, spec.num_experts, local) 

525 ] 

526 rows = [row[:3] + [sum(values), min(values), max(values), load] for row, values, load in zip(raw, counts, loads)] 

527 return _summarize_dynamic_rows(rows, spec) 

528 

529 

530def _route_cache_signature(spec: MegaMoeSpec, device: Any) -> tuple[Any, ...]: 

531 """Keep descriptor reuse specific to the route topology and size contract.""" 

532 return (spec.ep_size, id(spec.ep_group), spec.rank_id, spec.num_experts, spec.top_k, spec.hidden_size, 

533 spec.dispatch_mode, spec.auto_tokens, spec.max_local_num_tokens, torch.device(device)) 

534 

535 

536def _read_cached_route( 

537 payload: torch.Tensor, spec: MegaMoeSpec, workspace: MegaMoeWorkspace | None, device: Any, 

538) -> tuple[tuple[int, int, int], list[list[int]] | None, RouteMetadata | None]: 

539 """Reuse only an identical, previously validated all-rank wire table.""" 

540 if workspace is None or not spec.auto_tokens or spec.ep_size * spec.num_experts > _HOST_ROUTE_MAX_COUNTS: 

541 summary, raw = _read_dynamic_route(payload, spec) 

542 return summary, raw, None 

543 cached = getattr(workspace, "route_cache", None) 

544 same_topology = isinstance(cached, _HostRouteCache) and cached.signature == _route_cache_signature(spec, device) 

545 remaining = getattr(workspace, "route_cache_probe_after", 0) 

546 if same_topology and remaining > 0: 

547 workspace.route_cache_probe_after = remaining - 1 

548 return _dynamic_summary(payload, spec), None, None 

549 workspace.route_cache_probe_after = 0 

550 raw = payload.tolist() 

551 if same_topology and cached.raw == raw: 

552 cached.wait(device) 

553 return cached.summary, raw, cached.metadata 

554 summary = _summarize_host_route(raw, spec) 

555 if same_topology and spec.ep_size * spec.num_experts > _HOST_DESCRIPTOR_MAX_COUNTS: 

556 workspace.route_cache_probe_after = _ROUTE_CACHE_REPROBE_STEPS 

557 return summary, raw, None 

558 

559 

560def _cache_host_route( 

561 workspace: MegaMoeWorkspace | None, spec: MegaMoeSpec, raw: list[list[int]], 

562 summary: tuple[int, int, int], metadata: RouteMetadata, device: Any, 

563) -> None: 

564 """Retain descriptors only after confirming a repeated validated wire table.""" 

565 if workspace is None or not spec.auto_tokens: 

566 return 

567 signature = _route_cache_signature(spec, device) 

568 previous = getattr(workspace, "route_cache", None) 

569 repeated = isinstance(previous, _HostRouteCache) and previous.signature == signature and previous.raw == raw 

570 retained = metadata if repeated else None 

571 stream, ready = None, None 

572 if repeated and torch.device(device).type == "npu": 

573 stream = torch.npu.current_stream(device) 

574 ready = torch.npu.Event() 

575 ready.record(stream) 

576 workspace.route_cache = _HostRouteCache(signature, raw, summary, retained, stream, ready) 

577 

578 

579def _dynamic_route_metadata( 

580 counts: torch.Tensor, gathered: torch.Tensor, host_rows: list[list[int]] | None, 

581 spec: MegaMoeSpec, capacity: int, 

582) -> RouteMetadata: 

583 """Keep wider cache misses on device instead of rebuilding large Host descriptors.""" 

584 if host_rows is not None and spec.ep_size * spec.num_experts <= _HOST_DESCRIPTOR_MAX_COUNTS: 

585 return _host_route_metadata(host_rows, spec, capacity, counts.device) 

586 counts_by_source = gathered[:, 3:].contiguous() 

587 start = spec.rank_id * spec.local_experts 

588 received_counts = counts_by_source[:, start:start + spec.local_experts].contiguous() 

589 # Retained dispatch sizes must not alias a caller-owned Router count buffer. 

590 local_counts = counts_by_source[spec.rank_id] if spec.auto_tokens else counts 

591 return _compute_route_metadata(local_counts, counts_by_source, received_counts, spec, capacity) 

592 

593 

594def _permute_dynamic_pull( 

595 hidden_states: torch.Tensor, topk_ids: torch.Tensor, spec: MegaMoeSpec, workspace: MegaMoeWorkspace | None, 

596) -> tuple[torch.Tensor, torch.Tensor]: 

597 """Take the source view only after automatic growth and lease acquisition.""" 

598 if workspace is None or not workspace.in_use or workspace.source_buffer is None: 

599 raise RuntimeError("Pull permutation requires an initialized, claimed workspace.") 

600 tokens = hidden_states.shape[0] 

601 output = workspace.source_buffer[:tokens * spec.top_k] 

602 return (_permute_topk_input_out(hidden_states, topk_ids, output) if tokens 

603 else (output, topk_ids.reshape(-1))) 

604 

605 

606def _can_stage_pull(hidden_states: torch.Tensor, spec: MegaMoeSpec) -> bool: 

607 """Bound ordinary-memory permutation staging while capacity agreement is pending.""" 

608 return (spec.auto_tokens and spec.dispatch_mode == "pull" and hidden_states.device.type == "npu" 

609 and hidden_states.numel() * spec.top_k * hidden_states.element_size() <= _PULL_STAGING_MAX_BYTES) 

610 

611 

612def _stage_dynamic_permutation( 

613 hidden_states: torch.Tensor, topk_ids: torch.Tensor, spec: MegaMoeSpec, invalid: bool, 

614 workspace: MegaMoeWorkspace | None = None, 

615) -> tuple[torch.Tensor, torch.Tensor] | None: 

616 """Stage only valid automatic routes whose ordinary storage survives heap growth.""" 

617 if invalid or not spec.auto_tokens: 

618 return None 

619 if (spec.dispatch_mode == "push" 

620 or _can_stage_pull(hidden_states, spec) and not getattr(workspace, "route_cache_probe_after", 0)): 

621 return _permute_topk_input(hidden_states, topk_ids) 

622 return None 

623 

624 

625def _publish_staged_pull( 

626 permuted: tuple[torch.Tensor, torch.Tensor], hidden_states: torch.Tensor, 

627 spec: MegaMoeSpec, workspace: MegaMoeWorkspace | None, 

628) -> tuple[torch.Tensor, torch.Tensor]: 

629 """Publish staged rows only after every rank validated and committed its workspace.""" 

630 if workspace is None or not workspace.in_use or workspace.source_buffer is None: 

631 raise RuntimeError("Pull publication requires an initialized, claimed workspace.") 

632 output = workspace.source_buffer[:hidden_states.shape[0] * spec.top_k] 

633 output.copy_(permuted[0]) 

634 return output, permuted[1] 

635 

636 

637def _dynamic_header( 

638 values: tuple[int, int, int], dtype: torch.dtype, device: Any, workspace: MegaMoeWorkspace | None, 

639) -> torch.Tensor: 

640 """Avoid repeated Host allocation and upload of an unchanged control header.""" 

641 key = (values, dtype, torch.device(device)) 

642 cached = getattr(workspace, "route_header_cache", None) 

643 repeated = isinstance(cached, _RouteHeaderCache) and cached.key == key 

644 if repeated and cached.tensor is not None: 

645 cached.wait(device) 

646 return cached.tensor 

647 header = _upload_route_tensor(torch.tensor(values, dtype=dtype), device) 

648 if workspace is not None: 

649 stream, ready = None, None 

650 if repeated and torch.device(device).type == "npu": 

651 stream = torch.npu.current_stream(device) 

652 ready = torch.npu.Event() 

653 ready.record(stream) 

654 workspace.route_header_cache = _RouteHeaderCache(key, header if repeated else None, stream, ready) 

655 return header 

656 

657 

658def _dynamic_route_payload( 

659 hidden_states: torch.Tensor, flat_ids: torch.Tensor, spec: MegaMoeSpec, 

660 supplied_counts: torch.Tensor | None, autograd_mask: int, workspace: MegaMoeWorkspace | None = None, 

661) -> tuple[torch.Tensor, torch.Tensor, bool]: 

662 """Pack one group-consistent wire format, including coordinated overflow flags.""" 

663 tokens = hidden_states.shape[0] 

664 flexible = spec.auto_tokens 

665 invalid = (tokens * max(spec.hidden_size, spec.top_k) > 2**31 - 1 if flexible 

666 else tokens > spec.max_local_num_tokens) 

667 counts = (torch.zeros(spec.num_experts, dtype=torch.int32, device=hidden_states.device) if invalid 

668 else _resolve_counts(flat_ids, supplied_counts, spec)) 

669 payload_dtype = torch.int32 if flexible else torch.int64 

670 # An invalid source must still join the same collective before all ranks reject it. 

671 header_tokens = min(tokens, 2**31 - 1) if flexible else tokens 

672 header = _dynamic_header((int(invalid), header_tokens, autograd_mask), payload_dtype, 

673 hidden_states.device, workspace if flexible else None) 

674 payload = torch.cat((header, counts.to(payload_dtype)))