DeepInfer/scripts/tile_kernels_patch/patch_tile_kernels.py

"""Patch tile_kernels so `tile_kernels.torch` imports without TileLang.

Two changes, both minimal and meant to be proposed upstream:

1. `__init__.py` becomes lazy (PEP 562). It used to do
   `from . import (config, engram, mhc, moe, quant, ...)`, which imports the compiled TileLang
   kernels for all of them. Now each submodule is resolved on first access, so
   `import tile_kernels.torch` touches only pure torch.
2. `config.py` stops importing `tilelang.language`. It only used it for three constants
   (`T.max_value` of two float formats), which are now literals.

Both keep behaviour identical wherever TileLang is installed; the script asserts that.

    python patch_tile_kernels.py <site-packages/tile_kernels> [--revert]
"""

import argparse
import importlib
import shutil
import sys
from pathlib import Path

NEW_INIT = (Path(__file__).parent / "tk_patch_init.py").read_text()

OLD_IMPORT = "from tilelang import language as T\n"

OLD_BLOCK = """_DEFAULT_CLAMP_MIN_VALUES = {
    ('e4m3', False): 1e-4,
    ('e4m3', True): float(T.max_value(T.float8_e4m3fn)) * (2**-9),
    ('e2m1', False): float(T.max_value(T.float4_e2m1fn) * (2**-126)),
    ('e2m1', True): float(T.max_value(T.float4_e2m1fn)) * (2**-9),
}
"""

NEW_BLOCK = """# Format maxima as literals (448.0 for e4m3, 6.0 for e2m1). They were previously read from
# tilelang, which meant importing this module required the compiler. The torch reference path
# needs these accessors but not TileLang, and it has to work on a host where TileLang is absent.
_E4M3_MAX = 448.0
_E2M1_MAX = 6.0
_DEFAULT_CLAMP_MIN_VALUES = {
    ('e4m3', False): 1e-4,
    ('e4m3', True): _E4M3_MAX * (2**-9),
    ('e2m1', False): _E2M1_MAX * (2**-126),
    ('e2m1', True): _E2M1_MAX * (2**-9),
}
"""


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("target", type=Path, help="the installed tile_kernels package directory")
    parser.add_argument("--revert", action="store_true")
    args = parser.parse_args()

    init_py = args.target / "__init__.py"
    config_py = args.target / "config.py"
    for path in (init_py, config_py):
        if not path.is_file():
            print(f"missing: {path}", file=sys.stderr)
            return 2

    if args.revert:
        for path in (init_py, config_py):
            backup = path.with_suffix(".py.orig")
            if not backup.is_file():
                print(f"no backup for {path}", file=sys.stderr)
                return 1
            shutil.move(str(backup), str(path))
            print(f"reverted {path.name}")
        return 0

    # Idempotent: re-running this script (after a reinstall, or by accident) must not fail.
    # Each marker only appears in the patched form of its file.
    if "def __getattr__(name: str)" in init_py.read_text() and "_E4M3_MAX" in config_py.read_text():
        print("already patched; nothing to do")
        return 0

    # Record the constants before touching anything, so the replacement can be checked against
    # what the original code computed rather than against what I expect it to compute.
    sys.path.insert(0, str(args.target.parent))
    before = dict(importlib.import_module("tile_kernels.config")._DEFAULT_CLAMP_MIN_VALUES)
    for name in [m for m in sys.modules if m.split(".")[0] == "tile_kernels"]:
        del sys.modules[name]

    for path in (init_py, config_py):
        backup = path.with_suffix(".py.orig")
        if not backup.is_file():
            shutil.copy2(path, backup)

    init_py.write_text(NEW_INIT)

    text = config_py.read_text()
    if OLD_IMPORT not in text:
        print("config.py: import line not found; was it already patched?", file=sys.stderr)
        return 1
    if OLD_BLOCK not in text:
        print("config.py: clamp block not found", file=sys.stderr)
        return 1
    config_py.write_text(text.replace(OLD_IMPORT, "").replace(OLD_BLOCK, NEW_BLOCK))

    for name in [m for m in sys.modules if m.split(".")[0] == "tile_kernels"]:
        del sys.modules[name]
    after = dict(importlib.import_module("tile_kernels.config")._DEFAULT_CLAMP_MIN_VALUES)

    print("constants before / after:")
    ok = True
    for key in sorted(before):
        same = before[key] == after[key]
        ok &= same
        print(f"  {key}: {before[key]!r} -> {after[key]!r} {'OK' if same else 'DIFFERS'}")
    print("\npatched" if ok else "\nCONSTANT MISMATCH")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())