loopyt 39b086d781 feat(h20): route tile_kernels calls through a compat layer, so the model runs without TileLang
The references were importable and numerically validated, but nothing pointed the model at them:
the forward pass still called the compiled kernels, which raise ModuleNotFoundError on a host
without TileLang. This adds the routing layer and swaps the call sites.

`deepinfer/tk_compat/` exposes the *compiled* signatures -- out-buffers, in-place conventions,
argument order -- because that is what the 25 call sites are written against, and adapts to the
reference's shape internally. Selection is `DEEPINFER_KERNELS=auto|cuda|ref`; `auto` follows
whether TileLang imports.

The adapters are not pass-throughs. Each difference was found by running both backends:

* `per_token_cast` takes `num_per_channels=n`, the reference takes `block_size=(1, n)`.
* `apply_rotary` takes a concatenated cos|sin cache and rotates in place returning None; the
  reference takes cos and sin separately and returns a new tensor. The adapter splits, calls, and
  copies back, because the call site relies on the in-place contract.
* `engram_gate` takes the pre-multiplied weight and `(eps, clamp_value)`; the reference needs the
  two weights separately and `(clamp_value, eps)`. Since a product cannot be un-multiplied,
  `fused_weight` returns the pair on the reference path and the gate unpacks it.
* `mhc_pre_big_fuse_fwd` has no single reference: its five steps are composed from the reference's
  own components.
* `mhc_reduce_partials_and_rmsnorm_fwd` likewise; rebuilt from the reference's formula.

`moe_topk_gate(force_random=...)` raises rather than silently ignoring the flag. The reference has
no equivalent, and dropping it would route tokens differently while looking healthy.

Verification, all on the real call signatures: 18 cases compared between `cuda` and `ref` modes,
`ALL_MATCH`, 12 bitwise identical including every quantised byte and every routing index; plus a
smoke run that executes all 23 APIs with TileLang blocked, `ALL_CALLS_OK`.

End to end still has not completed, and the reason is not this change: six GPUs on the host are
held by other containers (82 GiB of 139 each), so the 8-way EP run OOMs in allocate_userbuffer and
the surviving ranks wait at a barrier. Recorded in the doc along with the command to rerun it.
2026-10-22 18:10:14 +08:00

75 lines
2.7 KiB
Python
Executable File

"""List every `tile_kernels.<module>.<name>` deepinfer uses, so the compat layer can be complete.
Worth doing as a scan rather than discovering gaps one traceback at a time: the e2e run is the
only thing that would catch a missing attribute, and it costs a minute per attempt.
"""
import ast
import pathlib
import sys
ROOT = pathlib.Path(__file__).resolve().parents[2] / "deepinfer"
SKIP = ("tilelang_ops/cuda", "tilelang_ops/asc", "tk_compat")
# Names reached through a module-level `from deepinfer import tk_compat as tile_kernels`.
used: dict[str, dict[str, set[str]]] = {}
# Names imported directly from a submodule: `from deepinfer.tk_compat.quant import per_token_cast`.
imported: dict[str, set[str]] = {}
for path in sorted(ROOT.rglob("*.py")):
rel = path.as_posix()
if any(s in rel for s in SKIP):
continue
try:
tree = ast.parse(path.read_text(encoding="utf-8"))
except (SyntaxError, UnicodeDecodeError):
continue
for node in ast.walk(tree):
if isinstance(node, ast.ImportFrom) and node.module and node.module.startswith("deepinfer.tk_compat"):
parts = node.module.split(".")
if len(parts) >= 3:
sub = parts[2]
for a in node.names:
imported.setdefault(sub, set()).add(a.name)
if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Attribute):
inner = node.value
if isinstance(inner.value, ast.Name) and inner.value.id == "tile_kernels":
used.setdefault(inner.attr, {}).setdefault(node.attr, set()).add(
rel.replace(str(ROOT.parent) + "/", "")
)
print("=== 通过 `tile_kernels.<mod>.<name>` 访问的")
for sub in sorted(used):
print(f"\n{ sub }:")
for name in sorted(used[sub]):
files = ", ".join(sorted(used[sub][name]))
print(f" {name:<38} {files}")
print("\n=== `from deepinfer.tk_compat.<mod> import ...` 的")
for sub in sorted(imported):
print(f"\n{sub}: {', '.join(sorted(imported[sub]))}")
# Cross-check against what the compat modules actually export.
sys.path.insert(0, str(ROOT.parent))
import importlib
print("\n=== 对照 compat 实际导出")
missing = []
for sub, names in sorted(used.items()):
try:
module = importlib.import_module(f"deepinfer.tk_compat.{sub}")
except Exception as exc: # noqa: BLE001
print(f" {sub}: 无法导入 ({exc})")
continue
absent = [n for n in names if not hasattr(module, n)]
if absent:
missing.append((sub, absent))
print(f" {sub}: MISSING {absent}")
else:
print(f" {sub}: ok")
if missing:
print(f"\n缺口 {sum(len(a) for _, a in missing)} 个")
sys.exit(1)
print("\n无缺口")