Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / workspace.py: 94%
237 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
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"""Owned directional workspace for managed Torch MegaMoe execution."""
17from __future__ import annotations
19import math
20import threading
21import weakref
22from collections.abc import Mapping
23from contextlib import nullcontext
24from dataclasses import dataclass, field
25from typing import Any
27import torch
29from hyper_multicore import _collective, shmem
30from hyper_multicore.scheduler.config import (
31 MIN_EVENT_CAPACITY,
32 event_workspace_bytes,
33 mega_moe_event_capacity,
34)
35from hyper_multicore.shmem import _lifecycle
37from .spec import MegaMoeSpec, _resolve_receive_capacity
39# The composed Cube-only GMM and SwiGLU-grad kernels do not use these legacy
40# tensor arguments. Keep non-empty ABI placeholders; native tiling reserves
41# the separate CANN library workspace needed by the enclosing operator.
42_GMM_WORKSPACE_BYTES = 512
43_SWIGLU_GRAD_WORKSPACE_BYTES = 512
44_WORKSPACE_ALIGNMENT = 512
47def _workspace_allocation_specs(
48 *,
49 dispatch_mode: str,
50 routed_slots: int,
51 receive_capacity: int,
52 hidden_size: int,
53 ep_size: int,
54 num_experts: int,
55 dtype: torch.dtype,
56) -> tuple[tuple[str, _lifecycle._AllocationSpec], ...]: # pylint: disable=protected-access
57 """Return the single ordered physical Allocation description for a workspace."""
58 data_field = "source_buffer" if dispatch_mode == "pull" else "expert_buffer"
59 data_rows = routed_slots if dispatch_mode == "pull" else receive_capacity
60 event_bytes = event_workspace_bytes(ep_size, num_experts)
61 allocation_spec = _lifecycle._AllocationSpec # pylint: disable=protected-access
62 return (
63 (data_field, allocation_spec((data_rows, hidden_size), dtype, _WORKSPACE_ALIGNMENT)),
64 ("routed_buffer", allocation_spec((routed_slots, hidden_size), dtype, _WORKSPACE_ALIGNMENT)),
65 ("forward_event_counters", allocation_spec((event_bytes,), torch.uint8, _WORKSPACE_ALIGNMENT)),
66 ("backward_event_counters", allocation_spec((event_bytes,), torch.uint8, _WORKSPACE_ALIGNMENT)),
67 )
70def _spec_workspace_bytes(
71 specification: Mapping[str, Any], dtype: torch.dtype, *, receive_capacity: int | None = None,
72) -> int:
73 """Return planned symmetric bytes for one execution-resource group."""
74 local_tokens = specification["local_num_tokens"]
75 routed_slots = local_tokens * specification["top_k"]
76 capacity = _resolve_receive_capacity(
77 specification["initial_capacity_factor"],
78 routed_slots,
79 specification["ep_size"],
80 )
81 if receive_capacity is not None:
82 capacity = receive_capacity
83 if specification.get("dispatch_mode", "push") == "pull":
84 capacity = routed_slots
85 allocations = _workspace_allocation_specs(
86 dispatch_mode=specification.get("dispatch_mode", "push"),
87 routed_slots=routed_slots,
88 receive_capacity=capacity,
89 hidden_size=specification["hidden_size"],
90 ep_size=specification["ep_size"],
91 num_experts=specification["num_experts"],
92 dtype=dtype,
93 )
94 return sum(
95 spec.dtype.itemsize * math.prod(spec.shape)
96 + spec.alignment_bytes - 1
97 for _, spec in allocations
98 )
101@dataclass
102class MegaMoeWorkspace:
103 """Buffers owned by one standalone or explicitly shared resource group."""
105 shared: bool
106 dtype: Any | None = None
107 device: Any | None = None
108 expert_capacity: int = 0
109 routed_slots: int = 0
110 dispatch_mode: str = "push"
111 source_buffer: Any | None = None
112 expert_buffer: Any | None = None
113 routed_buffer: Any | None = None
114 forward_event_counters: Any | None = None
115 backward_event_counters: Any | None = None
116 gmm_workspace: Any | None = None
117 swiglu_grad_workspace: Any | None = None
118 completion_event: Any | None = None
119 in_use: bool = False
120 closing: bool = False
121 used: bool = False
122 event_counter_bytes: int = MIN_EVENT_CAPACITY * 4
123 forward_ready_initialized: bool = False
124 backward_ready_initialized: bool = False
125 capacity_floor: int = 0
126 heap_manager: Any | None = field(default=None, repr=False)
127 route_cache: Any | None = field(default=None, repr=False)
128 route_header_cache: Any | None = field(default=None, repr=False)
129 route_cache_probe_after: int = field(default=0, repr=False)
130 lock: Any = field(default_factory=threading.Lock, repr=False)
131 graphs: Any = field(default_factory=weakref.WeakSet, repr=False)
133 def describe_restart_allocations(
134 self,
135 spec: MegaMoeSpec,
136 dtype: torch.dtype,
137 receive_capacity: int,
138 ) -> tuple[_lifecycle._AllocationReplacement, ...]: # pylint: disable=protected-access
139 """Project materialized symmetric fields into ordered restart replacements."""
140 if self.completion_event is None:
141 return ()
142 targets = _workspace_allocation_specs(
143 dispatch_mode=spec.dispatch_mode,
144 routed_slots=spec.routed_slots,
145 receive_capacity=receive_capacity,
146 hidden_size=spec.hidden_size,
147 ep_size=spec.ep_size,
148 num_experts=spec.num_experts,
149 dtype=dtype,
150 )
151 replacement = _lifecycle._AllocationReplacement # pylint: disable=protected-access
152 return tuple(replacement(getattr(self, field_name), target) for field_name, target in targets)
154 def validate_restart_allocations(
155 self,
156 spec: MegaMoeSpec,
157 dtype: torch.dtype,
158 current_capacity: int,
159 replacements: tuple[_lifecycle._AllocationReplacement, ...], # pylint: disable=protected-access
160 ) -> bool:
161 """Validate the read-only workspace state captured for restart preflight."""
162 if self.in_use:
163 raise RuntimeError("active workspace lease prevents heap growth")
164 if self.closing:
165 raise RuntimeError("closing workspace prevents heap growth")
166 materialized = self.completion_event is not None
167 if not materialized:
168 if replacements:
169 raise RuntimeError("lazy workspace must not describe restart Allocations")
170 return False
171 expected = _workspace_allocation_specs(
172 dispatch_mode=spec.dispatch_mode,
173 routed_slots=spec.routed_slots,
174 receive_capacity=current_capacity,
175 hidden_size=spec.hidden_size,
176 ep_size=spec.ep_size,
177 num_experts=spec.num_experts,
178 dtype=dtype,
179 )
180 if len(replacements) != len(expected):
181 raise RuntimeError("materialized workspace has an incomplete restart Allocation mapping")
182 for (field_name, current_spec), replacement in zip(expected, replacements):
183 tensor = getattr(self, field_name)
184 if replacement.current is not tensor or tensor is None:
185 raise RuntimeError("workspace restart Allocation mapping changed during preflight")
186 if tuple(tensor.shape) != current_spec.shape or tensor.dtype != current_spec.dtype:
187 raise RuntimeError("workspace symmetric Allocation does not match its committed layout")
188 return True
190 def bind_restarted_allocations(
191 self,
192 spec: MegaMoeSpec,
193 dtype: torch.dtype,
194 device: Any,
195 receive_capacity: int,
196 tensors: tuple[torch.Tensor, ...],
197 ) -> None:
198 """Bind validated replacement Tensors without allocating or running collectives."""
199 targets = _workspace_allocation_specs(
200 dispatch_mode=spec.dispatch_mode,
201 routed_slots=spec.routed_slots,
202 receive_capacity=receive_capacity,
203 hidden_size=spec.hidden_size,
204 ep_size=spec.ep_size,
205 num_experts=spec.num_experts,
206 dtype=dtype,
207 )
208 self.source_buffer = None
209 self.expert_buffer = None
210 for tensor, (field_name, _) in zip(tensors, targets):
211 setattr(self, field_name, tensor)
212 self.dispatch_mode = spec.dispatch_mode
213 self.dtype = dtype
214 self.device = device
215 self.expert_capacity = receive_capacity
216 self.routed_slots = spec.routed_slots
217 self.event_counter_bytes = mega_moe_event_capacity(spec.num_experts, spec.ep_size) * 4
218 self.forward_ready_initialized = False
219 self.backward_ready_initialized = False
221 def validate_restarted_allocations(
222 self,
223 spec: MegaMoeSpec,
224 dtype: torch.dtype,
225 receive_capacity: int,
226 tensors: tuple[torch.Tensor, ...],
227 ) -> None:
228 """Validate a complete replacement slice without modifying workspace state."""
229 targets = _workspace_allocation_specs(
230 dispatch_mode=spec.dispatch_mode,
231 routed_slots=spec.routed_slots,
232 receive_capacity=receive_capacity,
233 hidden_size=spec.hidden_size,
234 ep_size=spec.ep_size,
235 num_experts=spec.num_experts,
236 dtype=dtype,
237 )
238 if len(tensors) != len(targets):
239 raise RuntimeError("SHMEM restart returned an unexpected Allocation count")
240 for tensor, (_, target) in zip(tensors, targets):
241 if tuple(tensor.shape) != target.shape or tensor.dtype != target.dtype:
242 raise RuntimeError("SHMEM restart returned an unexpected Allocation specification")
244 def track_graph(self, context: Any) -> None:
245 """Retain a weak lease while this context can still run backward.
247 Args:
248 context: Autograd context that may still access workspace buffers.
249 """
250 self.graphs.add(context)
252 def can_close(self) -> bool:
253 """Check active calls and weak graph leases without unpacking saved tensors."""
254 return not self.in_use and not self.graphs
256 def release_graph(self, context: Any) -> None:
257 """Retire this context's lease after successful non-retained backward.
259 Args:
260 context: Autograd context whose final backward completed.
261 """
262 self.graphs.discard(context)
264 def ensure(self, spec: MegaMoeSpec, dtype: Any, device: Any) -> None:
265 """Allocate the fixed configured route capacity once.
267 Args:
268 spec: First-forward-bound shape and capacity specification.
269 dtype: Element dtype for symmetric data buffers.
270 device: Device that owns local and symmetric buffers.
271 """
272 with self._access():
273 self._ensure_buffers(spec, dtype, device)
275 def _access(self) -> Any:
276 """Exclude submissions during a root Runtime restart."""
277 return self.heap_manager.access() if self.heap_manager is not None else nullcontext()
279 def _ensure_buffers(self, spec: MegaMoeSpec, dtype: Any, device: Any) -> None:
280 """Allocate this workspace at its current monotonic capacity."""
281 if self.closing:
282 raise RuntimeError("cannot allocate a closing MegaMoe workspace.")
283 requested_capacity = max(spec.receive_capacity, self.capacity_floor)
284 if self.expert_buffer is not None or self.source_buffer is not None:
285 compatible = (
286 self.dispatch_mode == spec.dispatch_mode
287 and self.dtype == dtype
288 and self.device == device
289 and self.routed_slots == spec.routed_slots
290 and self.expert_capacity == requested_capacity
291 )
292 if not compatible:
293 raise ValueError("MegaMoe workspace cannot change device, dtype, routed shape, or capacity.")
294 return
296 allocations = _workspace_allocation_specs(
297 dispatch_mode=spec.dispatch_mode,
298 routed_slots=spec.routed_slots,
299 receive_capacity=requested_capacity,
300 hidden_size=spec.hidden_size,
301 ep_size=spec.ep_size,
302 num_experts=spec.num_experts,
303 dtype=dtype,
304 )
305 self.dispatch_mode = spec.dispatch_mode
306 self.dtype = dtype
307 self.device = device
308 self.expert_capacity = requested_capacity
309 self.routed_slots = spec.routed_slots
310 self.event_counter_bytes = mega_moe_event_capacity(spec.num_experts, spec.ep_size) * 4
311 try:
312 for field_name, allocation in allocations:
313 setattr(self, field_name, shmem.empty(
314 allocation.shape,
315 dtype=allocation.dtype,
316 alignment=allocation.alignment_bytes,
317 ))
318 self.gmm_workspace = torch.empty(
319 (_GMM_WORKSPACE_BYTES,),
320 dtype=torch.uint8,
321 device=device,
322 )
323 self.swiglu_grad_workspace = torch.empty(
324 (_SWIGLU_GRAD_WORKSPACE_BYTES,),
325 dtype=torch.uint8,
326 device=device,
327 )
328 self.completion_event = torch.npu.Event()
329 except Exception:
330 self._free_symmetric_tensors()
331 self._free_local_tensors()
332 raise
334 def prepare_event_counters(self, *, forward: bool) -> Any:
335 """Reset per-call counters while preserving the ready generation."""
336 if forward:
337 field_name = "forward_event_counters"
338 events = self.forward_event_counters
339 ready_initialized = self.forward_ready_initialized
340 else:
341 field_name = "backward_event_counters"
342 events = self.backward_event_counters
343 ready_initialized = self.backward_ready_initialized
344 if events is None:
345 raise RuntimeError(f"MegaMoe workspace {field_name} is not initialized.")
346 events[:self.event_counter_bytes].zero_()
347 if events.numel() == self.event_counter_bytes or ready_initialized:
348 return events
350 # Initialize the persistent generation before any peer can signal it.
351 events[self.event_counter_bytes:].zero_()
352 shmem.host_barrier()
353 if forward:
354 self.forward_ready_initialized = True
355 else:
356 self.backward_ready_initialized = True
357 return events
359 def wait_for_reuse(self) -> None:
360 """Order count exchange after the previous workspace lease."""
361 with self._access(), self.lock:
362 if self.used:
363 torch.npu.current_stream(self.device).wait_event(self.completion_event)
365 def claim(self) -> None:
366 """Claim the serial workspace and order it after the previous stream."""
367 with self._access(), self.lock:
368 if self.closing:
369 raise RuntimeError("cannot claim a closing MegaMoe workspace.")
370 if self.in_use:
371 raise RuntimeError("MegaMoe execution resources do not support concurrent calls.")
372 if self.used:
373 torch.npu.current_stream().wait_event(self.completion_event)
374 self.in_use = True
376 def release(self) -> None:
377 """Record completion ordering and release the serial workspace lease."""
378 with self._access(), self.lock:
379 if not self.in_use:
380 raise RuntimeError("MegaMoe workspace was released without an active lease.")
381 self.completion_event.record(torch.npu.current_stream())
382 self.in_use = False
383 self.used = True
385 def _free_symmetric_tensors(self) -> None:
386 """Free each unique SHMEM allocation and invalidate its tensor view."""
387 if self.closing:
388 remaining = tuple(tensor is not None for tensor in (
389 self.source_buffer, self.expert_buffer, self.routed_buffer,
390 self.forward_event_counters, self.backward_event_counters))
391 peers = _collective.exchange(
392 remaining,
393 group=self._root_group(),
394 failure_prefix="MegaMoe symmetric workspace release failed while converging rank results",
395 )
396 if any(peer != remaining for peer in peers):
397 raise RuntimeError("partial symmetric teardown differs across ranks; restart the process")
398 for field_name in (
399 "source_buffer",
400 "expert_buffer",
401 "routed_buffer",
402 "forward_event_counters",
403 "backward_event_counters",
404 ):
405 tensor = getattr(self, field_name)
406 if tensor is None:
407 continue
408 shmem.free(tensor)
409 setattr(self, field_name, None)
410 self.expert_capacity = 0
411 self.routed_slots = 0
412 self.forward_ready_initialized = False
413 self.backward_ready_initialized = False
415 def _free_local_tensors(self) -> None:
416 """Release local-only workspaces after all queued kernels complete."""
417 for field_name in ("gmm_workspace", "swiglu_grad_workspace"):
418 tensor = getattr(self, field_name)
419 if tensor is None:
420 continue
421 tensor.untyped_storage().resize_(0)
422 setattr(self, field_name, None)
424 def _root_group(self) -> Any:
425 """Require an owning manager for collective teardown without a default-group fallback."""
426 if self.heap_manager is None:
427 raise RuntimeError("MegaMoe workspace collective teardown requires an owning HeapManager")
428 return self.heap_manager.root_group
430 def close(self) -> None:
431 """Synchronize and release physical buffers on the owning HeapManager's Root."""
432 with self._access(), self.lock:
433 if self.in_use:
434 raise RuntimeError("cannot close MegaMoe workspace during an active call.")
435 if not self.can_close():
436 raise RuntimeError("cannot close MegaMoe workspace while backward graphs still need its buffers.")
437 empty = all(tensor is None for tensor in (
438 self.source_buffer,
439 self.expert_buffer,
440 self.routed_buffer,
441 self.forward_event_counters,
442 self.backward_event_counters,
443 self.gmm_workspace,
444 self.swiglu_grad_workspace,
445 self.completion_event,
446 ))
447 root_group = self._root_group()
448 self.closing = True
449 self.route_cache = None
450 self.route_header_cache = None
451 self.route_cache_probe_after = 0
452 # An empty rank must still participate when peers retain buffers after a failure.
453 peers = _collective.exchange(
454 empty, group=root_group, failure_prefix="MegaMoe workspace close failed while converging rank results",
455 )
456 if all(peers):
457 self.closing = True
458 return
459 _collective.collective_call(
460 lambda: None if empty else torch.npu.synchronize(self.device),
461 stage="MegaMoe workspace synchronization",
462 group=root_group,
463 )
464 # Workspace teardown requires each collective barrier to complete before
465 # the following free or local Tensor release.
466 shmem.host_barrier()
467 _collective.collective_call(
468 self._free_symmetric_tensors,
469 stage="MegaMoe symmetric workspace release",
470 group=root_group,
471 )
472 shmem.host_barrier()
473 _collective.collective_call(
474 self._free_local_tensors,
475 stage="MegaMoe local workspace release",
476 group=root_group,
477 )
478 self.dtype = None
479 self.device = None
480 self.completion_event = None
481 self.used = False
482 self.closing = True