mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
39b086d781
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.
111 lines
4.0 KiB
Python
Executable File
111 lines
4.0 KiB
Python
Executable File
"""Point deepinfer's `tile_kernels` imports at the compat layer.
|
|
|
|
Two rewrites, and only on files the model forward imports:
|
|
|
|
import tile_kernels -> from deepinfer import tk_compat as tile_kernels
|
|
from tile_kernels.<mod> import <names> -> from deepinfer.tk_compat.<mod> import <names>
|
|
|
|
`from tile_kernels.<mod>.<submod> import x` (e.g. `tile_kernels.mhc.post_kernel`) is rewritten to
|
|
`from deepinfer.tk_compat.<mod> import x`, because the compat layer flattens the per-implementation
|
|
submodules into one module per group.
|
|
|
|
Files under `tilelang_ops/asc/` and `tilelang_ops/cuda/` are left alone: those are backend
|
|
implementations that only ever load when TileLang is present, and routing them through the
|
|
reference path would be wrong rather than merely unnecessary.
|
|
|
|
python replace_imports.py --check # report what would change
|
|
python replace_imports.py # rewrite
|
|
python replace_imports.py --revert # restore from .import-orig backups
|
|
"""
|
|
|
|
import argparse
|
|
import re
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[2] / "deepinfer"
|
|
# Backend implementations only load when TileLang is present, so routing them through the
|
|
# reference path would be wrong rather than merely unnecessary. Matched as path segments because
|
|
# the files sit several directories deep.
|
|
SKIP_PARTS = {("tilelang_ops", "asc"), ("tilelang_ops", "cuda"), ("tk_compat",)}
|
|
|
|
BARE = "import tile_kernels\n"
|
|
BARE_NEW = "from deepinfer import tk_compat as tile_kernels\n"
|
|
|
|
# Anchored to the start of the line, i.e. module-level only. A function-local
|
|
# `import tile_kernels` is how the tilelang_ops reference implementations load the *compiled*
|
|
# kernels to compare their Triton ports against; routing those through the reference path would
|
|
# compare a thing to itself.
|
|
BARE_RE = re.compile(r"^import tile_kernels$", re.MULTILINE)
|
|
|
|
# from tile_kernels.a.b import x,y / from tile_kernels.a import x
|
|
FROM_RE = re.compile(r"^from (tile_kernels(?:\.\w+)+) import (.+)$", re.MULTILINE)
|
|
|
|
|
|
def targeted_files() -> list[Path]:
|
|
out = []
|
|
for path in sorted(ROOT.rglob("*.py")):
|
|
parts = path.relative_to(ROOT).parts
|
|
if any(all(seg in parts for seg in group) for group in SKIP_PARTS):
|
|
continue
|
|
text = path.read_text()
|
|
if BARE_RE.search(text) or re.search(r"^from tile_kernels", text, re.MULTILINE):
|
|
out.append(path)
|
|
return out
|
|
|
|
|
|
def rewrite(text: str) -> str:
|
|
text = BARE_RE.sub(BARE_NEW.rstrip("\n"), text)
|
|
|
|
def sub(match: re.Match[str]) -> str:
|
|
module, names = match.group(1), match.group(2)
|
|
parts = module.split(".")[1:] # drop the leading `tile_kernels`
|
|
# `tile_kernels.quant` -> quant; `tile_kernels.mhc.post_kernel` -> mhc
|
|
group = parts[0]
|
|
return f"from deepinfer.tk_compat.{group} import {names}"
|
|
|
|
return FROM_RE.sub(sub, text)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--check", action="store_true")
|
|
parser.add_argument("--revert", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
files = targeted_files()
|
|
if args.revert:
|
|
restored = 0
|
|
for path in files:
|
|
backup = path.with_suffix(".py.import-orig")
|
|
if backup.is_file():
|
|
shutil.move(str(backup), str(path))
|
|
restored += 1
|
|
print(f"restored {restored} file(s)")
|
|
return 0
|
|
|
|
changed = 0
|
|
for path in files:
|
|
original = path.read_text()
|
|
updated = rewrite(original)
|
|
if updated == original:
|
|
continue
|
|
rel = path.relative_to(ROOT.parent).as_posix()
|
|
if args.check:
|
|
print(f"would rewrite {rel}")
|
|
else:
|
|
backup = path.with_suffix(".py.import-orig")
|
|
if not backup.is_file():
|
|
shutil.copy2(path, backup)
|
|
path.write_text(updated)
|
|
print(f"rewrote {rel}")
|
|
changed += 1
|
|
|
|
print(f"\n{changed} file(s) {'would change' if args.check else 'changed'}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|