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