mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 16:56: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.
70 lines
3.0 KiB
Bash
Executable File
70 lines
3.0 KiB
Bash
Executable File
#!/bin/bash
|
|
# Make `tile_kernels.torch` importable on a host without TileLang.
|
|
#
|
|
# Why this is needed: `tile_kernels.torch` holds pure-torch reference implementations of nearly
|
|
# every kernel the model forward calls (quantisation, SwiGLU, MoE routing, RoPE, MHC, Engram),
|
|
# and none of them reference TileLang. But the package's `__init__` re-exported every compiled
|
|
# kernel eagerly, so `import tile_kernels.torch` dragged in the whole TileLang stack and failed
|
|
# where that stack cannot be built -- which is exactly the case on a non-NVIDIA card.
|
|
#
|
|
# The three edits, applied to the installed package:
|
|
# 1. top-level `__init__` resolves its re-exports on first access (PEP 562)
|
|
# 2. `quant/__init__` does the same, because importing `quant.common` runs it first
|
|
# 3. `quant/common.py` stops requiring TileLang for its host-side half (config dataclasses,
|
|
# shape helpers) while keeping the `@T.macro` fragments for the kernels
|
|
#
|
|
# Nothing changes on a host that has TileLang: `scripts/verify_tile_kernels_patch.py` checks
|
|
# that the surface is identical.
|
|
#
|
|
# Usage, inside the target environment:
|
|
# bash scripts/enable_torch_kernel_refs.sh # patch in place
|
|
# bash scripts/enable_torch_kernel_refs.sh --revert # restore the originals
|
|
# PYTHON=/path/to/python bash scripts/enable_torch_kernel_refs.sh
|
|
#
|
|
# This patches a third-party package, which is why it is a script rather than a checkout-local
|
|
# shim: the intent is for these changes to go upstream to FlagOS. Until then, re-running this
|
|
# after a reinstall is the whole maintenance burden, and --revert makes it reversible.
|
|
|
|
set -euo pipefail
|
|
|
|
PYTHON="${PYTHON:-python3}"
|
|
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
|
|
|
TARGET="$("$PYTHON" - <<'PY'
|
|
import pathlib
|
|
|
|
import tile_kernels
|
|
|
|
print(pathlib.Path(tile_kernels.__file__).parent)
|
|
PY
|
|
)"
|
|
|
|
say() { printf '\n\033[1m== %s\033[0m\n' "$*"; }
|
|
|
|
say "tile_kernels at $TARGET"
|
|
|
|
if [ "${1:-}" = "--revert" ]; then
|
|
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/lazify_package_init.py" \
|
|
"$TARGET/quant/__init__.py" --revert
|
|
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_quant_common.py" "$TARGET" --revert
|
|
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_tile_kernels.py" "$TARGET" --revert
|
|
say "reverted"
|
|
exit 0
|
|
fi
|
|
|
|
say "1/3 lazy top-level __init__"
|
|
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_tile_kernels.py" "$TARGET"
|
|
|
|
say "2/3 lazy quant/__init__"
|
|
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/lazify_package_init.py" "$TARGET/quant/__init__.py"
|
|
|
|
say "3/3 quant/common.py without a hard TileLang import"
|
|
"$PYTHON" "$REPO_ROOT/scripts/tile_kernels_patch/patch_quant_common.py" "$TARGET"
|
|
|
|
say "verify: unchanged where TileLang is installed"
|
|
"$PYTHON" "$REPO_ROOT/scripts/verify_tile_kernels_patch.py"
|
|
|
|
say "done -- the reference path is now importable without TileLang"
|
|
echo " check it with: DEEPINFER_SIMULATE_NO_TILELANG=1 python -c 'import tile_kernels.torch'"
|
|
echo " (that env var needs the sitecustomize from scripts/simulate_no_tilelang/)"
|