mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
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.
75 lines
2.7 KiB
Python
Executable File
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无缺口")
|