Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / framework_spec.py: 16%
201 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Resolve inline codegen declarations from structure, not family render specs.
17This module replaces per-family ``render_spec.py`` files. Every fact a
18generated artifact needs is derived from the family's *existing runtime
19providers* plus the frozen plan, never from a family-name dispatch table:
21* a YAML ``replace_module._target_`` names the family's own
22 ``@module_replacement`` factory; inspecting that factory's source reveals
23 which generic component it constructs (``RMSNorm`` / ``GroupedExperts`` /
24 ``GQAAttention``) and, for fused attention, the kernel entry it hands over;
25* a YAML ``replace_module._target_`` that names a generic component class
26 directly (under ``hyper_parallel.components.``) resolves through that
27 component's own wrap-style constructor -- the same declaration the native
28 ``_target_replacement_factory`` builds, so no family factory is required;
29* a YAML ``local_compute_fn._target_`` that names a known MoE archetype
30 factory resolves to that archetype's strategy kind / target class / boundary
31 sub-patterns / EP body keys;
32* a CP ``inner_wrapper`` target under the ``models/<family>/adapter``
33 convention resolves to the framework CP attention strategy (imports only).
35The source class name an artifact must model (``Qwen3MoeRMSNorm``) is not
36duplicated here either: it is read from the rule's own ``module_type``, which
37the YAML already carries for its type check. No model-family codegen
38declaration remains.
40Heavy resolution (importing the real generic component to render a fused
41attention class) stays lazy: it only runs when a fused-attention target is
42actually present, so importing this module never pulls torch / torch_npu.
43Resolved specs are cached so repeated lookups stay cheap and stable (callers
44may rely on object identity across calls).
45"""
47from __future__ import annotations
49from dataclasses import asdict
50import importlib
51import inspect
52from typing import Any, Iterable, Mapping
54from hyper_parallel.codegen.inline.ir import ImportPatch, ModuleSnippetPatch
55from hyper_parallel.codegen.inline.recognition import (
56 ReplacementArgument,
57 ReplacementTarget,
58 adapter_path_parts,
59 archetype_for_factory,
60 class_name_of,
61 family_from_target,
62 replacement_target,
63)
64from hyper_parallel.codegen.inline.spec_bundle import ReplacementSpec, StrategySpec
66#: The framework CP attention strategy kind (structure-keyed, not family-named).
67CP_ATTENTION_KIND = "cp_attention"
69_EP_IMPORTS = (
70 ImportPatch("hyper_parallel.distributed.expert_parallel.experts", ("moe_ep_forward",)),
71)
74def _target_path(value: Any) -> str | None:
75 """Return the dotted ``_target_`` of a projected override value."""
76 if value is None:
77 return None
78 if isinstance(value, str):
79 return value
80 if isinstance(value, Mapping):
81 target = value.get("_target_") or value.get("target")
82 return str(target) if target else None
83 return None
86def _family_provider(model_type: str | None, provider: str) -> Any | None:
87 """Return one family adapter provider module, or ``None`` if unregistered."""
88 if not model_type:
89 return None
90 from hyper_parallel.models.registry import get_model_adapter # pylint: disable=C0415
92 adapter = get_model_adapter(model_type)
93 loader = getattr(adapter, provider, None) if adapter is not None else None
94 if loader is None:
95 return None
96 return loader()
99def _family_key(model_type: str | None, target: str | None) -> str | None:
100 """Canonical family key for one lookup.
102 Every spelling of an identity (``qwen3_moe`` / ``Qwen3MoeForCausalLM``)
103 normalizes to the registered ``model_type``, so repeated lookups share one
104 resolved spec -- callers may rely on object identity. A target under the
105 ``models/<family>/adapter`` convention supplies the family when the caller
106 did not name one. An unregistered identity is kept verbatim so resolution
107 still fails closed (``None``) instead of falling back to another family.
108 """
109 if model_type is None:
110 model_type = family_from_target(target)
111 if not model_type:
112 return None
113 from hyper_parallel.models.registry import get_model_adapter # pylint: disable=C0415
115 adapter = get_model_adapter(model_type)
116 return adapter.model_type if adapter is not None else model_type
119def _factory_function(target: str, model_type: str | None):
120 """Resolve a YAML replacement ``_target_`` to the family's real factory.
122 Returns ``None`` when the target is not a factory exported by this family's
123 replacements provider — a target the framework cannot structurally resolve.
124 """
125 module = _family_provider(model_type, "replacements")
126 if module is None:
127 return None
128 prefix = module.__name__ + "."
129 if not target.startswith(prefix):
130 return None
131 name = target[len(prefix):]
132 if "." in name:
133 return None
134 return getattr(module, name, None)
137#: YAML ``replace_module`` targets under this prefix name a generic component
138#: class directly (``hyper_parallel.components.modules.engram.EngramModule``)
139#: -- the declaration the native path builds through
140#: ``_target_replacement_factory`` as ``Cls(module=..., module_fqn=...,
141#: context=...)``.
142_COMPONENTS_PREFIX = "hyper_parallel.components."
145def _direct_component_class(target: str) -> type | None:
146 """Import the generic component class a target names directly.
148 Fail closed: only a ``hyper_parallel.components.*`` path that resolves to
149 a class yields a component. Any other target (a family factory path, a
150 user helper, a typo) is not a component declaration and returns ``None``.
151 """
152 if not target.startswith(_COMPONENTS_PREFIX):
153 return None
154 module_name, _, attr = target.rpartition(".")
155 if not module_name or not attr:
156 return None
157 try:
158 module = importlib.import_module(module_name)
159 except ImportError:
160 return None
161 component = getattr(module, attr, None)
162 return component if isinstance(component, type) else None
165def _component_replacement_target(component: type) -> ReplacementTarget | None:
166 """Read a wrap-style component's own constructor as its replacement target.
168 The structural mirror of ``_target_replacement_factory``: a component
169 whose constructor accepts the wrapped ``module`` declares the same
170 replacement a family factory does -- the matched source instance is
171 handed over (``wraps_source``), and each remaining framework keyword the
172 native factory passes (``module_fqn`` / ``context``) applies the
173 component's own constructor default.
174 """
175 parameters = inspect.signature(component).parameters
176 if "module" not in parameters:
177 return None
178 arguments = tuple(
179 ReplacementArgument(name, None)
180 for name in ("module_fqn", "context")
181 if name in parameters
182 )
183 return ReplacementTarget(
184 component=component.__name__,
185 module=component.__module__,
186 wraps_source=True,
187 arguments=arguments,
188 )
191def _component_class(target: ReplacementTarget) -> type:
192 """Import the component class one factory constructs.
194 Lazy by construction: the component's own module is imported only when a
195 rule that constructs it is actually present (``components.modules``
196 re-exports lazily, so this pulls in one submodule, never the whole set).
197 """
198 module = importlib.import_module(target.module)
199 component = getattr(module, target.component, None)
200 if not isinstance(component, type):
201 raise RuntimeError(
202 f"codegen: {target.module!r} does not export the class "
203 f"{target.component!r} the replacement factory constructs"
204 )
205 return component
208def _render_arguments(
209 target: ReplacementTarget, component: type
210) -> tuple[str, ...] | None:
211 """Render the constructor keywords the generated call passes.
213 Arguments the factory binds to its own parameters have no value at
214 generation time, so the component's own constructor default is emitted --
215 the framework never restates a component's keyword shape. ``None`` means the
216 component does not accept one of the keywords the factory passes, i.e. the
217 factory constructs an interface codegen cannot sink.
218 """
219 parameters = inspect.signature(component).parameters
220 rendered: list[str] = []
221 for argument in target.arguments:
222 if argument.value is not None:
223 rendered.append(f"{argument.name}={argument.value}")
224 continue
225 if argument.name not in parameters:
226 return None
227 rendered.append(f"{argument.name}={parameters[argument.name].default!r}")
228 return tuple(rendered)
231def _interface_function(factory: Any, target: ReplacementTarget) -> Any:
232 """The family kernel entry the factory hands its component."""
233 interface = getattr(factory, "__globals__", {}).get(target.interface)
234 if not callable(interface):
235 raise RuntimeError(
236 f"codegen: replacement factory {getattr(factory, '__name__', factory)!r} "
237 f"hands over {target.interface!r}, which its module does not define"
238 )
239 return interface
242def _build_replacement_spec(
243 factory: Any, target: ReplacementTarget, old_ctor: str
244) -> ReplacementSpec | None:
245 """Build the source-level replacement one factory declares.
247 Every field comes from the factory's own structure: the class it constructs
248 (and the module it imports it from), whether it hands the matched module over
249 as a wrap, the constructor keywords it passes, and the family interface it
250 hands to the component -- which is what makes the component's class rendered
251 (visible in the artifact) instead of merely imported. ``None`` means the
252 construction is not one codegen can sink into the generated ``__init__``.
253 """
254 component = _component_class(target)
255 keyword_args = _render_arguments(target, component)
256 if keyword_args is None:
257 return None
258 # A component that receives a family interface is rendered into the artifact,
259 # so its class is defined there -- importing it as well would be dead weight
260 # the rendered definition shadows.
261 imports: list[ImportPatch] = (
262 [] if target.interface else [ImportPatch(target.module, (target.component,))]
263 )
264 snippets: tuple[ModuleSnippetPatch, ...] = ()
265 if target.interface:
266 from hyper_parallel.codegen.inline.attention import ( # pylint: disable=C0415
267 render_attention_class,
268 )
270 expanded = render_attention_class(
271 component, interface=_interface_function(factory, target)
272 )
273 imports.extend(
274 ImportPatch(module="", names=(), raw=line) for line in expanded.imports
275 )
276 snippets = (ModuleSnippetPatch(expanded.source),)
277 mode = "wrap_source" if target.wraps_source else "name"
278 return ReplacementSpec(
279 old_ctor=old_ctor,
280 new_ctor=target.component,
281 mode=mode,
282 keyword_args=keyword_args,
283 imports=tuple(imports),
284 remove_class=not target.wraps_source,
285 snippets=snippets,
286 replacement_note=(
287 f"{target.component} wraps {old_ctor}; the original class is kept "
288 "as the wrapper input."
289 if target.wraps_source
290 else f"{target.component} replaces {old_ctor}."
291 ),
292 )
295#: Resolved-spec caches. Keyed by everything the resolution reads, so a
296#: repeated lookup returns the same object (callers may compare identity) and
297#: the expensive fused-attention render runs once.
298_REPLACEMENT_CACHE: dict[tuple[str | None, Any, str | None], ReplacementSpec | None] = {}
299_STRATEGY_CACHE: dict[tuple[str | None, str | None, bool], StrategySpec | None] = {}
302def replacement_spec_for(
303 target: str | None,
304 model_type: str | None = None,
305 *,
306 module_type: str | None = None,
307) -> ReplacementSpec | None:
308 """Resolve a ``replace_module`` target to its structural component spec.
310 ``module_type`` is the matched source module's dotted type path from the
311 YAML rule; the source class name (``old_ctor``) is its class name — a
312 property of the matched source module, never a hand-written literal.
313 """
314 if target is None:
315 return None
316 if not adapter_path_parts(target):
317 return _component_replacement_spec(target, model_type, module_type=module_type)
318 model_type = _family_key(model_type, target)
319 factory = _factory_function(target, model_type)
320 if factory is None:
321 return None
322 # The spec is derived from the *factory's* structure, so the factory is
323 # part of the key: a target re-pointed at a different factory (or a family
324 # re-registered with one) must not reuse the previous resolution.
325 key = (model_type, factory, module_type)
326 if key in _REPLACEMENT_CACHE:
327 return _REPLACEMENT_CACHE[key]
328 spec = _resolve_replacement_spec(factory, model_type, module_type, target)
329 _REPLACEMENT_CACHE[key] = spec
330 return spec
333def _resolve_replacement_spec(
334 factory: Any, model_type: str | None, module_type: str | None, target: str
335) -> ReplacementSpec | None:
336 """One construction path for every replacement: read the factory.
338 Whatever component a family's factory constructs, the spec is derived from
339 that factory's own source, so a new family (or a new component inside one)
340 needs no codegen declaration. A factory with no constructor call to sink --
341 or a rule with no ``module_type`` naming the source class it replaces -- has
342 no spec, and the rule is left to the native path.
343 """
344 del model_type, target
345 # The source class the constructor call replaces: without the rule's
346 # ``module_type`` there is no call site to rewrite, and guessing one would
347 # silently target the wrong class.
348 old_ctor = class_name_of(module_type)
349 if not old_ctor:
350 return None
351 replacement = replacement_target(factory)
352 if replacement is None:
353 return None
354 return _build_replacement_spec(factory, replacement, old_ctor)
357def _component_replacement_spec(
358 target: str, model_type: str | None, *, module_type: str | None
359) -> ReplacementSpec | None:
360 """Resolve a target that names a generic component class directly.
362 The component's own constructor is the declaration: the spec is derived
363 from it exactly as the factory channel derives one from the factory's
364 construction call, so a wrap-style component needs no family factory to
365 be sunk into the generated ``__init__``.
366 """
367 component = _direct_component_class(target)
368 if component is None:
369 return None
370 key = (model_type, component, module_type)
371 if key in _REPLACEMENT_CACHE:
372 return _REPLACEMENT_CACHE[key]
373 spec = None
374 old_ctor = class_name_of(module_type)
375 if old_ctor:
376 replacement = _component_replacement_target(component)
377 if replacement is not None:
378 # No interface: a direct component hands over no family kernel,
379 # so the factory argument of _build_replacement_spec is unused.
380 spec = _build_replacement_spec(None, replacement, old_ctor)
381 _REPLACEMENT_CACHE[key] = spec
382 return spec
385def strategy_spec_for(
386 target: str | None,
387 model_type: str | None = None,
388 *,
389 inner_wrapper: bool = False,
390) -> StrategySpec | None:
391 """Resolve a strategy target from structure.
393 A known MoE archetype factory (EP) resolves to its structural strategy; a
394 CP ``inner_wrapper`` target under the adapter convention resolves to the
395 framework CP attention strategy (imports only). Anything else is
396 unresolvable and yields ``None`` — recognition never guesses a strategy.
397 """
398 if target is None:
399 return None
400 model_type = _family_key(model_type, target)
401 key = (target, model_type, inner_wrapper)
402 if key in _STRATEGY_CACHE:
403 return _STRATEGY_CACHE[key]
404 spec = _resolve_strategy_spec(target, model_type, inner_wrapper)
405 _STRATEGY_CACHE[key] = spec
406 return spec
409def _resolve_strategy_spec(
410 target: str, model_type: str | None, inner_wrapper: bool
411) -> StrategySpec | None:
412 archetype = archetype_for_factory(target)
413 if archetype is not None:
414 return _ep_strategy_spec(archetype)
415 if inner_wrapper and adapter_path_parts(target):
416 return StrategySpec(kind=CP_ATTENTION_KIND, imports=())
417 return None
420def _ep_strategy_spec(archetype: Any) -> StrategySpec:
421 """Build the EP strategy spec: the inlined forward plus its imports.
423 The inlined EP forward orchestrates through the instance-attribute channel
424 (``self.ep_enable`` / ``self._hp_ep_group``) like every other generated
425 forward, so it declares no module-level snippet.
426 """
427 from hyper_parallel.codegen.inline.templates import moe_ep_forward_body # pylint: disable=C0415
429 return StrategySpec(
430 kind=archetype.kind,
431 target_class=archetype.target_class,
432 body_template=moe_ep_forward_body(
433 router_kind=archetype.router_kind, shared=archetype.shared
434 ),
435 imports=_EP_IMPORTS,
436 strip_boundary_subpatterns=archetype.strip_boundary_subpatterns,
437 )
440def interface_component_class(model_type: str | None) -> str | None:
441 """The class this family's factories render, because one hands over an interface.
443 Structural: the component is named by the replacement factory itself, so
444 codegen carries no fused-attention class literal. A rendered component
445 keeps the component's own ``forward`` and therefore takes a *compiled
446 boundary* (``hp_install_boundaries``) rather than external state — see
447 :func:`external_state_classes_for`.
448 """
449 module = _family_provider(model_type, "replacements")
450 if module is None:
451 return None
452 for name in dir(module):
453 factory = getattr(module, name)
454 if not callable(factory) or getattr(factory, "_hp_module_replacement", False) is not True:
455 continue
456 target = replacement_target(factory)
457 if target is not None and target.interface:
458 return target.component
459 return None
462def external_state_classes_for(
463 model_type: str | None,
464 injections: Iterable[Mapping[str, Any]] | None = None,
465) -> tuple[str, ...]:
466 """Classes whose ``forward`` an inline strategy body *replaces*.
468 Structure-proven, not declared: every strategy target this plan injects is
469 resolved to its :class:`StrategySpec`, and a spec that carries a
470 ``body_template`` rewrites that class's whole ``forward`` with an inline
471 orchestration body (the MoE EP shell). Such a class owns its collectives,
472 so it must keep its parameter sharding without a compiled boundary — it is
473 skipped by the boundary / compute / inner-wrapper installers and reads its
474 parallel channel from install-time instance attributes instead.
476 A *rendered* component class (the class this family's replacement
477 factories build, e.g. the fused attention) is deliberately NOT one of
478 these: its forward is the component's own plus the emitted boundary form,
479 so the shared compiled-boundary path
480 (``hp_install_boundaries`` / ``hp_apply_inner_wrapper``) owns it.
481 """
482 inlined: set[str] = set()
483 for record in injections or ():
484 spec = strategy_spec_for(
485 _target_path(record.get("local_compute_fn")), model_type
486 )
487 if spec is None or spec.body_template is None or spec.target_class is None:
488 continue
489 inlined.add(spec.target_class)
490 return tuple(sorted(inlined))
493def declaration_summary(
494 model_type: str | None,
495 overrides: Iterable[Mapping[str, Any]] | None = None,
496) -> dict[str, Any]:
497 """Canonical, JSON-safe summary of the declarations an override set resolves to.
499 Adapter declarations live outside the codegen implementation tree, so their
500 resolved content must invalidate cached artifacts. Summarizing exactly the
501 YAML targets this override set will inline keeps that invalidation faithful
502 without a per-family declaration site.
503 """
504 replacement_specs: dict[str, Any] = {}
505 strategy_specs: dict[str, Any] = {}
506 for entry in overrides or ():
507 replace_target = _target_path(entry.get("replace_module"))
508 if replace_target is not None:
509 spec = replacement_spec_for(
510 replace_target, model_type, module_type=entry.get("module_type")
511 )
512 replacement_specs[replace_target] = asdict(spec) if spec is not None else None
513 local_target = _target_path(entry.get("local_compute_fn"))
514 if local_target is not None:
515 spec = strategy_spec_for(local_target, model_type)
516 strategy_specs[local_target] = asdict(spec) if spec is not None else None
517 inner_target = _target_path(entry.get("inner_wrapper"))
518 if inner_target is not None:
519 spec = strategy_spec_for(inner_target, model_type, inner_wrapper=True)
520 strategy_specs[inner_target] = asdict(spec) if spec is not None else None
521 return {"replacement_specs": replacement_specs, "strategy_specs": strategy_specs}
524__all__ = [
525 "CP_ATTENTION_KIND",
526 "declaration_summary",
527 "external_state_classes_for",
528 "interface_component_class",
529 "replacement_spec_for",
530 "strategy_spec_for",
531]