DeepInfer/scripts/tile_kernels_patch/patch_quant_common.py

"""Patch tile_kernels.quant.common so its host-side half imports without TileLang.

The module mixes two things: configuration dataclasses and shape helpers (pure Python, needed
by `tile_kernels.torch`), and TileLang `@T.macro` kernel fragments (only used by the compiled
kernels). An eager `from tilelang import language as T` plus module-level `@T.macro` decorators
made the whole file unimportable without TileLang, which in turn blocked `tile_kernels.torch`.

Edits, each asserted:

1. The TileLang import becomes optional.
2. The four host-visible annotations that mention `T.dtype` are quoted, so they are not
   evaluated at definition time. Quoting them locally is deliberate: a module-wide
   `from __future__ import annotations` also stringifies the annotations *inside* the
   `@T.macro` bodies, and TileLang's macro parser resolves those itself -- it then fails with
   `NameError: name 'float32' is not defined`.
3. The three members that return TileLang dtype handles resolve `T` on call instead. They
   describe device-side formats, so they still need TileLang -- but only if called.
4. The `@T.macro` block is guarded, since decorators run at import. A guarded block still
   assigns module globals, so `from .common import *` in the kernel modules is unaffected
   wherever TileLang is installed.

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

import argparse
import shutil
import subprocess
import sys
from pathlib import Path

OLD_IMPORT = "from tilelang import language as T\n"
NEW_IMPORT = """try:
    from tilelang import language as T
except ImportError:  # torch reference path only: no TileLang, so no macros below
    T = None
"""

# Annotations evaluated at definition time. Quoted so a missing TileLang does not matter.
ANNOTATIONS = [
    (
        "def get_best_vectorize_size(dtype: T.dtype) -> int:",
        'def get_best_vectorize_size(dtype: "T.dtype") -> int:',
    ),
    (
        "    def dtype(self) -> T.dtype:",
        '    def dtype(self) -> "T.dtype":',
    ),
    (
        "    def sf_dtype(self) -> T.dtype:",
        '    def sf_dtype(self) -> "T.dtype":',
    ),
    (
        "def get_sf_inv_dtype(value_dtype: T.dtype, out_config: CastOutputConfig) -> T.dtype:",
        'def get_sf_inv_dtype(value_dtype: "T.dtype", out_config: CastOutputConfig) -> "T.dtype":',
    ),
]

OLD_DTYPE = """    @property
    def dtype(self) -> "T.dtype":
        return T.dtype(self.torch_dtype) if self.torch_dtype != torch.int8 else T.float4_e2m1fn
"""
NEW_DTYPE = """    @property
    def dtype(self) -> "T.dtype":
        # A TileLang dtype handle describes a device-side format, so it resolves on call rather
        # than at import: the torch reference path never needs it and must not need TileLang.
        from tilelang import language as T

        return T.dtype(self.torch_dtype) if self.torch_dtype != torch.int8 else T.float4_e2m1fn
"""

OLD_SF_DTYPE = """    @property
    def sf_dtype(self) -> "T.dtype":
        return T.dtype(self.sf_torch_dtype)
"""
NEW_SF_DTYPE = """    @property
    def sf_dtype(self) -> "T.dtype":
        from tilelang import language as T

        return T.dtype(self.sf_torch_dtype)
"""

OLD_SF_INV = """    if value_dtype == T.bfloat16 and out_config.round_sf:
"""
NEW_SF_INV = """    from tilelang import language as T

    if value_dtype == T.bfloat16 and out_config.round_sf:
"""

MACRO_START = "@T.macro\n"
MACRO_END = "    tensor[row, col] = sf\n"


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()

    path = args.target / "quant" / "common.py"
    backup = path.with_suffix(".py.orig2")
    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("reverted")
        return 0

    if not backup.is_file():
        shutil.copy2(path, backup)

    text = path.read_text()
    # Idempotent: `T = None` only appears in the patched form, so seeing it means the whole patch
    # already applied and a re-run should be a no-op rather than an assertion error.
    if "    T = None\n" in text:
        print("already patched; nothing to do")
        return 0
    steps = []

    def apply(name: str, old: str, new: str) -> None:
        nonlocal text
        found = text.count(old)
        if found != 1:
            raise SystemExit(f"{name}: expected 1 occurrence, found {found}")
        text = text.replace(old, new, 1)
        steps.append(name)

    apply("optional tilelang import", OLD_IMPORT, NEW_IMPORT)
    for old, new in ANNOTATIONS:
        apply(f"quoted annotation: {old.strip()[:48]}", old, new)
    apply("BaseCastConfig.dtype", OLD_DTYPE, NEW_DTYPE)
    apply("BaseCastConfig.sf_dtype", OLD_SF_DTYPE, NEW_SF_DTYPE)
    apply("get_sf_inv_dtype", OLD_SF_INV, NEW_SF_INV)

    # Guard the macro block: keep the definitions out of import time without changing what the
    # kernel modules see, because a guarded block still sets module globals.
    start = text.index(MACRO_START)
    end = text.index(MACRO_END) + len(MACRO_END)
    block = text[start:end]
    indented = "".join(f"    {line}" if line.strip() else line for line in block.splitlines(True))
    text = text[:start] + "if T is not None:\n" + indented + text[end:]
    steps.append("guarded @T.macro block")

    path.write_text(text)
    compile(path.read_text(), str(path), "exec")
    for step in steps:
        print(f"  applied: {step}")

    # The macros must still exist when TileLang is installed, otherwise the kernel modules break.
    root = str(args.target.parent)
    probe_code = (
        "import sys\n"
        f"sys.path.insert(0, {root!r})\n"
        "from tile_kernels.quant.common import (\n"
        "    BaseCastConfig, get_sf_inv_dtype, get_sf_and_inv,\n"
        "    load_sf, store_sf, transform_sf, get_sf_index)\n"
        "print('macros + host api present')\n"
    )
    probe = subprocess.run(
        [sys.executable, "-c", probe_code],
        capture_output=True,
        text=True,
    )
    print(probe.stdout.strip() or "\n".join(probe.stderr.strip().splitlines()[-3:]))
    return probe.returncode


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