Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / parallel.py: 0%

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"""Lower the frozen parallel plan into explicit forward code. 

16 

17The lowerer sinks forward wrappers into the source. Which form a boundary 

18class's ``forward`` takes is decided by 

19:func:`hyper_parallel.codegen.plan.boundary_forms.classify_boundary_form` — 

20the same pure classifier the runtime re-runs at install time and preflight 

21re-runs at train time, so the three stages can never disagree about a 

22boundary's form. 

23 

24Class vs FQN: the frozen plan is keyed by boundary FQN 

25(``model.layers.0.self_attn``), but a ``class.forward`` is shared by many 

26instances. The lowerer groups boundary FQNs by their class (meta's 

27``boundary_classes``), enforces that every FQN of a class carries the SAME 

28boundary contract, and emits one rewritten ``forward`` per class. 

29 

30The original ``forward`` is kept 

31verbatim as the class-private ``_forward_impl`` (extracted by a zero-width 

32insert before the ``def forward`` — a fresh-source operation, so the offsets 

33stay valid when combined with the body replacement), and the public 

34``forward`` is rewritten in one of the classified forms: 

35 

36- ``identity`` (pruned): every declared transition is identity on the plan's 

37 active axes — the class forward is NOT rewritten at all. The runtime's 

38 install path covers the exact no-op / to-local semantics with its generic 

39 wrapper, so pruning is byte-for-byte equivalent to the previous rewrite. 

40- ``tp_collective`` (static bare operators): every declared transition is 

41 identity or TP-lowerable, and the source-side structural gates pass 

42 (single-value return, no ``*args``/``**kwargs``, declared inputs bind to 

43 forward parameters). The forward names its form with the class attribute 

44 ``_hp_boundary_form = "tp_collective"`` and calls bare operators on 

45 ``self._hp_tp`` (bound at install time after the runtime re-validates 

46 the compiled plan against the live mesh; on mismatch the runtime replaces 

47 the forward with the generic engine — the baked calls never run unvalidated). 

48- ``generic``: redistribute in, ``_forward_impl``, redistribute out, through 

49 the instance-bound compiled plan (``self._hp_boundary``, bound once by 

50 ``hp_install_boundaries``), with a declarative comment block annotating 

51 what each side's transitions lower to. 

52- ``region`` (EP MoE local region): redistribute in, 

53 ``__hp_compute__`` (bound by ``hp_parallelize``), re-wrap locals per 

54 ``out_src``, redistribute out. 

55 

56A CP inner-wrapper boundary is ``generic`` in the generated code: the wrapper 

57is installed at runtime by ``hp_apply_inner_wrapper`` (it mutates 

58``target.forward`` in place — not a bindable callable), and runs inside the 

59``_forward_impl`` call. 

60 

61The lowerer is a pure function over the source text + the frozen plan. It 

62returns the patched text. Structural work (finding a class / its forward / 

63the import block) is delegated to :mod:`astkit`. 

64""" 

65from __future__ import annotations 

66 

67import json 

68import logging 

69from typing import Any, Optional 

70 

71from hyper_parallel.codegen.astkit.edits import ( 

72 TextEdit, 

73 apply_edits, 

74 insert_after_imports, 

75 replace_function_body, 

76) 

77from hyper_parallel.codegen.astkit.index import ( 

78 FunctionInfo, 

79 SourceIndex, 

80 build_source_index, 

81 returns_any_value, 

82 returns_single_value, 

83) 

84from hyper_parallel.codegen.plan.boundary_forms import ( 

85 FORM_GENERIC, 

86 FORM_IDENTITY, 

87 FORM_REGION, 

88 FORM_TP_COLLECTIVE, 

89 BoundaryForm, 

90 TransitionOp, 

91 classify_boundary_form, 

92 is_boundary_entry, 

93) 

94from hyper_parallel.codegen.plan.freeze import parse_placement 

95 

96#: Injection keys that describe how a boundary's compute is regioned. 

97_INJECTION_KEYS = ("inner_wrapper", "inner_target", "inner_out_src", "local_compute_fn") 

98 

99 

100def lower_forward_boundaries( 

101 source_text: str, 

102 frozen_plan: Any, 

103 *, 

104 boundary_classes: Optional[dict[str, str]] = None, 

105 module_name: str = "", 

106) -> str: 

107 """Rewrite each boundary class's ``forward`` to call the codegen runtime. 

108 

109 ``frozen_plan`` is the meta/param-plan dict or a ``CodegenMeta``; the 

110 plan entries are read via ``_plan_field``. ``boundary_classes`` 

111 is the FQN -> class name map (``meta.boundary_classes``); when omitted it 

112 degrades to the class name looked up from the source index itself, which is 

113 only correct if every boundary FQN's leaf class name equals the key's tail — 

114 ``_class_for_fqn``. The plan's active axes (``mesh_dim_names``) drive 

115 :func:`classify_boundary_form`; the classified form plus the source-side 

116 structural gates decide which template each class gets (see the module 

117 docstring). 

118 

119 Returns the patched source text. This function only rewrites the 

120 ``forward`` bodies; the rewritten bodies reference the ``_hp_boundary`` 

121 / ``_hp_tp`` instance attributes that ``hp_install_boundaries`` binds 

122 at runtime — no per-boundary global constant is emitted. 

123 """ 

124 edits: list[TextEdit] = [] 

125 needs_runtime_helper_import = False 

126 for class_name, form, emitted, func, injection, deferred_bias in iter_emitted_forms( 

127 source_text, 

128 frozen_plan, 

129 boundary_classes=boundary_classes, 

130 module_name=module_name, 

131 ): 

132 if emitted == FORM_IDENTITY: 

133 # Pruned: the runtime's install path covers the exact no-op / 

134 # to-local semantics with its generic wrapper, so leaving the 

135 # original forward untouched is behavior-preserving. 

136 if deferred_bias: 

137 # D-22: an unrewritten forward never calls the exit hook, so 

138 # the suppressed bias would silently vanish — fail fast 

139 # instead (the planner defers only rowwise-bias exits that 

140 # reduce, which classify as non-identity, so this is a 

141 # mis-declared plan, not a supported shape). 

142 raise NotImplementedError( 

143 f"lower_forward_boundaries: boundary class {class_name!r} " 

144 "prunes to identity but declares deferred bias params; an " 

145 "identity forward has no exit hook to re-add them" 

146 ) 

147 continue 

148 # Every emitted forward body (generic, local-region) ends with the 

149 # runtime helper ``hp_to_local_if_dtensor``; that helper is a 

150 # dependency of the templates, not of any declaration, so the lowerer 

151 # must import it itself — otherwise the artifact is not self-contained. 

152 needs_runtime_helper_import = True 

153 body = _build_forward_body( 

154 source_text, injection, func, form, emitted, deferred_bias 

155 ) 

156 # Keep the original forward as ``_forward_impl``. The 

157 # extracted method is inserted BEFORE the ``def forward`` (zero-width 

158 # edit at ``def_offset``), and the rewritten body replaces the original 

159 # body at ``body_start`` — the two spans never overlap, so the 

160 # incremental-shift validation in ``_apply_edits`` passes. The static 

161 # form additionally prefixes the class marker naming the form. 

162 marker = _TP_FORM_MARKER_LINE if emitted == FORM_TP_COLLECTIVE else None 

163 impl = _build_forward_impl_edit(source_text, func, prefix=marker) 

164 if impl is not None: 

165 edits.append(impl) 

166 edits.extend(replace_function_body(source_text, func, body, indent="")) 

167 if not edits: 

168 return source_text 

169 

170 nested = _validate_no_circular(edits) 

171 if nested: 

172 raise ValueError( 

173 "lower_forward_boundaries: edits overlap (circular class nesting); " 

174 "cannot safely rewrite: " + ", ".join(nested) 

175 ) 

176 patched = apply_edits(source_text, edits) 

177 if needs_runtime_helper_import: 

178 patched = _add_runtime_helper_import(patched) 

179 return patched 

180 

181 

182#: Runtime helper the emitted forward templates call. The templates (and 

183#: this constant) must stay in sync — see ``_render_local_compute`` and 

184#: ``_render_output_redistribute``. 

185_RUNTIME_HELPER_IMPORT = ( 

186 "from hyper_parallel.codegen.runtime import hp_to_local_if_dtensor" 

187) 

188 

189 

190def _add_runtime_helper_import(source_text: str) -> str: 

191 """Ensure the forward templates' runtime helper is imported. 

192 

193 The helper is a template dependency, so nothing in the declarations 

194 contributes it; without this the generated artifact raises ``NameError`` at 

195 the first forward call. No-op when the source already imports it. 

196 """ 

197 if "import hp_to_local_if_dtensor" in source_text: 

198 return source_text 

199 index = build_source_index(source_text) 

200 edit = insert_after_imports( 

201 source_text, 

202 index.import_end, 

203 _RUNTIME_HELPER_IMPORT, 

204 default_offset=index.docstring_end or 0, 

205 ) 

206 return apply_edits(source_text, [edit]) if edit is not None else source_text 

207 

208 

209#: Class attribute a ``tp_collective`` forward carries, naming its form so the 

210#: runtime install path and preflight can recognize the static template. 

211TP_FORM_ATTRIBUTE = "_hp_boundary_form" 

212#: The marker value naming the statically lowered TP-collective form. 

213TP_FORM_MARKER = "tp_collective" 

214#: The emitted marker statement (method indent, followed by a blank line). 

215_TP_FORM_MARKER_LINE = f' {TP_FORM_ATTRIBUTE} = "{TP_FORM_MARKER}"\n\n' 

216 

217 

218def iter_emitted_forms( 

219 source_text: str, 

220 frozen_plan: Any, 

221 *, 

222 boundary_classes: Optional[dict[str, str]] = None, 

223 module_name: str = "", 

224): 

225 """Yield every boundary class's emit decision over ``source_text``. 

226 

227 Single decision path shared by the emitter 

228 (:func:`lower_forward_boundaries`) and the train-time preflight 

229 (``check/preflight.verify_boundary_forms``): classification 

230 (:func:`classify_boundary_form`) plus the structural gates 

231 (:func:`resolve_emitted_form`), per boundary class. Yields 

232 ``(class_name, form, emitted, func, injection, deferred_bias)`` — 

233 ``form`` the classifier's verdict, ``emitted`` the form actually 

234 rendered after the gates, ``func`` the class's ``forward`` 

235 :class:`FunctionInfo`, ``injection`` the class's frozen injection rule 

236 (``None`` when absent), and ``deferred_bias`` whether the class's 

237 canonical entry carries D-22 deferred bias params (the emitted exit 

238 then calls ``self._hp_deferred_bias(outputs)`` after the output 

239 redistribution). 

240 

241 Boundary classes are iterated in sorted ``class_name`` order. The 

242 grouping walks ``param_plan``, whose insertion order is NOT stable 

243 across the write/reload cycle: ``meta.param_plan`` is persisted with 

244 ``json.dump(..., sort_keys=True)``, which reorders the keys 

245 alphabetically, so reload-ing the meta yields a different class order 

246 than the in-memory freeze did. Sorting here makes the emitted byte 

247 stream a pure function of which classes are boundaries, independent of 

248 how ``param_plan`` reached us. 

249 """ 

250 index = build_source_index(source_text) 

251 param_plan = _plan_field(frozen_plan, "param_plan") or {} 

252 injections = _plan_field(frozen_plan, "injections") or [] 

253 injection_by_class = _index_injections_by_class(injections, boundary_classes, index) 

254 mesh_dim_names = tuple(_plan_field(frozen_plan, "mesh_dim_names") or ()) 

255 

256 groups = _group_boundaries_by_class( 

257 param_plan, boundary_classes, index, module_name=module_name 

258 ) 

259 for class_name in sorted(groups): 

260 group = groups[class_name] 

261 # Fail fast when one class's FQNs disagree on the contract: one shared 

262 # ``forward`` cannot carry two boundary plans. The canonical entry is 

263 # the per-class contract every FQN agreed on. 

264 canonical_entry = _verify_class_contract(group, class_name) 

265 func = _require_forward(index, class_name) 

266 injection = injection_by_class.get(class_name) 

267 form = classify_boundary_form(canonical_entry, mesh_dim_names, injection) 

268 emitted = resolve_emitted_form(source_text, func, form) 

269 deferred_bias = bool(canonical_entry.get("deferred_bias_params")) 

270 yield class_name, form, emitted, func, injection, deferred_bias 

271 

272 

273def resolve_emitted_form( 

274 source_text: str, 

275 func: FunctionInfo, 

276 form: BoundaryForm, 

277) -> str: 

278 """The form the emitter actually renders for one classified boundary. 

279 

280 The classifier's ``tp_collective`` is an *optimistic* verdict: the static 

281 template additionally requires the source-side structural gates (no 

282 ``*args``, value returns, declared inputs binding to forward parameters, 

283 output ops targeting the first output). A boundary that fails a gate 

284 renders as ``generic`` instead. This resolver is the single decision 

285 shared by the emitter and preflight, so the emitted file and the train-time 

286 check can never disagree about which template a class got. 

287 """ 

288 if form.form in (FORM_IDENTITY, FORM_REGION, FORM_GENERIC): 

289 return form.form 

290 if _static_gates_pass(source_text, func, form): 

291 return FORM_TP_COLLECTIVE 

292 return FORM_GENERIC 

293 

294 

295def _static_gates_pass(source_text: str, func: FunctionInfo, form: BoundaryForm) -> bool: 

296 """Whether the static bare-operator template is admissible for ``func``. 

297 

298 Gates (each mirrors a concrete way the generic engine is more general 

299 than a static rewrite): 

300 

301 1. No ``*args`` — a positional variable tail cannot be enumerated 

302 statically. ``**kwargs`` is allowed and is re-passed by name. 

303 2. Every ``return`` carries a value. Single-value returns receive output 

304 ops directly; tuple returns receive index-0 output ops and are 

305 rebuilt, matching the compiled plan's output-index behavior. 

306 3. Every declared input name binds to a forward parameter — the generic 

307 engine skips an ``in_src`` name the caller never passed (``_get_arg`` 

308 default ``None``); the static template would reference an unbound 

309 variable instead, so a declared name outside the signature disqualifies 

310 the static form. 

311 4. Every non-identity output op targets the first declared output — with a 

312 single-value return the compiled plan applies index-0 ops only; an op 

313 mapped to a later output position cannot be expressed statically. 

314 """ 

315 if func.vararg_name is not None: 

316 return False 

317 if not returns_any_value(source_text, func): 

318 return False 

319 signature_names = set(func.param_names) | set(func.kwonly_names) 

320 if not set(form.in_declared) <= signature_names: 

321 return False 

322 out_names = form.out_names 

323 for name in form.out_ops: 

324 # An op name missing from the declared order compiles to index 0 

325 # (``_compile_output_plan``'s ``name_to_idx.get(name, 0)``), which the 

326 # single-value template covers; a later position does not. 

327 if name in out_names and out_names.index(name) != 0: 

328 return False 

329 return True 

330 

331 

332# --------------------------------------------------------------------------- 

333# grouping / consistency 

334# --------------------------------------------------------------------------- 

335 

336def _group_boundaries_by_class( 

337 param_plan: dict[str, Any], 

338 boundary_classes: Optional[dict[str, str]], 

339 index: SourceIndex, 

340 *, 

341 module_name: str, 

342) -> dict[str, list[tuple[str, dict[str, Any]]]]: 

343 """Group boundary FQNs by the class that owns their ``forward``. 

344 

345 Returns ``{class_name: [(fqn, entry), ...]}`` in a stable order (FQN sorted 

346 for determinism). 

347 

348 A boundary whose class is NOT defined in this source file — ``index`` is 

349 the source index, so ``index.find_class(name) is None`` means the class is 

350 imported (``nn.Embedding`` / ``nn.Linear`` / ``nn.LayerNorm`` for 

351 ``embed_tokens`` / ``lm_head`` / ``norm``) — has no ``def forward`` span 

352 to rewrite, so its FQN is dropped from the grouping and the boundary is 

353 left un-rewritten here. Its boundary contract is still executed at runtime 

354 by ``hp_install_boundaries``: the in/out 

355 redistribution for an imported class is installed as a forward wrapper on 

356 the live module, so the contract is not dropped. Skipping the rewrite here 

357 is exactly what lets the real plan (which marks ``embed_tokens`` / 

358 ``lm_head`` as boundaries) emit without failing on a class that has no 

359 source ``forward``. 

360 

361 The same drop applies to a class that IS defined in this source file but 

362 inherits its ``forward`` from an imported base (an attention subclassing 

363 a component base class): the class body has no ``def forward`` span to 

364 rewrite, and live instances carry no ``_forward_impl``, so the runtime 

365 install path takes its imported-class branch all the same. 

366 

367 An FQN whose param-plan entry carries the explicit ``skip_source_lowering`` 

368 marker (set by meta normalization when the class is shared between 

369 boundary and non-boundary instances — DeepSeek-V3 ``DeepseekV3MLP``: 

370 dense decoder layers and shared experts) is skipped: a class-level 

371 ``forward`` rewrite is inherited by every instance, so the non-boundary 

372 instance would run a boundary body without the attributes 

373 ``hp_install_boundaries`` binds only on boundary instances. The 

374 per-instance runtime wrapper is installed instead and executes the same 

375 boundary contract. 

376 

377 An FQN whose class CANNOT be resolved fails fast — a boundary whose forward 

378 is never rewritten would silently drop its redistribution. 

379 """ 

380 logger = logging.getLogger(__name__) 

381 groups: dict[str, list[tuple[str, dict[str, Any]]]] = {} 

382 for fqn, entry in param_plan.items(): 

383 if not is_boundary_entry(entry): 

384 continue 

385 if entry.get("skip_source_lowering"): 

386 logger.info( 

387 "lower_forward_boundaries: boundary %s is marked " 

388 "skip_source_lowering (class shared with a non-boundary " 

389 "instance); the per-instance runtime wrapper will be " 

390 "installed instead", 

391 fqn, 

392 ) 

393 continue 

394 class_name = _class_for_fqn(fqn, boundary_classes, index, module_name) 

395 if class_name is None: 

396 raise ValueError( 

397 f"lower_forward_boundaries: cannot resolve the class for boundary " 

398 f"{fqn!r} — the frozen plan has no class mapping and the source " 

399 "index does not contain a class whose name matches; a boundary " 

400 "whose forward is not rewritten would silently drop its " 

401 "redistribution" 

402 ) 

403 if index.find_class(class_name) is None: 

404 # Imported class (Embedding/Linear/LayerNorm): no source ``forward`` 

405 # to rewrite; runtime ``hp_install_boundaries`` installs the 

406 # in/out boundary wrapper. Not added to ``groups``. 

407 continue 

408 if index.find_forward(class_name) is None: 

409 # In-source class with an inherited forward (an attention 

410 # subclassing an imported base): no class-body ``def forward`` 

411 # span to rewrite, the same shape an imported class has. Live 

412 # instances carry no ``_forward_impl``, so hp_install_boundaries 

413 # takes its imported-class branch and installs the per-instance 

414 # entry/exit wrapper. 

415 logger.info( 

416 "lower_forward_boundaries: boundary %s is a %s instance whose " 

417 "forward is inherited (no class-body ``def forward``); the " 

418 "per-instance runtime wrapper will be installed instead", 

419 fqn, 

420 class_name, 

421 ) 

422 continue 

423 groups.setdefault(class_name, []).append((fqn, entry)) 

424 for class_name, items in groups.items(): 

425 items.sort(key=lambda item: item[0]) 

426 return groups 

427 

428 

429def _class_for_fqn( 

430 fqn: str, 

431 boundary_classes: Optional[dict[str, str]], 

432 index: SourceIndex, 

433 module_name: str, 

434) -> Optional[str]: 

435 """Resolve a boundary FQN to a top-level class name. 

436 

437 Resolution order: 

438 1. the explicit ``boundary_classes`` map (freeze-time truth); 

439 2. a source-index class whose name matches the FQN's leaf segment; 

440 3. a source-index class whose name matches after stripping a known module 

441 prefix (``module_name``) — remote-code modeling files carry the bare 

442 class name, and the FQN leaf is usually it; 

443 4. ``None`` (caller fails fast). 

444 """ 

445 if boundary_classes: 

446 cls = boundary_classes.get(fqn) 

447 if cls: 

448 return cls 

449 leaf = fqn.rsplit(".", 1)[-1] 

450 if leaf in index.classes: 

451 return leaf 

452 # A remote-code file may name the class without the module-name prefix on 

453 # the FQN leaf; only accept it when the leaf matches a class AND the leaf 

454 # is a plausible modeling class (no lambdas / generics). 

455 if module_name: 

456 candidate = module_name.rsplit(".", 1)[-1] 

457 if candidate in index.classes and _name_matches(candidate, leaf): 

458 return candidate 

459 return None 

460 

461 

462def _name_matches(class_name: str, leaf: str) -> bool: 

463 """Loose match between a module-name class and an FQN leaf. 

464 

465 The FQN leaf is the module attribute (``self_attn`` / ``mlp``); the class 

466 is its type (``Qwen3MoeAttention``). They usually differ, so this only 

467 accepts the trivial identity. The strong mapping is ``boundary_classes``; 

468 this fallback is deliberately conservative. 

469 """ 

470 return class_name == leaf 

471 

472 

473def _verify_class_contract( 

474 group: list[tuple[str, dict[str, Any]]], class_name: str 

475) -> dict[str, Any]: 

476 """Fail fast unless every FQN of ``class_name`` carries the same contract. 

477 

478 One ``class.forward`` is shared by every instance of the class, so all the 

479 class's boundaries must be describable by the SAME frozen entry (modulo the 

480 ``params`` tree, which is instance-specific and drives the native Phase A 

481 parameter sharding, not the forward rewrite). Returns the canonical entry. 

482 """ 

483 canonical: Optional[str] = None 

484 canonical_entry: dict[str, Any] = {} 

485 for fqn, entry in group: 

486 entry = _without_params(entry) 

487 blob = json.dumps(entry, sort_keys=True) 

488 if canonical is None: 

489 canonical = blob 

490 canonical_entry = entry 

491 elif blob != canonical: 

492 raise ValueError( 

493 f"lower_forward_boundaries: class {class_name!r} has FQNs with " 

494 f"differing boundary contracts ({fqn!r}); a shared forward " 

495 "cannot carry two contracts — split the plan or align the " 

496 "per-rule contracts" 

497 ) 

498 return canonical_entry 

499 

500 

501def _without_params(entry: dict[str, Any]) -> dict[str, Any]: 

502 """Drop the instance-specific ``params`` tree from a contract comparison.""" 

503 if "params" not in entry: 

504 return entry 

505 return {k: v for k, v in entry.items() if k != "params"} 

506 

507 

508def _index_injections_by_class( 

509 injections: list[dict[str, Any]], 

510 boundary_classes: Optional[dict[str, str]], 

511 index: SourceIndex, 

512) -> dict[str, dict[str, Any]]: 

513 """Fold a class's injection declarations (inner wrapper / local compute). 

514 

515 Returns ``{class_name: {key: value}}`` — the union of every injection rule 

516 whose ``match`` lands on a FQN of that class. A class with conflicting 

517 values for one key (two FQNs declared different local compute fns) fails 

518 fast, for the same reason as ``_verify_class_contract``. 

519 """ 

520 by_class: dict[str, dict[str, Any]] = {} 

521 for rule in injections: 

522 match = rule.get("match") 

523 if not isinstance(match, str) or not match: 

524 continue 

525 class_name = _class_for_fqn(match, boundary_classes, index, "") 

526 if class_name is None: 

527 continue 

528 payload = {k: rule[k] for k in _INJECTION_KEYS if k in rule} 

529 if not payload: 

530 continue 

531 existing = by_class.setdefault(class_name, {}) 

532 for key, value in payload.items(): 

533 if key in existing and _json_stable(existing[key]) != _json_stable(value): 

534 raise ValueError( 

535 f"lower_forward_boundaries: class {class_name!r} has " 

536 f"conflicting injection values for {key!r}" 

537 ) 

538 existing[key] = value 

539 return by_class 

540 

541 

542def _json_stable(value: Any) -> str: 

543 return json.dumps(value, sort_keys=True, default=str) 

544 

545 

546# --------------------------------------------------------------------------- 

547# toggle / switch forward template (parallel, non-default generation mode) 

548# --------------------------------------------------------------------------- 

549#: Class marker value naming the switch-based (toggle) template. Distinct from 

550#: ``TP_FORM_MARKER`` ("tp_collective") so the runtime install path and preflight 

551#: can tell the two opt-in templates apart. Like the static marker it names the 

552#: template a boundary class actually got. 

553TOGGLE_FORM_MARKER = "toggle" 

554#: The emitted marker statement (method indent, then a blank line). 

555_TOGGLE_FORM_MARKER_LINE = f' {TP_FORM_ATTRIBUTE} = "{TOGGLE_FORM_MARKER}"\n\n' 

556 

557 

558def lower_forward_boundaries_toggle( 

559 source_text: str, 

560 frozen_plan: Any, 

561 *, 

562 boundary_classes: Optional[dict[str, str]] = None, 

563 module_name: str = "", 

564) -> str: 

565 """Rewrite each boundary ``forward`` into per-dimension-switched segments. 

566 

567 A SEPARATE generation mode from :func:`lower_forward_boundaries`: it emits, 

568 for every source-defined boundary class, a forward that keeps the *original 

569 HF forward body verbatim* as ``_forward_impl`` and wraps it with independent 

570 ``if tp_enable`` / ``if cp_enable`` / ``if ep_enable`` segments that slice / 

571 communicate / merge along each dimension the frozen plan shards on. Only ONE 

572 mode is chosen per generation (see ``emit.modeling``), so a generated artifact 

573 never carries a second parallel execution path. 

574 

575 What the switch bodies read is bound by the block-1 runtime path: 

576 

577 - ``mesh_context`` — the codegen mesh context already published as a module 

578 global by ``runtime.publish_codegen_mesh_context``; the toggle booleans and 

579 the submesh handles defer to it. Referencing a module global keeps the 

580 artifact importable without a mesh (the name is only resolved at call time). 

581 - ``self._hp_tp`` — the bare TP operator object bound by 

582 ``hp_install_boundaries``'s collectives step for the ``tp_enable`` body. 

583 

584 The per-axis decomposition is re-derived here from ``frozen_plan`` (the 

585 frozen entry's full per-axis placement dicts), independently of the merged 

586 ``tp_collective`` classification, because the switch template needs one 

587 guarded segment per *axis* rather than one merged verdict over all axes. 

588 """ 

589 edits: list[TextEdit] = [] 

590 for class_name, _, emitted, func, injection, deferred_bias in iter_emitted_forms( 

591 source_text, 

592 frozen_plan, 

593 boundary_classes=boundary_classes, 

594 module_name=module_name, 

595 ): 

596 del injection 

597 if emitted == FORM_IDENTITY: 

598 # Pruned as in the parallel mode: identity boundaries need no rewrite. 

599 continue 

600 if emitted == FORM_REGION: 

601 raise NotImplementedError( 

602 "toggle forward: the region (EP MoE local) form is not emitted " 

603 f"for {_class_for_fqn(class_name, boundary_classes, build_source_index(source_text), module_name)!r}; " 

604 "the toggle template assumes a dense (non-MoE) boundary. Split the " 

605 "plan or keep the default parallel mode for MoE topologies." 

606 ) 

607 if deferred_bias: 

608 # D-22: the toggle template decomposes the exit into per-axis 

609 # guarded segments and has no single boundary-exit point to hang 

610 # the deferred-bias re-add on; emitting nothing would silently 

611 # drop the suppressed bias. Fail fast and point at the mode that 

612 # supports deferral. 

613 raise NotImplementedError( 

614 f"toggle forward: boundary class {class_name!r} declares " 

615 "deferred bias params (D-22); the toggle template has no " 

616 "boundary exit hook to re-add them — use the default " 

617 "parallel mode for plans that defer rowwise biases" 

618 ) 

619 body = _render_toggle_forward( 

620 source_text, 

621 frozen_plan, 

622 class_name, 

623 func, 

624 boundary_classes=boundary_classes, 

625 module_name=module_name, 

626 ) 

627 impl = _build_forward_impl_edit(source_text, func, prefix=_TOGGLE_FORM_MARKER_LINE) 

628 if impl is not None: 

629 edits.append(impl) 

630 edits.extend(replace_function_body(source_text, func, body, indent="")) 

631 if not edits: 

632 return source_text 

633 

634 nested = _validate_no_circular(edits) 

635 if nested: 

636 raise ValueError( 

637 "lower_forward_boundaries_toggle: edits overlap (circular class " 

638 "nesting); cannot safely rewrite: " + ", ".join(nested) 

639 ) 

640 return apply_edits(source_text, edits) 

641 

642 

643def _render_toggle_forward( 

644 source_text: str, 

645 frozen_plan: Any, 

646 class_name: str, 

647 func: FunctionInfo, 

648 *, 

649 boundary_classes: Optional[dict[str, str]] = None, 

650 module_name: str = "", 

651) -> str: 

652 """Render one switch-based forward body for a boundary class. 

653 

654 The body keeps the original forward as ``_forward_impl`` and inserts one 

655 guarded segment per enabled dimension. TP calls the block-1-bound 

656 ``self._hp_tp``; the ``tp_enable`` / ``cp_enable`` / ``ep_enable`` 

657 booleans are bound on the instance by the install path and defer to 

658 ``mesh_context``. 

659 """ 

660 inbound, outbound = _per_axis_toggle_ops( 

661 frozen_plan, class_name, source_text, 

662 boundary_classes=boundary_classes, module_name=module_name, 

663 ) 

664 

665 lines = [ 

666 " # [HYPER TOGGLE] keep the original HF forward as ``_forward_impl``", 

667 " # and insert per-dimension slice / collective / merge segments,", 

668 " # each guarded by its own {{tp|cp|ep}}_enable switch.", 

669 f" # [HYPER INLINE] {_fusion_note_for(class_name)}", 

670 "", 

671 ] 

672 for axis in ("tp", "cp", "ep"): 

673 axis_ops = inbound.get(axis) 

674 if not axis_ops: 

675 continue 

676 lines.append(f" if self.{axis}_enable:") 

677 for name, op in sorted(axis_ops.items()): 

678 expr = _render_tp_call_expr(name, op) 

679 lines.append(f" {name} = {expr}") 

680 lines.append("") 

681 lines.append( 

682 " # [HYPER TOGGLE] original forward body (now ``_forward_impl``)" 

683 ) 

684 lines.append(_render_forward_call(func)) 

685 lines.append("") 

686 for axis in reversed(("ep", "cp", "tp")): 

687 axis_ops = outbound.get(axis) 

688 if not axis_ops: 

689 continue 

690 lines.append(f" if self.{axis}_enable:") 

691 for name, op in sorted(axis_ops.items()): 

692 if op.kind == "all_reduce": 

693 lines.append(" outputs = self._hp_tp.all_reduce(outputs)") 

694 elif op.kind == "all_gather": 

695 lines.append( 

696 f" outputs = self._hp_tp.all_gather(outputs, dim={op.tensor_dim})" 

697 ) 

698 else: 

699 lines.append( 

700 f" outputs = self._hp_tp.reduce_scatter(outputs, dim={op.tensor_dim})" 

701 ) 

702 lines.append("") 

703 lines.append(" return outputs") 

704 return "\n".join(lines) 

705 

706 

707def _per_axis_toggle_ops( 

708 frozen_plan: Any, 

709 class_name: str, 

710 source_text: str, 

711 *, 

712 boundary_classes: Optional[dict[str, str]] = None, 

713 module_name: str = "", 

714) -> tuple[dict[str, dict[str, TransitionOp]], dict[str, dict[str, TransitionOp]]]: 

715 """Decompose one class's boundary into per-axis switch ops for the toggle body. 

716 

717 Returns ``(in, out)``, each ``{axis: {name: TransitionOp}}`` for only the 

718 axes whose placements actually change (op ``kind != identity``). The 

719 decomposition re-classifies each declared name on each of the plan's active 

720 axes independently so a boundary can guard ``tp`` and ``cp`` work in separate 

721 blocks. A class's forward is shared by every instance, so only the first 

722 boundary FQN that resolves to ``class_name`` is decomposed. 

723 """ 

724 param_plan = _plan_field(frozen_plan, "param_plan") or {} 

725 axes = tuple(_plan_field(frozen_plan, "mesh_dim_names") or ()) 

726 if not axes or not param_plan: 

727 return {}, {} 

728 index = build_source_index(source_text) 

729 for fqn, entry in param_plan.items(): 

730 if not is_boundary_entry(entry): 

731 continue 

732 if _class_for_fqn(fqn, boundary_classes, index, module_name) != class_name: 

733 continue 

734 inbound = _per_axis_side(entry, "in_src", "in_dst", axes) 

735 outbound = _per_axis_side(entry, "out_src", "out_dst", axes) 

736 return inbound, outbound 

737 return {}, {} 

738 

739 

740def _per_axis_side( 

741 entry: dict[str, Any], src_field: str, dst_field: str, axes: tuple[str, ...] 

742) -> dict[str, dict[str, TransitionOp]]: 

743 """Per-axis non-identity ops for one side, from the frozen placement strings. 

744 

745 Classifies each declared name on each single active axis purely from the 

746 canonical placement strings (``"R"`` / ``"S(n)"`` / ``"P(sum)"``) that the 

747 freeze wrote — deliberately NOT via ``classify_tp_transition`` (which pulls 

748 ``hyper_parallel.distributed`` / ``accelerate`` and trips over the tests' 

749 ``torch_npu`` stub in isolated runs). So a boundary can carry an independent 

750 ``tp`` and ``cp`` segment. ``ep`` changes (MoE expert routing) are left for 

751 the caller to reject via the region/dense guard. 

752 """ 

753 src_raw = entry.get(src_field) or {} 

754 dst_raw = entry.get(dst_field) or {} 

755 result: dict[str, dict[str, TransitionOp]] = {} 

756 for axis in axes: 

757 axis_ops: dict[str, TransitionOp] = {} 

758 for name in sorted(set(src_raw) | set(dst_raw)): 

759 src_pl = src_raw.get(name) or {} 

760 dst_pl = dst_raw.get(name) or {} 

761 src = src_pl.get(axis, "R") 

762 dst = dst_pl.get(axis, "R") 

763 if not isinstance(src, str) or not isinstance(dst, str) or src == dst: 

764 continue 

765 op = _toggle_placement_op(src, dst) 

766 if op is not None: 

767 axis_ops[name] = op 

768 if axis_ops: 

769 result[axis] = axis_ops 

770 return result 

771 

772 

773def _toggle_placement_op(src: str, dst: str) -> Optional[TransitionOp]: 

774 """One per-axis transition op from a valid frozen placement pair, or None. 

775 

776 Delegates the placement semantics to native 

777 ``classify_tp_transition`` (the single source of truth the boundary-form 

778 classifier already uses), invoked on the single active axis. Frozen 

779 strings are parsed with ``freeze.parse_placement``; any transition the 

780 native path cannot lower (non-``sum`` partial, unsupported pair) or an 

781 identity yields ``None`` (left to the boundary's compiled-plan/DTensor 

782 fallback). 

783 """ 

784 from hyper_parallel.distributed._builder.tp_collective_lowering import ( # pylint: disable=C0415 

785 classify_tp_transition, 

786 ) 

787 

788 result = classify_tp_transition( 

789 [parse_placement(src)], [parse_placement(dst)], ("tp",) 

790 ) 

791 if result is None or result["kind"] == "identity": 

792 return None 

793 return TransitionOp(kind=result["kind"], tensor_dim=result.get("tensor_dim")) 

794 

795 

796def _fusion_note_for(class_name: str) -> str: 

797 """One-line replacement note placed at the top of a switch forward.""" 

798 return f"fused module {class_name}: the switch template keeps the original HF " \ 

799 "forward as its backbone and only adds dimension-guarded segments." 

800 

801 

802# --------------------------------------------------------------------------- 

803# forward body construction 

804# --------------------------------------------------------------------------- 

805 

806def _require_forward(index: SourceIndex, class_name: str) -> FunctionInfo: 

807 """Return the ``forward`` ``FunctionInfo`` for a boundary class or raise.""" 

808 func = index.find_forward(class_name) 

809 if func is None: 

810 raise ValueError( 

811 f"lower_forward_boundaries: class {class_name!r} has no forward " 

812 "method; a boundary forward must exist to be rewritten" 

813 ) 

814 if func.body_start is None or func.body_end is None: 

815 raise ValueError( 

816 f"lower_forward_boundaries: class {class_name!r}.forward has no " 

817 "replaceable body (single-line/pass); cannot sink a boundary" 

818 ) 

819 return func 

820 

821 

822def _build_forward_impl_edit( 

823 source_text: str, 

824 func: FunctionInfo, 

825 prefix: Optional[str] = None, 

826) -> Optional[TextEdit]: 

827 """Build an edit that extracts a method as ``_forward_impl``. 

828 

829 The original method (signature + body) is copied out of the source and 

830 inserted as the private ``_forward_impl`` right before the rewritten 

831 ``def forward``. The copy keeps the original signature and body byte-for- 

832 byte, re-indented to the method's own indent, with the ``def forward`` 

833 renamed to ``def _forward_impl``. ``prefix`` (the static form's class 

834 marker) is prepended to the same insertion — one zero-width edit at 

835 ``def_offset``, so the edit set never carries two edits at one offset. 

836 

837 Returns ``None`` when the forward has no replaceable body (single-line/ 

838 pass) — there is nothing meaningful to extract, and ``forward`` itself was 

839 already rejected by ``_require_forward`` for the same reason, so this is a 

840 defensive fallback. 

841 """ 

842 if func.body_start is None or func.body_end is None: 

843 return None 

844 # The extracted method owns the signature + body; DECORATORS stay attached to 

845 # the original ``forward`` (which is being rewritten) — copying them onto 

846 # ``_forward_impl`` too would duplicate a side effect (e.g. 

847 # ``@torch.no_grad()``) on a method that must be pure. ``func.def_offset`` 

848 # anchors the FIRST line of the method (a decorator when one is present), so 

849 # scan forward from it past decorator lines and blanks to the ``def`` keyword 

850 # — this finds the signature line for both the plain and the decorated case 

851 # without offset arithmetic (a naive backward scan from ``body_start`` finds 

852 # the newline that *ends* the ``def`` line, landing on the first body line). 

853 def_line_start = func.def_offset 

854 for line in source_text[def_line_start:func.body_start].splitlines(keepends=True): 

855 stripped = line.lstrip() 

856 if stripped.startswith("def "): 

857 break 

858 def_line_start += len(line) 

859 span = source_text[def_line_start:func.body_end] 

860 lines = span.splitlines(keepends=True) 

861 if not lines: 

862 return None 

863 # ``_forward_impl`` is a VERBATIM copy of the original method (signature + 

864 # body, byte-for-byte, keeping the method indent on the ``def`` and the body 

865 # nesting) with only the name changed -- so nothing here needs dedenting or 

866 # re-indenting, and the copy can never drift from the original source. 

867 for i, line in enumerate(lines): 

868 stripped = line.lstrip() 

869 if stripped.startswith(f"def {func.name}("): 

870 lines[i] = line.replace(f"def {func.name}(", "def _forward_impl(", 1) 

871 break 

872 impl = "".join(lines) 

873 # The copy ends at ``body_end``, which is clamped to EOF for a final 

874 # statement with no trailing newline — so the extracted method can lose its 

875 # terminating newline and land joined onto the ``def forward`` that follows. 

876 # A method must be separated from the next one by a newline, so pad. 

877 if not impl.endswith("\n"): 

878 impl += "\n" 

879 if prefix: 

880 impl = prefix + impl 

881 return TextEdit(func.def_offset, func.def_offset, impl) 

882 

883 

884def _build_forward_body( 

885 source_text: str, 

886 injection: Optional[dict[str, Any]], 

887 func: FunctionInfo, 

888 form: BoundaryForm, 

889 emitted: str, 

890 deferred_bias: bool = False, 

891) -> str: 

892 """Render the rewritten forward body for one boundary class. 

893 

894 ``emitted`` is :func:`resolve_emitted_form`'s verdict (the classifier's 

895 form after the static gates). The redistribute calls go through the 

896 instance-bound compiled plan (``self._hp_boundary``, installed by 

897 ``hp_install_boundaries``); the static ``tp_collective`` form calls bare 

898 operators on ``self._hp_tp`` instead; ``__hp_compute__`` is bound by 

899 ``hp_bind_compute`` for the local-region shape. Either way the class 

900 stays runtime-importable without a mesh. 

901 

902 ``deferred_bias`` (D-22) appends the exit hook call 

903 ``self._hp_deferred_bias(outputs)`` after the output redistribution — 

904 the runtime binds that hook when it installs the native 

905 suppression/restore pair. 

906 """ 

907 if emitted == FORM_TP_COLLECTIVE: 

908 return _render_tp_collective_forward(source_text, func, form, deferred_bias) 

909 

910 input_redist = _render_input_redistribute(func) 

911 out_redist = _render_output_redistribute(deferred_bias) 

912 

913 local = injection.get("local_compute_fn") if injection else None 

914 if local is not None: 

915 return _render_local_compute(local, input_redist, out_redist, form) 

916 # A generic boundary (including an inner-wrapper one — the CP wrapper is 

917 # installed at runtime by ``hp_apply_inner_wrapper`` and weaves into 

918 # this forward) redistributes through the compiled plan, annotated with 

919 # what each side's transitions lower to. 

920 return _render_plain_redistribute(input_redist, out_redist, form) 

921 

922 

923def _render_input_redistribute(func: FunctionInfo) -> str: 

924 """Render the input-side redistribute call on the compiled boundary plan. 

925 

926 The generated forward re-binds args/kwargs against its own signature, so a 

927 positional ``(args, kwargs)`` pair reaches the compiled plan as the 

928 boundary's input side. Index binding happens once at install time 

929 (``hp_install_boundaries`` → ``_bind_input_indices``); the pair is still 

930 built from the signature's own parameter names (not the plan's) — for a 

931 ``def forward(self, hidden_states)`` that is ``((hidden_states,), {})``. 

932 

933 The full parameter profile is preserved on the re-pass: positional names go 

934 in the args tuple (a trailing ``*args`` expands into it), keyword-only 

935 parameters and a fixed ``**kwargs`` go into the kwargs dict. The boundary 

936 plan rewrites only the args/kwargs it owns and forwards the rest through 

937 unchanged, so ``self._forward_impl(*args, **kwargs)`` receives every 

938 argument class the original call site used — a keyword-only ``required`` 

939 or an HF ``**kwargs`` pass-through no longer drops out. 

940 """ 

941 args = list(func.param_names) 

942 if func.vararg_name: 

943 args.append(f"*{func.vararg_name}") 

944 kwargs_parts = [f'"{name}": {name}' for name in func.kwonly_names] 

945 if func.kwarg_name: 

946 kwargs_parts.append(f"**{func.kwarg_name}") 

947 kwargs_literal = "{" + ", ".join(kwargs_parts) + "}" if kwargs_parts else "{}" 

948 # The args body is spliced into a tuple literal; a single element needs a 

949 # trailing comma (also for a lone ``*vararg``) so it stays a tuple rather 

950 # than parenthesized scalars / a generator expression. 

951 args_body = ", ".join(args) + ("," if len(args) == 1 else "") 

952 return ( 

953 " args, kwargs = self._hp_boundary.redistribute_inputs(\n" 

954 f" (({args_body}), {kwargs_literal}),\n" 

955 " )" 

956 ) 

957