Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Generate a self-contained attention class from real component source.

17The generated artifact preserves the genuine construction and fused QKV weight

18layout (``linear_qkv`` + ``InterleaveQKV`` checkpoint conversion) by copying the

19real ``GQAAttention`` methods verbatim; ``forward`` is emitted exactly as the

20native class writes it, keeping its ``self.attention_interface(...)`` call. The

21kernel entry function that the real class normally receives as

22``attention_interface`` is inlined from its real source (plan §6) so the

23artifact stays readable down to the single hardware kernel

24``torch_npu.npu_fusion_attention``.

26Neither CP nor TP orchestration is inlined into this rendered class, and the

27class is not an external-state class: it keeps the component's own forward, so

28the emitted boundary form wraps ``_forward_impl`` around it (``emit.parallel``

29lowering + ``hp_install_boundaries``) and the declared CP inner wrapper is

30installed by ``hp_apply_inner_wrapper`` — the same shared mechanisms native

31uses at apply time.

33This generator is platform-agnostic ``core`` code: the module type and kernel

34entry are handed in, never imported here.

35"""

37from __future__ import annotations

39import ast

40import builtins

41import inspect

42import symtable

43import textwrap

44from types import ModuleType

45from typing import Any, Callable, Iterable

47from hyper_parallel.codegen.inline.expansion import ExpandedSource, expand_function

50#: The generated forward is the native GQAAttention.forward verbatim — its

51#: ``self.attention_interface(...)`` call site is kept unmodified, and neither CP

52#: nor TP orchestration is inlined into it. The rendered class keeps the

53#: component's own forward, so the emitted boundary form owns the TP

54#: redistribution (``emit.parallel`` lowering, installed by

55#: ``hp_install_boundaries``) and the CP inner wrapper is installed by

56#: ``hp_apply_inner_wrapper``.

59_FUSION_REPLACEMENT_NOTE = (

60 "# [HYPER INLINE] fused replacement for the original model's per-head\n"

61 "# projections: q_proj/k_proj/v_proj are packed into a single linear_qkv\n"

62 "# weight, whose checkpoint layout is converted via InterleaveQKV.\n"

63)

66def _fusion_replacement_note(class_name: str) -> str:

67 """Annotate that the emitted fused module replaces the original projections.

69 The artifact keeps the real fused class (``linear_qkv`` takes over the

70 weight the original model spread across ``q_proj``/``k_proj``/``v_proj``),

71 so a reader of the generated model file sees a genuine class rather than a

72 black-box substitution, with a comment naming what the fusion swallowed.

73 """

74 return f"# [HYPER INLINE] {class_name}: {_FUSION_REPLACEMENT_NOTE}"

77def _build_forward(module_type: type) -> ast.FunctionDef:

78 """Derive forward from real source verbatim.

80 The body is the real ``forward`` exactly as written (its native

81 ``self.attention_interface(...)`` call site is left completely unmodified).

82 Neither CP nor TP orchestration is inlined here: the emitted boundary form

83 redistributes around this forward, and the CP wrapper swaps

84 ``attention_interface`` at install time — exactly as native's applier does

85 at apply time.

86 """

87 source = textwrap.dedent(inspect.getsource(module_type.forward))

88 tree = ast.parse(source)

89 forward = tree.body[0]

90 if not isinstance(forward, ast.FunctionDef) or forward.name != "forward":

91 raise ValueError(f"codegen attention: {module_type.__name__}.forward is not a plain def")

93 _strip_annotations(forward)

94 ast.fix_missing_locations(forward)

95 return forward

98def _strip_annotations(statement: ast.FunctionDef) -> None:

99 """Drop type hints so the artifact has no import-time resolution burden.

100

101 Only annotations are cleared; ``*args`` and ``**kwargs`` parameters are kept

102 because real forward bodies may reference them.

103 """

104 statement.returns = None

105 for arg in (*statement.args.posonlyargs, *statement.args.args, *statement.args.kwonlyargs):

106 arg.annotation = None

107 if statement.args.vararg is not None:

108 statement.args.vararg.annotation = None

109 if statement.args.kwarg is not None:

110 statement.args.kwarg.annotation = None

111

112

113def _drop_docstring(statement: ast.FunctionDef) -> None:

114 """Drop a leading string-expression docstring from an emitted definition."""

115 if not statement.body:

116 return

117 first = statement.body[0]

118 if (

119 isinstance(first, ast.Expr)

120 and isinstance(first.value, ast.Constant)

121 and isinstance(first.value.value, str)

122 ):

123 del statement.body[0]

124

125

126def _referenced_global_names(source: str) -> set[str]:

127 """Module-scope names ``source`` reads from every scope it defines.

128

129 ``symtable`` resolves the scopes, so parameters, locals, loop variables,

130 comprehension targets and attribute accesses (``x.y`` contributes only

131 ``x``) never show up as global reads.

132 """

133 referenced: set[str] = set()

134 pending = [symtable.symtable(source, "<codegen attention>", "exec")]

135 while pending:

136 table = pending.pop()

137 referenced.update(

138 symbol.get_name()

139 for symbol in table.get_symbols()

140 if symbol.is_global() and symbol.is_referenced()

141 )

142 pending.extend(table.get_children())

143 return referenced

144

145

146def _module_bound_names(source: str) -> set[str]:

147 """Names the module-level statements of ``source`` bind."""

148 table = symtable.symtable(source, "<codegen attention>", "exec")

149 return {

150 symbol.get_name()

151 for symbol in table.get_symbols()

152 if symbol.is_assigned() or symbol.is_imported() or symbol.is_namespace()

153 }

154

155

156def _import_bound_names(lines: Iterable[str]) -> set[str]:

157 """Names a sequence of ``import`` lines binds at module scope."""

158 names: set[str] = set()

159 for line in lines:

160 for alias in ast.parse(line).body[0].names:

161 names.add(alias.asname or alias.name.split(".")[0])

162 return names

163

164

165def _is_module_global(name: str) -> bool:

166 """Whether ``name`` resolves in any module without an import of ours.

167

168 Builtins do, and so do dunders such as ``__name__``. A component definition

169 of the same name wins over this: it is looked up first.

170 """

171 return hasattr(builtins, name) or name.startswith("__")

172

173

174def _owner_definitions(owner: ModuleType) -> dict[str, ast.stmt]:

175 """Module-level definitions of the component's own module, by bound name."""

176 definitions: dict[str, ast.stmt] = {}

177 for node in ast.parse(inspect.getsource(owner)).body:

178 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):

179 definitions[node.name] = node

180 continue

181 if isinstance(node, ast.Assign):

182 targets: list[ast.expr] = list(node.targets)

183 elif isinstance(node, ast.AnnAssign):

184 targets = [node.target]

185 else:

186 continue

187 for target in targets:

188 if isinstance(target, ast.Name):

189 definitions[target.id] = node

190 return definitions

191

192

193def _owner_import_lines(owner: ModuleType) -> dict[str, str]:

194 """Import statement of the component's own module, per name it binds.

195

196 The artifact must import what the component's *module* imports (a re-export

197 such as ``WeightConverter`` only exists under the path that module uses), so

198 each line is copied from that module's own source rather than reconstructed

199 from the value's ``__module__``.

200 """

201 lines: dict[str, str] = {}

202 for node in ast.parse(inspect.getsource(owner)).body:

203 if isinstance(node, ast.Import):

204 for alias in node.names:

205 bound = alias.asname or alias.name.split(".")[0]

206 line = f"import {alias.name}" + (f" as {alias.asname}" if alias.asname else "")

207 lines.setdefault(bound, line)

208 elif isinstance(node, ast.ImportFrom) and not node.level:

209 for alias in node.names:

210 bound = alias.asname or alias.name

211 line = f"from {node.module} import {alias.name}" + (

212 f" as {alias.asname}" if alias.asname else ""

213 )

214 lines.setdefault(bound, line)

215 return lines

216

217

218def _definition_source(node: ast.stmt) -> str:

219 """Copy one module-level definition, dropping hints and docstrings."""

220 if isinstance(node, ast.FunctionDef):

221 _strip_annotations(node)

222 _drop_docstring(node)

223 return ast.unparse(node)

224

225

226def _component_dependencies(

227 module_type: type, code: str, *, provided: Iterable[str]

228) -> tuple[str, tuple[str, ...]]:

229 """Inline and import every name ``code`` reads from the component's module.

230

231 Copying a class body copies its references too, so a module-level helper

232 added to the component after this renderer was written would be emitted as an

233 undefined name. Each global name the copied code reads is therefore resolved

234 against the component's own source module (``inspect.getmodule``, never an

235 import of the component here): a module-level definition is inlined -- its own

236 references first, so a helper referencing another helper completes the

237 closure -- and a name that module imports is imported the same way. A name

238 that resolves to neither fails generation rather than emitting an artifact

239 with a latent ``NameError``.

240

241 Args:

242 module_type: The real component whose source module is the access path.

243 code: Emitted artifact text (the inlined class) to close over.

244 provided: Names other emitted parts already bind in the artifact.

245

246 Returns:

247 The inlined definitions and the import lines the copied code needs.

248 """

249 owner = inspect.getmodule(module_type)

250 if owner is None:

251 raise ValueError(

252 f"codegen attention: cannot resolve the source module of {module_type.__name__!r}"

253 )

254 definitions = _owner_definitions(owner)

255 owner_imports = _owner_import_lines(owner)

256 inlined: list[str] = []

257 imports: list[str] = []

258 resolved = set(provided)

259 unresolved: list[str] = []

260

261 def resolve(name: str) -> None:

262 """Inline or import ``name``, or record it as unresolvable."""

263 if name in resolved:

264 return

265 node = definitions.get(name)

266 if node is not None:

267 resolved.add(name) # before recursing: a helper may reference itself

268 for dependency in sorted(_referenced_global_names(ast.unparse(node))):

269 resolve(dependency)

270 inlined.append(_definition_source(node))

271 return

272 line = owner_imports.get(name)

273 if line is not None:

274 resolved.add(name)

275 if line not in imports:

276 imports.append(line)

277 return

278 if _is_module_global(name):

279 resolved.add(name)

280 return

281 unresolved.append(name)

282

283 for name in sorted(_referenced_global_names(code)):

284 resolve(name)

285 if unresolved:

286 raise ValueError(

287 f"codegen attention: {module_type.__name__} references "

288 f"{sorted(set(unresolved))} which neither {owner.__name__!r} defines nor "

289 "imports; the generated class would carry an undefined name"

290 )

291 return "\n\n\n".join(inlined), tuple(imports)

292

293

294def _class_methods(module_type: type) -> list[ast.stmt]:

295 """Copy every real method except ``forward``, dropping the module decorator."""

296 source = textwrap.dedent(inspect.getsource(module_type))

297 tree = ast.parse(source)

298 for node in tree.body:

299 if isinstance(node, ast.ClassDef) and node.name == module_type.__name__:

300 methods = [

301 statement

302 for statement in node.body

303 if not (

304 isinstance(statement, ast.FunctionDef) and statement.name == "forward"

305 )

306 ]

307 for statement in methods:

308 statement.decorator_list = []

309 if isinstance(statement, ast.FunctionDef):

310 _strip_annotations(statement)

311 return methods

312 raise ValueError(f"codegen attention: no class {module_type.__name__!r} in its source module")

313

314

315def render_attention_class(module_type: type, *, interface: Callable[..., Any]) -> ExpandedSource:

316 """Render a self-contained attention class with a visible TP boundary.

317

318 The copied class methods and the module-level helpers they call are closed

319 over the component's own source module, so the artifact carries everything it

320 references: helpers are inlined and imports derived there, never listed here.

321

322 Args:

323 module_type: The real grouped-query attention component (e.g. ``modules.GQAAttention``).

324 interface: The attention kernel entry function to inline (e.g.

325 ``run_qwen3_moe_flash_attention``); its closure is expanded into the

326 artifact and it must stop at the hardware kernel.

327 """

328 expanded_interface = expand_function(interface)

329 forward = _build_forward(module_type)

330 class_node = ast.ClassDef(

331 name=module_type.__name__,

332 bases=[ast.Name(id="nn.Module", ctx=ast.Load())],

333 keywords=[],

334 body=_class_methods(module_type) + [forward],

335 decorator_list=[],

336 )

337 built_class = _fusion_replacement_note(module_type.__name__) + ast.unparse(class_node)

338 shared_defs = expanded_interface.source

339 helpers, imports = _component_dependencies(

340 module_type,

341 built_class,

342 provided=(_module_bound_names(shared_defs) | _import_bound_names(expanded_interface.imports)),

343 )

344 return ExpandedSource(

345 (*expanded_interface.imports, *imports),

346 "\n\n\n".join(part for part in (helpers, shared_defs, built_class) if part),

347 )

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

143 statements