mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 09:16:20 +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.
95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
"""Route `tile_kernels` calls to either the compiled kernels or the pure-torch references.
|
|
|
|
Why this exists: the model forward calls `tile_kernels` directly in 13 modules (quantisation,
|
|
SwiGLU, MoE routing, RoPE, MHC, Engram). Those kernels are TileLang programs, so on an
|
|
accelerator where TileLang cannot be built, nothing runs. `tile_kernels.torch` carries a
|
|
reference implementation of nearly all of them, and this package makes the choice a one-line
|
|
change at each call site:
|
|
|
|
import tile_kernels -> from deepinfer import tk_compat as tile_kernels
|
|
|
|
Call sites keep writing `tile_kernels.quant.per_token_cast(...)`, so no other line changes. The
|
|
signatures exposed here are the *compiled* ones -- including the `out=` buffers and the in-place
|
|
conventions -- because that is what the call sites are written against; each module adapts to the
|
|
reference's shape (different keyword names, return values instead of out parameters) internally.
|
|
|
|
Selection is by `DEEPINFER_KERNELS`, defaulting to `auto`:
|
|
|
|
* ``auto`` -- compiled if TileLang is importable, references otherwise
|
|
* ``cuda`` -- always compiled; fails loudly if TileLang is absent
|
|
* ``ref`` -- always references, even where TileLang exists (for A/B testing)
|
|
|
|
Every operator implemented here has been checked numerically against the compiled version; the
|
|
one known gap is `moe_topk_gate`'s `force_random`, which the reference does not accept -- see
|
|
`deepinfer.tk_compat.moe`.
|
|
"""
|
|
|
|
import importlib
|
|
import os
|
|
|
|
_SUBMODULES = frozenset({"config", "engram", "mhc", "moe", "quant", "torch", "transform", "utils"})
|
|
|
|
# `tile_kernels` re-exports the launch knobs at the top level, and `worker/runtime_config.py`
|
|
# calls `tile_kernels.set_pdl(...)`. Keeping that surface is what lets the call sites read the
|
|
# same after the import swap, so these resolve through `config` on first access.
|
|
_CONFIG_EXPORTS = frozenset(
|
|
{
|
|
"get_amax_clamp_for_quant",
|
|
"get_deterministic_algorithms",
|
|
"get_device_num_sms",
|
|
"get_num_ai_cores",
|
|
"get_num_cube_cores",
|
|
"get_num_sms",
|
|
"get_num_vec_cores",
|
|
"get_pdl",
|
|
"get_token_alignment",
|
|
"is_ascend",
|
|
"set_amax_clamp_for_quant",
|
|
"set_num_sms",
|
|
"set_pdl",
|
|
"set_token_alignment",
|
|
"use_deterministic_algorithms",
|
|
}
|
|
)
|
|
|
|
|
|
def use_references() -> bool:
|
|
"""Whether calls should resolve to the pure-torch references."""
|
|
mode = os.getenv("DEEPINFER_KERNELS", "auto").strip().lower()
|
|
if mode == "ref":
|
|
return True
|
|
if mode == "cuda":
|
|
return False
|
|
if mode != "auto":
|
|
raise ValueError(f"DEEPINFER_KERNELS must be auto, cuda or ref, got {mode!r}")
|
|
return not compiled_available()
|
|
|
|
|
|
def compiled_available() -> bool:
|
|
"""Whether the compiled kernels can actually run here.
|
|
|
|
TileLang is the gate, not `tile_kernels`: the package's own Python is importable anywhere,
|
|
which is what makes the reference path reachable in the first place.
|
|
"""
|
|
try:
|
|
import tilelang # noqa: F401
|
|
except ImportError:
|
|
return False
|
|
return True
|
|
|
|
|
|
def __getattr__(name: str):
|
|
if name in _SUBMODULES:
|
|
module = importlib.import_module(f".{name}", __name__)
|
|
globals()[name] = module
|
|
return module
|
|
if name in _CONFIG_EXPORTS:
|
|
value = getattr(importlib.import_module(".config", __name__), name)
|
|
globals()[name] = value
|
|
return value
|
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
|
|
|
|
def __dir__() -> list[str]:
|
|
return sorted(set(globals()) | _SUBMODULES | _CONFIG_EXPORTS)
|