Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / scheduler / config.py: 96%
229 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"""
16Runtime configuration: topology parameters, C++ ABI structs, and constants.
18Combines:
19- TaskSplitValue / init_task_split_value — TP/EP/seq topology + per-rank counters
20- ctypes.Structure mirrors of runtime_head.hpp — passed as binary blobs to C++ kernels
21"""
22import ctypes
23from dataclasses import dataclass
24from enum import IntEnum
25from typing import Any
27# ── Constants (match runtime_head.hpp exactly) ────────────────────────────────
28MAX_TENSOR_DIMS = 4
29MAX_INPUTS_PER_TASK = 4
30MAX_OUTPUTS_PER_TASK = 4
32MIN_EVENT_CAPACITY = 1024
33# Minimum Cube worker slots in the serialized runtime.
34MIN_CUBE_WORKER_CAPACITY = 24
35MIN_GROUP_LIST_CAPACITY = 512
36GROUP_LIST_CACHE_LINE_BYTES = 128
37ATOMIC_ADD_VALUE_LEN = 8
38READY_CACHE_LINE_BYTES = 64
39INVALID_PROFILE_DESC_ID = 0xFFFFFFFF
40INVALID_PROFILE_OWNER_ID = 0xFFFFFFFF
41DEFAULT_DEPENDENCY_POLL_INTERVAL_US = 0
42FAST_DEPENDENCY_POLL_INTERVAL_US = 5
43EVENT_INVALID_ID = 0xFFFFFFFF
44RUNTIME_PROTOCOL_VERSION = 1
47def mega_moe_event_capacity(num_experts: int, ep_size: int) -> int:
48 """Reserve graph events and the complete atomic write at the last trigger.
50 Args:
51 num_experts: Global expert count.
52 ep_size: Expert-parallel group size.
54 Returns:
55 Cache-line-aligned event slots, with a minimum of 1024.
56 """
57 if ep_size <= 0 or num_experts <= 0 or num_experts % ep_size:
58 raise ValueError("num_experts must be positive and divisible by ep_size")
59 # Ready and worker completion follow termination with separate atomic lanes.
60 final_event = num_experts + 3 * (num_experts // ep_size) + 4 + ATOMIC_ADD_VALUE_LEN
61 required = final_event + ATOMIC_ADD_VALUE_LEN
62 return max(MIN_EVENT_CAPACITY, (required + 15) // 16 * 16)
65# ── Enums ─────────────────────────────────────────────────────────────────────
66class TaskAiCoreType(IntEnum):
67 """Worker core categories encoded in RuntimeConfig."""
69 TASK_AICORE_INVALID = 0
70 TASK_AICORE_CUBE = 1
71 TASK_AICORE_VECTOR = 2
72 TASK_AICORE_MIX = 3
75class TaskType(IntEnum):
76 """Task operation kinds understood by the Device scheduler."""
78 TASK_TERMINATE = 0
79 TASK_BEGIN_TASK_GRAPH = 10
80 TASK_ADD_CUSTOM = 101
81 TASK_SWI_GLU = 102
82 TASK_MATMUL = 103
83 TASK_GROUPED_MATMUL = 104
84 TASK_SHMEM_PUT_MEM_SIGNAL = 105
85 TASK_SWI_GLU_GRAD = 106
86 TASK_MHC_POST = 107
87 TASK_MHC_NORM_CAST = 108
88 TASK_MHC_PROJECTION = 109
89 TASK_MHC_MAPPING = 110
90 TASK_RMS_NORM = 111
91 TASK_MHC_INPUT_MIX = 112
92 TASK_RMS_NORM_GRAD = 113
93 TASK_MHC_GRAD_PREV_A = 114
94 TASK_MHC_GRAD_MAPPING = 115
95 TASK_MHC_GRAD_PHI_RMS = 116
96 TASK_MHC_GRAD_PREV_X_AND_POST = 117
97 TASK_MHC_POST_GRAD = 118
98 TASK_MHC_GRAD_PREV_A_AND_MAPPING = 119
99 TASK_GATE_SOFTPLUS = 120
100 TASK_GATE_SQRT = 121
101 TASK_GATE_ADD_BIAS = 122
102 TASK_GATE_TOPK = 123
103 TASK_GATE_GATHER = 124
104 TASK_GATE_REDUCE_SUM = 125
105 TASK_GATE_ADD_EPSILON = 126
106 TASK_GATE_DIV = 127
107 TASK_GATE_MUL_SCALE = 128
108 TASK_GATE_CAST_INDEX = 129
109 TASK_GATE_GRAD_MULS_SCALE = 130
110 TASK_GATE_GRAD_BROADCAST_DENOMINATOR = 131
111 TASK_GATE_GRAD_NEG = 132
112 TASK_GATE_GRAD_DIV_SELECTED = 133
113 TASK_GATE_GRAD_DIV_SELECTED_RATIO = 134
114 TASK_GATE_GRAD_MUL_CROSS = 135
115 TASK_GATE_GRAD_DIV_DIRECT = 136
116 TASK_GATE_GRAD_REDUCE_SUM = 137
117 TASK_GATE_GRAD_BROADCAST_ROW_SUM = 138
118 TASK_GATE_GRAD_ADD_SELECTED = 139
119 TASK_GATE_GRAD_ZEROS = 140
120 TASK_SHMEM_GET_MEM = 141
123class EventType(IntEnum):
124 """Dependency and trigger event operations used by scheduled tasks."""
126 EVENT_EMPTY = 900
127 EVENT_LAUNCH_TASKS = 901
128 EVENT_LAUNCH_MASSIVE_TASKS = 902
129 EVENT_LAUNCH_DEPENDENT_TASKS = 903
130 EVENT_END_OF_TASK_GRAPH = 910
131 EVENT_TERMINATION = 911
132 EVENT_INVALID = 999
135class DynamicType(IntEnum):
136 """Runtime dynamic-data operations applied before task execution."""
138 DYNAMIC_EMPTY = 0
139 DYNAMIC_DSV3_MOE = 101
142# ── ctypes Structures (mirror runtime_head.hpp) ───────────────────────────────
144class TensorDescC(ctypes.Structure):
145 """Serialized tensor address and shape descriptor."""
147 _fields_ = [
148 ("tensor_type", ctypes.c_uint32),
149 ("num_dims", ctypes.c_uint32),
150 ("dim", ctypes.c_uint32 * MAX_TENSOR_DIMS),
151 ("stride", ctypes.c_uint32 * MAX_TENSOR_DIMS),
152 ("data_type", ctypes.c_uint32),
153 ("input_position", ctypes.c_uint32),
154 ("base_ptr_offset", ctypes.c_uint32),
155 ("transpose_flag", ctypes.c_uint32),
156 ("dynamic_shape", ctypes.c_uint32),
157 ("dynamic_dim", ctypes.c_uint32),
158 ]
161class TaskDescC(ctypes.Structure):
162 """Serialized task descriptor shared by Host and Device schedulers."""
163 _fields_ = [
164 ("task_type", ctypes.c_uint32),
165 ("task_aicore_type", ctypes.c_uint32),
166 ("num_inputs", ctypes.c_uint32),
167 ("num_outputs", ctypes.c_uint32),
168 ("trigger_event", ctypes.c_uint32),
169 ("dependent_event", ctypes.c_uint32),
170 ("inputs", TensorDescC * MAX_INPUTS_PER_TASK),
171 ("outputs", TensorDescC * MAX_OUTPUTS_PER_TASK),
172 ("tiling_data_position", ctypes.c_uint32),
173 ("tiling_data_offset", ctypes.c_uint32),
174 ("task_index", ctypes.c_uint32),
175 ("task_split_num", ctypes.c_uint32),
176 ("task_split_value", ctypes.c_uint32),
177 ("extra_value_0", ctypes.c_uint32),
178 ("extra_value_1", ctypes.c_uint32),
179 ("extra_value_2", ctypes.c_uint32),
180 ("profile_desc_id", ctypes.c_uint32),
181 ("profile_owner_id", ctypes.c_uint32),
182 ]
184 def __init__(self, *args: Any, **kwargs: Any) -> None:
185 """Initialize optional profiling metadata to explicit invalid sentinels."""
186 has_profile_desc = len(args) >= len(self._fields_) - 1 or "profile_desc_id" in kwargs
187 has_profile_owner = len(args) >= len(self._fields_) or "profile_owner_id" in kwargs
188 super().__init__(*args, **kwargs)
189 if not has_profile_desc:
190 self.profile_desc_id = INVALID_PROFILE_DESC_ID
191 if not has_profile_owner:
192 self.profile_owner_id = INVALID_PROFILE_OWNER_ID
195class EventDescC(ctypes.Structure):
196 """Serialized event operation descriptor."""
198 _fields_ = [
199 ("event_type", ctypes.c_uint32),
200 ("num_triggers", ctypes.c_uint32),
201 ("first_task_id", ctypes.c_uint32),
202 ("last_task_id", ctypes.c_uint32),
203 ]
206class DynamicDataC(ctypes.Structure):
207 """Serialized dynamic-data update descriptor."""
209 _fields_ = [
210 ("dynamic_type", ctypes.c_uint32),
211 ("dynamic_input_position", ctypes.c_uint32),
212 ("dynamic_group_size", ctypes.c_uint32),
213 ("dynamic_max_seq_len", ctypes.c_int32),
214 ]
217class RuntimeConfigC(ctypes.Structure):
218 """Common 64-byte header; allocated subclasses append graph-sized arrays."""
220 _fields_ = [
221 ("task_num", ctypes.c_uint32),
222 ("num_workers", ctypes.c_uint32),
223 ("task_capacity", ctypes.c_uint32),
224 ("event_capacity", ctypes.c_uint32),
225 ("ready_event", ctypes.c_uint32),
226 ("diagnostic_mode", ctypes.c_uint32),
227 ("aic_profile_record_capacity", ctypes.c_uint32),
228 ("aiv_profile_record_capacity", ctypes.c_uint32),
229 ("completion_event", ctypes.c_uint32),
230 ("protocol_version", ctypes.c_uint32),
231 ("_padding", ctypes.c_uint32 * 6),
232 ]
235class TilingDataC(ctypes.Structure):
236 """Tiling data for add_custom ops (not GMM)."""
237 _fields_ = [
238 ("smallCoreDataNum", ctypes.c_uint64),
239 ("bigCoreDataNum", ctypes.c_uint64),
240 ("ubPartDataNum", ctypes.c_uint64),
241 ("smallCoreTailDataNum", ctypes.c_uint64),
242 ("bigCoreTailDataNum", ctypes.c_uint64),
243 ("smallCoreLoopNum", ctypes.c_uint64),
244 ("bigCoreLoopNum", ctypes.c_uint64),
245 ("tailBlockNum", ctypes.c_uint64),
246 ]
249class SwiGluTilingDataC(ctypes.Structure):
250 """Serialized SwiGLU tiling values for one worker."""
252 _fields_ = [
253 ("is32BAligned", ctypes.c_uint32),
254 ("isDoubleBuffer", ctypes.c_uint32),
255 ("rowLen", ctypes.c_uint64),
256 ("colLen", ctypes.c_uint64),
257 ("baseRowLen", ctypes.c_uint32),
258 ("baseColLen", ctypes.c_uint32),
259 ("activateLeft", ctypes.c_uint32),
260 ("biasIsEmpty", ctypes.c_uint32),
261 ("quantScaleIsEmpty", ctypes.c_uint32),
262 ("activateScaleIsEmpty", ctypes.c_uint32),
263 ("swiColLen", ctypes.c_uint64),
264 ("perRowLen", ctypes.c_uint64),
265 ("modRowLen", ctypes.c_uint64),
266 ("usedCoreNum", ctypes.c_uint32),
267 ]
270# ── Topology + runtime counters ───────────────────────────────────────────────
272@dataclass
273class TaskSplitValue:
274 """
275 Hardware topology parameters + per-rank runtime counters.
277 Topology fields are user-supplied; derived properties compute
278 sequence/expert partition sizes. Runtime counters are reset by
279 init_task_split_value() before each rank's fill loop.
280 """
281 # ── User inputs ───────────────────────────────────────────────────────────
282 tp: int = 4
283 ep: int = 4
284 seq_size: int = 8192
285 all_expert_num: int = 32
286 top_k: int = 8
287 dispatch_mode: str = "push"
289 def __post_init__(self) -> None:
290 """Reject topology values that would invalidate runtime arithmetic."""
291 if self.dispatch_mode not in ("push", "pull"):
292 raise ValueError("dispatch_mode must be push or pull")
293 values = {
294 "tp": self.tp,
295 "ep": self.ep,
296 "seq_size": self.seq_size,
297 "all_expert_num": self.all_expert_num,
298 "top_k": self.top_k,
299 }
300 for name, value in values.items():
301 if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
302 raise ValueError(f"{name} must be a positive integer, got {value!r}.")
303 if self.top_k > self.all_expert_num:
304 raise ValueError(
305 f"top_k ({self.top_k}) cannot exceed all_expert_num ({self.all_expert_num})."
306 )
307 if self.all_expert_num % self.ep:
308 raise ValueError(
309 f"all_expert_num ({self.all_expert_num}) must be divisible by ep ({self.ep})."
310 )
312 # ── Derived properties ────────────────────────────────────────────────────
313 @property
314 def single_rank_expert_num(self) -> int:
315 """Number of experts assigned to a single rank."""
316 return self.all_expert_num // self.ep
318 @property
319 def seq_all(self) -> int:
320 """Total sequence length after accounting for EP and TP parallelism."""
321 return (self.seq_size * self.ep * self.top_k) // self.tp
323 @property
324 def per_expert_seq(self) -> int:
325 """Sequence length per expert (across all EP ranks)."""
326 return self.seq_all // self.top_k
328 @property
329 def per_rank_seq(self) -> int:
330 """Sequence length assigned to each rank."""
331 return self.seq_all // self.ep
333 @property
334 def per_expert_seq_to_other(self) -> int:
335 """Sequence length per expert when sending to other ranks."""
336 return self.seq_all // (self.ep * self.top_k)
338 @property
339 def all_event_num(self) -> int:
340 """Total number of events needed for the schedule."""
341 e = self.single_rank_expert_num
342 return 1 + self.all_expert_num + e + e + e
344 # ── Runtime counters (reset by init_task_split_value per rank) ────────────
345 rank_id: int = 0
346 pre_pre_event_num: int = 0
347 pre_event_num: int = 0
348 pre_task_num: int = 0
349 pre_cube_task_num: int = 0
350 pre_vector_task_num: int = 0
351 pre_mix_task_num: int = 0
354def init_task_split_value(tsv: TaskSplitValue) -> None:
355 """Reset per-rank runtime counters to zero.
357 Args:
358 tsv: Task split values whose runtime counters are reset.
359 """
360 tsv.pre_pre_event_num = 0
361 tsv.pre_event_num = 0
362 tsv.pre_task_num = 0
363 tsv.pre_cube_task_num = 0
364 tsv.pre_vector_task_num = 0
365 tsv.pre_mix_task_num = 0
368def _validate_num_cube_cores(num_cube_cores: int) -> None:
369 """Validate the hardware core count used by GMM task descriptors."""
370 if not isinstance(num_cube_cores, int) or isinstance(num_cube_cores, bool) or num_cube_cores <= 0:
371 raise ValueError(f"num_cube_cores must be a positive integer, got {num_cube_cores!r}.")
374def _validate_task_bounds(task: TaskDescC, descriptor_index: int) -> None:
375 """Validate the common task index range."""
376 if task.task_split_num == 0 or task.task_index >= task.task_split_num:
377 raise ValueError(
378 f"task {descriptor_index} has invalid task_index/task_split_num "
379 f"({task.task_index}/{task.task_split_num})."
380 )
383def _validate_swiglu_task(task: TaskDescC, descriptor_index: int, group_size: int) -> None:
384 """Validate SwiGLU expert partition arithmetic."""
385 if task.task_split_num < group_size or task.task_split_num % group_size:
386 raise ValueError(
387 f"SwiGLU task {descriptor_index} cannot be partitioned over {group_size} experts."
388 )
391def _validate_gmm_task(task: TaskDescC, descriptor_index: int, expected_tasks: int) -> None:
392 """Validate that every local expert owns one task per Cube core."""
393 if task.task_split_num != expected_tasks:
394 raise ValueError(
395 f"GMM task {descriptor_index} has task_split_num {task.task_split_num}; "
396 f"expected {expected_tasks}."
397 )
400def _validate_shmem_task(task: TaskDescC, descriptor_index: int, tsv: TaskSplitValue) -> None:
401 """Validate SHMEM task partitioning and remote-rank selection."""
402 if task.task_split_num < tsv.all_expert_num or task.task_split_num % tsv.all_expert_num:
403 raise ValueError(
404 f"SHMEM task {descriptor_index} cannot be partitioned over "
405 f"{tsv.all_expert_num} experts."
406 )
407 tasks_per_expert = task.task_split_num // tsv.all_expert_num
408 target_rank_offset = task.task_index // (tsv.single_rank_expert_num * tasks_per_expert)
409 if target_rank_offset >= tsv.ep:
410 raise ValueError(
411 f"SHMEM task {descriptor_index} targets rank offset {target_rank_offset}; "
412 f"ep is {tsv.ep}."
413 )
416def _validate_runtime_task_queues(cfg: RuntimeConfigC, event_capacity: int) -> None:
417 """Validate scheduled task IDs and their event ranges."""
418 queues = (cfg.cube_task_indices, cfg.vector_task_indices, cfg.mix_task_indices)
419 for indices, count in zip(queues, cfg.task_index_num[:3]):
420 if count < 0 or count > len(indices):
421 raise ValueError(f"invalid runtime task-index count {count}")
422 for task_id in indices[:count]:
423 if task_id < 0 or task_id >= len(cfg.all_tasks):
424 raise ValueError(f"invalid runtime task id {task_id}")
425 task = cfg.all_tasks[task_id]
426 dependency_out_of_range = (
427 task.dependent_event != 0xFFFFFFFF
428 and task.dependent_event >= event_capacity
429 )
430 trigger_out_of_range = task.trigger_event + ATOMIC_ADD_VALUE_LEN > event_capacity
431 if dependency_out_of_range or trigger_out_of_range:
432 raise ValueError(
433 f"task {task_id} references an event outside allocated capacity {event_capacity}"
434 )
437def validate_runtime_config(
438 cfg: RuntimeConfigC,
439 tsv: TaskSplitValue,
440 num_cube_cores: int,
441) -> None:
442 """Validate task arithmetic before serializing a runtime configuration.
444 The device worker intentionally assumes that generated task descriptors are
445 internally consistent. Rejecting an invalid schedule here avoids both
446 out-of-bounds descriptor access and device-side early exits that could omit
447 event-counter or SHMEM-signal updates.
449 Args:
450 cfg: Populated runtime configuration awaiting serialization.
451 tsv: Validated topology and task split values.
452 num_cube_cores: Number of Cube cores used to generate GMM tasks.
454 Raises:
455 ValueError: If the runtime configuration violates a device-side invariant.
456 """
457 _validate_num_cube_cores(num_cube_cores)
458 if cfg.num_workers != 2 * num_cube_cores:
459 raise ValueError(
460 f"num_workers ({cfg.num_workers}) must equal twice num_cube_cores ({num_cube_cores})."
461 )
462 if cfg.task_num > len(cfg.all_tasks):
463 raise ValueError(f"task_num ({cfg.task_num}) exceeds allocated task capacity ({len(cfg.all_tasks)}).")
464 event_capacity = len(cfg.all_event_num_triggers)
465 _validate_runtime_task_queues(cfg, event_capacity)
467 local_experts = tsv.single_rank_expert_num
468 group_size = cfg.dynamic_data.dynamic_group_size
469 if group_size != local_experts:
470 raise ValueError(
471 f"dynamic_group_size ({group_size}) must equal local expert count ({local_experts})."
472 )
474 for descriptor_index in range(cfg.task_num):
475 task = cfg.all_tasks[descriptor_index]
476 task_type = task.task_type
477 if task_type not in (
478 TaskType.TASK_SWI_GLU,
479 TaskType.TASK_SWI_GLU_GRAD,
480 TaskType.TASK_GROUPED_MATMUL,
481 TaskType.TASK_SHMEM_PUT_MEM_SIGNAL,
482 TaskType.TASK_SHMEM_GET_MEM,
483 ):
484 continue
485 _validate_task_bounds(task, descriptor_index)
487 if task_type in (TaskType.TASK_SWI_GLU, TaskType.TASK_SWI_GLU_GRAD):
488 _validate_swiglu_task(task, descriptor_index, group_size)
489 elif task_type == TaskType.TASK_GROUPED_MATMUL:
490 _validate_gmm_task(task, descriptor_index, num_cube_cores * local_experts)
491 else:
492 _validate_shmem_task(task, descriptor_index, tsv)
495def configure_ready_handshake(cfg: RuntimeConfigC, tsv: TaskSplitValue) -> None:
496 """Reserve entry readiness and completion of all participating workers.
498 Args:
499 cfg: Allocated graph configuration whose counters include ready padding.
500 tsv: Topology that determines whether peer readiness is needed.
501 """
502 ready_event = tsv.all_event_num + 3
503 completion_event = ready_event + ATOMIC_ADD_VALUE_LEN
504 if completion_event + ATOMIC_ADD_VALUE_LEN > cfg.event_capacity:
505 raise ValueError("ready handshake event exceeds event_capacity.")
506 cfg.ready_event = ready_event if tsv.ep > 1 else 0
507 cfg.protocol_version = RUNTIME_PROTOCOL_VERSION if tsv.dispatch_mode == "pull" else 0
508 cfg.completion_event = completion_event if tsv.dispatch_mode == "pull" else 0
509 cfg.all_event_num_triggers[ready_event] = 1
510 cfg.all_event_num_triggers[completion_event] = cfg.num_workers
513def event_workspace_bytes(ep_size: int, num_experts: int) -> int:
514 """Return graph counters plus separate ready and completion generations.
516 Args:
517 ep_size: Number of expert-parallel ranks sharing the symmetric arena.
518 num_experts: Global expert count used to size ordinary event counters.
520 Returns:
521 Required byte count for one direction's event workspace.
522 """
523 counter_bytes = mega_moe_event_capacity(num_experts, ep_size) * ctypes.sizeof(ctypes.c_int32)
524 if ep_size <= 1:
525 return counter_bytes
526 return counter_bytes + 2 * (ep_size + 1) * READY_CACHE_LINE_BYTES