Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / framework_spec.py: 57%
164 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 CP ``inner_wrapper`` target under the ``models/<family>/adapter``
30 convention resolves to the framework CP attention strategy — imports
31 and a recognizable inner forward rewrite additionally materializes an
32 in-artifact CP dispatch subclass;
34The source class name an artifact must model (``Qwen3MoeRMSNorm``) is not
35duplicated here either: it is read from the rule's own ``module_type``, which
36the YAML already carries for its type check. No model-family codegen
37declaration remains.
39Heavy resolution (importing the real generic component to render a fused
40attention class) stays lazy: it only runs when a fused-attention target is
41actually present, so importing this module never pulls torch / torch_npu.
42Resolved specs are cached so repeated lookups stay cheap and stable (callers
43may rely on object identity across calls).
44"""
46from __future__ import annotations
48import importlib
49import inspect
50from typing import Any, Callable, cast
52from hyper_parallel.codegen.inline.ir import (
53 ImportPatch,
54 ModuleSnippetPatch,
55 ReplacementSpec,
56 StrategySpec,
57)
58from hyper_parallel.codegen.inline.recognition import (
59 ReplacementArgument,
60 ReplacementTarget,
61 adapter_path_parts,
62 class_name_of,
63 family_from_target,
64 inner_wrapper_target,
65 replacement_target,
66)
68_CP_DISPATCH_SHELL = '''if self.cp_enable:
69 return {helper_name}(self, __hp_forward_args__)
71return self._forward_impl(__hp_forward_args__)
72'''
75def cp_dispatch_shell_body(helper_name: str) -> str:
76 """Render CP dispatch around the extracted original forward."""
77 return _CP_DISPATCH_SHELL.format(helper_name=helper_name)
79#: The framework CP attention strategy kind (structure-keyed, not family-named).
80CP_ATTENTION_KIND = "cp_attention"
82def _family_provider(model_type: str | None, provider: str) -> Any | None:
83 """Return one family adapter provider module, or ``None`` if unregistered."""
84 if not model_type:
85 return None
86 from hyper_parallel.models.registry import get_model_adapter # pylint: disable=C0415
88 adapter = get_model_adapter(model_type)
89 loader = getattr(adapter, provider, None) if adapter is not None else None
90 if loader is None:
91 return None
92 return loader()
95def _family_key(model_type: str | None, target: str | None) -> str | None:
96 """Canonical family key for one lookup.
98 Every spelling of an identity (``qwen3_moe`` / ``Qwen3MoeForCausalLM``)
99 normalizes to the registered ``model_type``, so repeated lookups share one
100 resolved spec -- callers may rely on object identity. A target under the
101 ``models/<family>/adapter`` convention supplies the family when the caller
102 did not name one. An unregistered identity is kept verbatim so resolution
103 still fails closed (``None``) instead of falling back to another family.
104 """
105 if model_type is None:
106 model_type = family_from_target(target)
107 if not model_type:
108 return None
109 from hyper_parallel.models.registry import get_model_adapter # pylint: disable=C0415
111 adapter = get_model_adapter(model_type)
112 return adapter.model_type if adapter is not None else model_type
115def _factory_function(target: str, model_type: str | None):
116 """Resolve a YAML replacement ``_target_`` to the family's real factory.
118 Returns ``None`` when the target is not a factory exported by this family's
119 replacements provider — a target the framework cannot structurally resolve.
120 """
121 module = _family_provider(model_type, "replacements")
122 if module is None:
123 return None
124 prefix = module.__name__ + "."
125 if not target.startswith(prefix):
126 return None
127 name = target[len(prefix):]
128 if "." in name:
129 return None
130 return getattr(module, name, None)
133#: Generic components are valid direct targets without a decorator. Model-owned
134#: classes are accepted only when their native replacement decorator marks them.
135_COMPONENTS_PREFIX = "hyper_parallel.components."
138def _direct_component_class(target: str) -> type | None:
139 """Import a generic or natively decorated replacement class."""
140 module_name, _, attr = target.rpartition(".")
141 if not module_name or not attr:
142 return None
143 try:
144 module = importlib.import_module(module_name)
145 except ImportError:
146 return None
147 component = getattr(module, attr, None)
148 if not isinstance(component, type):
149 return None
150 if (
151 target.startswith(_COMPONENTS_PREFIX)
152 or getattr(component, "_hp_module_replacement", False)
153 ):
154 return component
155 return None
158def _component_replacement_target(component: type) -> ReplacementTarget | None:
159 """Read a wrap-style component's own constructor as its replacement target.
161 The structural mirror of ``_target_replacement_factory``: a component
162 whose constructor accepts the wrapped ``module`` declares the same
163 replacement a family factory does -- the matched source instance is
164 handed over (``wraps_source``), and each remaining framework keyword the
165 native factory passes (``module_fqn`` / ``context``) applies the
166 component's own constructor default.
167 """
168 parameters = inspect.signature(component).parameters
169 if "module" not in parameters:
170 return None
171 arguments = tuple(
172 ReplacementArgument(name, value)
173 for name, value in (("module_fqn", "''"), ("context", "{}"))
174 if name in parameters
175 )
176 return ReplacementTarget(
177 component=component.__name__,
178 module=component.__module__,
179 wraps_source=True,
180 arguments=arguments,
181 )
184def _component_class(target: ReplacementTarget) -> type:
185 """Import the component class one factory constructs.
187 Lazy by construction: the component's own module is imported only when a
188 rule that constructs it is actually present (``components.modules``
189 re-exports lazily, so this pulls in one submodule, never the whole set).
190 """
191 module = importlib.import_module(target.module)
192 component = getattr(module, target.component, None)
193 if not isinstance(component, type):
194 raise RuntimeError(
195 f"codegen: {target.module!r} does not export the class "
196 f"{target.component!r} the replacement factory constructs"
197 )
198 return component
201def _render_arguments(
202 target: ReplacementTarget, component: type
203) -> tuple[str, ...] | None:
204 """Render the constructor keywords the generated call passes.
206 Arguments the factory binds to its own parameters have no value at
207 generation time, so the component's own constructor default is emitted --
208 the framework never restates a component's keyword shape. ``None`` means the
209 component does not accept one of the keywords the factory passes, i.e. the
210 factory constructs an interface codegen cannot sink.
211 """
212 parameters = inspect.signature(component).parameters
213 rendered: list[str] = []
214 for argument in target.arguments:
215 if argument.value is not None:
216 rendered.append(f"{argument.name}={argument.value}")
217 continue
218 if argument.name not in parameters:
219 return None
220 rendered.append(f"{argument.name}={parameters[argument.name].default!r}")
221 return tuple(rendered)
224def _interface_function(factory: Any, target: ReplacementTarget) -> Any:
225 """The family kernel entry the factory hands its component."""
226 interface = getattr(factory, "__globals__", {}).get(target.interface)
227 if not callable(interface):
228 raise RuntimeError(
229 f"codegen: replacement factory {getattr(factory, '__name__', factory)!r} "
230 f"hands over {target.interface!r}, which its module does not define"
231 )
232 return interface
235def _build_replacement_spec(
236 factory: Any,
237 target: ReplacementTarget,
238 old_ctor: str,
239 model_type: str | None,
240) -> ReplacementSpec | None:
241 """Build the source-level replacement one factory declares.
243 Every field comes from the factory's own structure: the class it constructs
244 (and the module it imports it from), whether it hands the matched module over
245 as a wrap, the constructor keywords it passes, and the family interface it
246 hands to the component -- which is what makes the component's class rendered
247 (visible in the artifact) instead of merely imported. ``None`` means the
248 construction is not one codegen can sink into the generated ``__init__``.
249 """
250 component = _component_class(target)
251 keyword_args = _render_arguments(target, component)
252 if keyword_args is None:
253 return None
254 # A component that receives a family interface is rendered into the artifact,
255 # so its class is defined there -- importing it as well would be dead weight
256 # the rendered definition shadows.
257 imports: list[ImportPatch] = (
258 [] if target.interface else [ImportPatch(target.module, (target.component,))]
259 )
260 snippets: tuple[ModuleSnippetPatch, ...] = ()
261 if target.interface:
262 provider = _family_provider(model_type, "codegen")
263 renderer = getattr(provider, "render_replacement_source", None)
264 if not callable(renderer):
265 return None
266 render_source = cast(Callable[..., Any], renderer)
267 expanded = render_source( # pylint: disable=not-callable
268 component, interface=_interface_function(factory, target)
269 )
270 imports.extend(
271 ImportPatch(module="", names=(), raw=line) for line in expanded.imports
272 )
273 snippets = (ModuleSnippetPatch(expanded.source),)
274 mode = "wrap_source" if target.wraps_source else "name"
275 return ReplacementSpec(
276 old_ctor=old_ctor,
277 new_ctor=target.component,
278 mode=mode,
279 keyword_args=keyword_args,
280 imports=tuple(imports),
281 remove_class=not target.wraps_source,
282 snippets=snippets,
283 replacement_note=(
284 f"{target.component} wraps {old_ctor}; the original class is kept "
285 "as the wrapper input."
286 if target.wraps_source
287 else f"{target.component} replaces {old_ctor}."
288 ),
289 )
292#: Resolved-spec caches. Keyed by everything the resolution reads, so a
293#: repeated lookup returns the same object (callers may compare identity) and
294#: the expensive fused-attention render runs once.
295_REPLACEMENT_CACHE: dict[tuple[str | None, Any, str | None], ReplacementSpec | None] = {}
296_STRATEGY_CACHE: dict[tuple[str | None, str | None, bool], StrategySpec | None] = {}
299def replacement_spec_for(
300 target: str | None,
301 model_type: str | None = None,
302 *,
303 module_type: str | None = None,
304) -> ReplacementSpec | None:
305 """Resolve a ``replace_module`` target to its structural component spec.
307 ``module_type`` is the matched source module's dotted type path from the
308 YAML rule; the source class name (``old_ctor``) is its class name — a
309 property of the matched source module, never a hand-written literal.
310 """
311 if target is None:
312 return None
313 if _direct_component_class(target) is not None:
314 return _component_replacement_spec(target, model_type, module_type=module_type)
315 if not adapter_path_parts(target):
316 return None
317 model_type = _family_key(model_type, target)
318 factory = _factory_function(target, model_type)
319 if factory is None:
320 return None
321 # The spec is derived from the *factory's* structure, so the factory is
322 # part of the key: a target re-pointed at a different factory (or a family
323 # re-registered with one) must not reuse the previous resolution.
324 key = (model_type, factory, module_type)
325 if key in _REPLACEMENT_CACHE:
326 return _REPLACEMENT_CACHE[key]
327 spec = _resolve_replacement_spec(factory, model_type, module_type, target)
328 _REPLACEMENT_CACHE[key] = spec
329 return spec
332def _resolve_replacement_spec(
333 factory: Any, model_type: str | None, module_type: str | None, target: str
334) -> ReplacementSpec | None:
335 """One construction path for every replacement: read the factory.
337 Whatever component a family's factory constructs, the spec is derived from
338 that factory's own source, so a new family (or a new component inside one)
339 needs no codegen declaration. A factory with no constructor call to sink --
340 or a rule with no ``module_type`` naming the source class it replaces -- has
341 no spec, and the rule is left to the native path.
342 """
343 del target
344 # The source class the constructor call replaces: without the rule's
345 # ``module_type`` there is no call site to rewrite, and guessing one would
346 # silently target the wrong class.
347 old_ctor = class_name_of(module_type)
348 if not old_ctor:
349 return None
350 replacement = replacement_target(factory)
351 if replacement is None:
352 return None
353 return _build_replacement_spec(factory, replacement, old_ctor, model_type)
356def _component_replacement_spec(
357 target: str, model_type: str | None, *, module_type: str | None
358) -> ReplacementSpec | None:
359 """Resolve a target that names a generic component class directly.
361 The component's own constructor is the declaration: the spec is derived
362 from it exactly as the factory channel derives one from the factory's
363 construction call, so a wrap-style component needs no family factory to
364 be sunk into the generated ``__init__``.
365 """
366 component = _direct_component_class(target)
367 if component is None:
368 return None
369 key = (model_type, component, module_type)
370 if key in _REPLACEMENT_CACHE:
371 return _REPLACEMENT_CACHE[key]
372 spec = None
373 old_ctor = class_name_of(module_type)
374 if old_ctor:
375 replacement = _component_replacement_target(component)
376 if replacement is not None:
377 # No interface: a direct component hands over no family kernel,
378 # so the factory argument of _build_replacement_spec is unused.
379 spec = _build_replacement_spec(None, replacement, old_ctor, model_type)
380 _REPLACEMENT_CACHE[key] = spec
381 return spec
384def strategy_spec_for(
385 target: str | None,
386 model_type: str | None = None,
387 *,
388 inner_wrapper: bool = False,
389) -> StrategySpec | None:
390 """Resolve a strategy target from structure.
392 A CP ``inner_wrapper`` target under the adapter convention resolves to the
393 framework CP attention strategy (imports only). EP mutations come from the
394 builder recorder instead of secondary structural inference. Anything else
395 is unresolvable and yields ``None`` — recognition never guesses a strategy.
396 """
397 if target is None:
398 return None
399 model_type = _family_key(model_type, target)
400 key = (target, model_type, inner_wrapper)
401 if key in _STRATEGY_CACHE:
402 return _STRATEGY_CACHE[key]
403 spec = _resolve_strategy_spec(target, inner_wrapper)
404 _STRATEGY_CACHE[key] = spec
405 return spec
408def _resolve_strategy_spec(target: str, inner_wrapper: bool) -> StrategySpec | None:
409 if inner_wrapper:
410 return _cp_strategy_spec(target)
411 return None
414def _cp_strategy_spec(target: str) -> StrategySpec:
415 """Build the CP strategy spec: an in-place dispatch shell plus imports.
417 A recognizable inner forward rewrite participates in inline CP: the
418 materialization pass derives an
419 in-artifact rewrite of the attention host whose forward falls back to
420 ``_forward_impl`` when CP is disabled and delegates to the helper when
421 active. The helper is the same kernel the runtime CP wrapper delegates
422 to (trainer and codegen share one CP implementation), and the import is
423 derived from the wrapper factory source, so this module never hardcodes a
424 family helper and
425 stays free of torch imports at resolution time. ``subclass_state_attr``
426 carries the family kernel's companion attribute so the runtime binding
427 satisfies the identical contract the native wrapper installs.
429 Families without an inline helper keep the original framework CP
430 strategy (imports only), preserving the runtime wrapper path.
431 """
432 resolved = inner_wrapper_target(target)
433 if resolved is None:
434 return StrategySpec(kind=CP_ATTENTION_KIND, imports=())
435 if resolved.inline_body is not None:
436 return StrategySpec(
437 kind=CP_ATTENTION_KIND,
438 subclass_shell=resolved.inline_body,
439 imports=(),
440 )
441 return StrategySpec(
442 kind=CP_ATTENTION_KIND,
443 subclass_shell=cp_dispatch_shell_body(resolved.helper_name),
444 subclass_state_attr=resolved.state_attr,
445 imports=(ImportPatch(resolved.helper_module, (resolved.helper_name,)),),
446 )
449__all__ = [
450 "CP_ATTENTION_KIND",
451 "cp_dispatch_shell_body",
452 "replacement_spec_for",
453 "strategy_spec_for",
454]