mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
cbb28a8d65
The bring-up plan for a card that cannot build TileLang is to run the model on `tile_kernels.torch` -- pure-torch implementations of nearly every kernel the forward pass calls -- and replace the hot ones with Triton afterwards. That plan was blocked at the import: `import tile_kernels.torch` dragged in the whole compiled stack. Three eager imports were responsible, and each is now lazy or optional: * the top-level `__init__` re-exported every compiled kernel; it resolves them on access * `quant/__init__` did the same for 13 modules, and `import quant.common` runs it first * `quant/common.py` mixed host-side config dataclasses with `@T.macro` fragments; the dataclasses now import without TileLang, the macros are behind a guard, and its three dtype-handle members resolve TileLang when called rather than when imported `config.py` also only needed TileLang for three constants, now literals. One thing that looks like it should work and does not: a module-wide `from __future__ import annotations`. It also stringifies the annotations inside `@T.macro` bodies, which TileLang's macro parser evaluates itself, and it fails there. Only the four host-side annotations are quoted, individually. The patches are idempotent and reversible, and `verify_tile_kernels_patch.py` proves the surface is unchanged where TileLang is installed -- the kernel modules' wildcard imports of `quant.common` still see the macros, `from pkg import sub` still works, config knobs still take effect. The other half of the question, whether the references are numerically trustworthy, is answered by `compare_with_cuda.py`: on the recipe the engine actually uses on Hopper, per_token_cast, cast_back, swiglu_forward_and_per_token_cast and engram_hash are bitwise identical, and apply_rotary differs by one bf16 ULP on one element in 65536. The quantisation chain being bitwise equal is what makes a torch-first bring-up safe rather than approximate; a cast that were merely close would change tokens downstream. Not yet done, and recorded in the doc: MHC and the MoE gate still need their own comparison before anyone assumes they match, and the routing layer that points deepinfer's 25 call sites at these references does not exist yet.
102 lines
3.6 KiB
Python
Executable File
102 lines
3.6 KiB
Python
Executable File
"""Check that the tile_kernels patch is behaviour-preserving.
|
|
|
|
The lazy `__init__` and the optional TileLang import must change nothing for a host that has
|
|
TileLang. What callers can do with the package -- attribute access, `from x import y`, wildcard
|
|
imports inside the kernel modules, the config knobs -- has to keep working, or the patch has
|
|
broken the normal path in order to enable an unusual one.
|
|
|
|
Run with TileLang available.
|
|
"""
|
|
|
|
import sys
|
|
import traceback
|
|
from collections.abc import Callable
|
|
|
|
FAILED = []
|
|
|
|
|
|
def check(name: str, fn: Callable[[], None]) -> None:
|
|
try:
|
|
fn()
|
|
print(f" ok {name}")
|
|
except Exception as exc:
|
|
print(f" FAIL {name}: {type(exc).__name__}: {str(exc).splitlines()[0][:110]}")
|
|
traceback.print_exc()
|
|
FAILED.append(name)
|
|
|
|
|
|
def main() -> int:
|
|
print("=== same surface as the eager __init__")
|
|
|
|
def submodules_are_attributes() -> None:
|
|
import tile_kernels
|
|
|
|
for name in ("config", "utils", "engram", "mhc", "modeling", "moe", "quant", "rand", "transform", "testing"):
|
|
module = getattr(tile_kernels, name)
|
|
assert module is not None, name
|
|
assert module.__name__ == f"tile_kernels.{name}", module.__name__
|
|
|
|
def submodule_imports_work() -> None:
|
|
# `from pkg import sub` goes through __getattr__ as well.
|
|
from tile_kernels import quant # noqa: F401
|
|
|
|
assert quant.per_token_cast is not None
|
|
|
|
def config_reexports() -> None:
|
|
import tile_kernels
|
|
|
|
for name in ("get_num_sms", "set_num_sms", "get_token_alignment", "set_token_alignment"):
|
|
assert callable(getattr(tile_kernels, name)), name
|
|
# `is_ascend` was never a top-level export, and deepinfer reaches it as
|
|
# `tile_kernels.config.is_ascend`; check the path that is actually used.
|
|
assert callable(tile_kernels.config.is_ascend)
|
|
# The knobs the engine sets must still take effect.
|
|
tile_kernels.set_token_alignment(64)
|
|
assert tile_kernels.get_token_alignment() == 64
|
|
tile_kernels.set_token_alignment(128)
|
|
assert tile_kernels.get_token_alignment() == 128
|
|
|
|
def star_import_of_common() -> None:
|
|
# Kernel modules do `from tile_kernels.quant.common import *`; the guarded macro block
|
|
# still has to publish the macros at module level.
|
|
namespace: dict = {}
|
|
exec("from tile_kernels.quant.common import *", namespace) # noqa: S102
|
|
for name in ("get_sf_and_inv", "get_sf_index", "load_sf", "transform_sf", "store_sf", "BaseCastConfig"):
|
|
assert name in namespace, name
|
|
|
|
def dir_is_populated() -> None:
|
|
import tile_kernels
|
|
|
|
names = dir(tile_kernels)
|
|
assert "quant" in names and "get_num_sms" in names, "dir() lost the eager surface"
|
|
|
|
def unknown_attribute_raises() -> None:
|
|
import tile_kernels
|
|
|
|
try:
|
|
_ = tile_kernels.definitely_not_exported
|
|
except AttributeError:
|
|
return
|
|
raise AssertionError("missing attribute did not raise")
|
|
|
|
for name, fn in [
|
|
("submodules are attributes", submodules_are_attributes),
|
|
("from pkg import sub", submodule_imports_work),
|
|
("config re-exports and knobs", config_reexports),
|
|
("from quant.common import *", star_import_of_common),
|
|
("dir() surface", dir_is_populated),
|
|
("unknown attribute raises", unknown_attribute_raises),
|
|
]:
|
|
check(name, fn)
|
|
|
|
print()
|
|
if FAILED:
|
|
print(f"BROKEN {len(FAILED)}: {FAILED}")
|
|
return 1
|
|
print("PATCH_SAFE")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|