mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 16:56: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.
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
"""`tile_kernels.utils` -- small helpers the model imports directly.
|
|
|
|
Pure Python on both paths: `align` / `ceil_div` are arithmetic, `dtype_to_str` is a lookup. There
|
|
is nothing to route, only to forward, so that a call site does not have to know whether the
|
|
compiled package is present.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
try:
|
|
from tile_kernels.utils import align, ceil_div, dtype_to_str, str_to_dtype
|
|
|
|
_HAVE_UTILS = True
|
|
except ImportError: # pragma: no cover - only without tile_kernels installed
|
|
_HAVE_UTILS = False
|
|
|
|
def align(value: int, alignment: int) -> int:
|
|
return (value + alignment - 1) // alignment * alignment
|
|
|
|
def ceil_div(a: int, b: int) -> int:
|
|
return -(-a // b)
|
|
|
|
# The compiled helper returns short names ('bf16', 'fp32'), not torch's ('bfloat16'). Callers
|
|
# build kernel names and dtype strings from them, so the fallback has to match exactly.
|
|
_SHORT = {
|
|
"torch.bfloat16": "bf16",
|
|
"torch.float16": "fp16",
|
|
"torch.float32": "fp32",
|
|
"torch.float8_e4m3fn": "e4m3",
|
|
"torch.float8_e5m2": "e5m2",
|
|
"torch.int32": "int32",
|
|
"torch.int64": "int64",
|
|
"torch.int8": "int8",
|
|
"torch.uint8": "uint8",
|
|
}
|
|
|
|
def dtype_to_str(dtype: Any) -> str:
|
|
key = str(dtype)
|
|
if key in _SHORT:
|
|
return _SHORT[key]
|
|
return key.removeprefix("torch.")
|
|
|
|
def str_to_dtype(name: str) -> Any:
|
|
import torch
|
|
|
|
long_names = {short: long.removeprefix("torch.") for long, short in _SHORT.items()}
|
|
if name in long_names:
|
|
return getattr(torch, long_names[name])
|
|
return getattr(torch, name)
|
|
|
|
|
|
__all__ = ["align", "ceil_div", "dtype_to_str", "str_to_dtype"]
|