mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
cbb28a8d65
The bring-up plan for a card that cannot build TileLang is to run the model on `tile_kernels.torch` -- pure-torch implementations of nearly every kernel the forward pass calls -- and replace the hot ones with Triton afterwards. That plan was blocked at the import: `import tile_kernels.torch` dragged in the whole compiled stack. Three eager imports were responsible, and each is now lazy or optional: * the top-level `__init__` re-exported every compiled kernel; it resolves them on access * `quant/__init__` did the same for 13 modules, and `import quant.common` runs it first * `quant/common.py` mixed host-side config dataclasses with `@T.macro` fragments; the dataclasses now import without TileLang, the macros are behind a guard, and its three dtype-handle members resolve TileLang when called rather than when imported `config.py` also only needed TileLang for three constants, now literals. One thing that looks like it should work and does not: a module-wide `from __future__ import annotations`. It also stringifies the annotations inside `@T.macro` bodies, which TileLang's macro parser evaluates itself, and it fails there. Only the four host-side annotations are quoted, individually. The patches are idempotent and reversible, and `verify_tile_kernels_patch.py` proves the surface is unchanged where TileLang is installed -- the kernel modules' wildcard imports of `quant.common` still see the macros, `from pkg import sub` still works, config knobs still take effect. The other half of the question, whether the references are numerically trustworthy, is answered by `compare_with_cuda.py`: on the recipe the engine actually uses on Hopper, per_token_cast, cast_back, swiglu_forward_and_per_token_cast and engram_hash are bitwise identical, and apply_rotary differs by one bf16 ULP on one element in 65536. The quantisation chain being bitwise equal is what makes a torch-first bring-up safe rather than approximate; a cast that were merely close would change tokens downstream. Not yet done, and recorded in the doc: MHC and the MoE gate still need their own comparison before anyone assumes they match, and the routing layer that points deepinfer's 25 call sites at these references does not exist yet.
116 lines
4.0 KiB
Python
Executable File
116 lines
4.0 KiB
Python
Executable File
"""Rewrite an eagerly-importing package `__init__.py` into a lazy one (PEP 562).
|
|
|
|
Some TileLang packages re-export every compiled kernel from their `__init__`. That makes the
|
|
package unimportable without TileLang even when the caller only wants a pure-Python submodule
|
|
-- importing `pkg.sub` runs `pkg/__init__.py` first. Resolving the re-exports on first access
|
|
fixes it without changing what callers see.
|
|
|
|
The export map is parsed from the original file rather than hand-written, so a name this script
|
|
does not know about cannot silently disappear.
|
|
|
|
`__all__` is emitted so `from pkg import *` keeps working: wildcard imports call `getattr` for
|
|
every name listed there, which is exactly what triggers the lazy resolution.
|
|
|
|
python lazify_package_init.py <path/to/__init__.py> [--revert]
|
|
"""
|
|
|
|
import argparse
|
|
import ast
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
TEMPLATE = '''"""Lazy re-exports, so importing a submodule does not import every compiled kernel.
|
|
|
|
Generated from the original eager `__init__.py`: each name below is resolved from its module on
|
|
first access. The reason is that `import {pkg}.sub` executes this file first, so an eager
|
|
re-export list made even the pure-Python submodules unavailable on a host without TileLang.
|
|
"""
|
|
|
|
import importlib
|
|
|
|
_EXPORTS = {{
|
|
{exports}
|
|
}}
|
|
|
|
__all__ = sorted(_EXPORTS)
|
|
|
|
|
|
def __getattr__(name: str):
|
|
module = _EXPORTS.get(name)
|
|
if module is None:
|
|
raise AttributeError(f"module {{__name__!r}} has no attribute {{name!r}}")
|
|
value = getattr(importlib.import_module(module, __name__), name)
|
|
globals()[name] = value
|
|
return value
|
|
|
|
|
|
def __dir__() -> list[str]:
|
|
return sorted(set(globals()) | set(_EXPORTS))
|
|
'''
|
|
|
|
|
|
def collect(tree: ast.Module) -> dict[str, str]:
|
|
"""Map each re-exported name to the module it comes from, as an absolute import string."""
|
|
exports: dict[str, str] = {}
|
|
for node in tree.body:
|
|
if isinstance(node, ast.ImportFrom) and node.level == 1 and node.module:
|
|
# `from .mod import a, b`
|
|
for alias in node.names:
|
|
if alias.name == "*":
|
|
raise SystemExit("wildcard import in __init__; cannot infer the export map")
|
|
exports[alias.asname or alias.name] = f".{node.module}"
|
|
elif isinstance(node, ast.ImportFrom) and node.level == 1 and node.module is None:
|
|
# `from . import mod` -- the export *is* the submodule.
|
|
for alias in node.names:
|
|
exports[alias.asname or alias.name] = f".{alias.name}"
|
|
elif isinstance(node, (ast.Import, ast.FunctionDef, ast.ClassDef, ast.Assign, ast.AnnAssign)):
|
|
raise SystemExit(f"unexpected top-level statement in __init__: {type(node).__name__}")
|
|
return exports
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("init_py", type=Path)
|
|
parser.add_argument("--revert", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
path = args.init_py
|
|
backup = path.with_suffix(".py.orig")
|
|
if not path.is_file():
|
|
print(f"missing: {path}", file=sys.stderr)
|
|
return 2
|
|
|
|
if args.revert:
|
|
if not backup.is_file():
|
|
print(f"no backup: {backup}", file=sys.stderr)
|
|
return 1
|
|
shutil.move(str(backup), str(path))
|
|
print(f"reverted {path}")
|
|
return 0
|
|
|
|
source = path.read_text()
|
|
if "_EXPORTS" in source:
|
|
print("already lazy; nothing to do")
|
|
return 0
|
|
exports = collect(ast.parse(source))
|
|
if not exports:
|
|
print("no re-exports found; nothing to do")
|
|
return 0
|
|
|
|
if not backup.is_file():
|
|
shutil.copy2(path, backup)
|
|
|
|
width = max(len(name) for name in exports)
|
|
body = "\n".join(f' "{name}":'.ljust(width + 6) + f' "{module}",' for name, module in sorted(exports.items()))
|
|
path.write_text(TEMPLATE.format(pkg=path.parent.name, exports=body))
|
|
compile(path.read_text(), str(path), "exec")
|
|
print(f"lazified {path}: {len(exports)} names")
|
|
for name, module in sorted(exports.items()):
|
|
print(f" {name:<45} {module}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|