mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 09:16:20 +08:00
32d3d2c6bf
Opens `h20-sm90-flaggems` to replace `tile_kernels`'s torch reference with FlagGems' Triton kernels. The point is not "runs without TileLang" -- that landed already. It is speed: the compat layer is plain torch today, so every activation in the hot path is going through eager PyTorch. Wrote the plan down because the tempting framing is wrong. `tile_kernels` and `flag_gems` are not substitutes; they overlap in three layers and only the first is a clean fit: * generic ops (quant, rotary, rms_norm, topk) -- FlagGems' home turf, 251 ops * DeepSeek-specific fused ops (mhc_*, engram_*, the five-way swiglu+quant) -- partial overlap, and almost none of the signatures match 1:1 * the engine kernel backend (tilelang_ops) -- already done, unrelated De-risked the highest-value one first, since per-token quantisation runs on every activation. `ab_quant.py`, 5 shapes: `per_token_group_quant_fp8(scale_ue8m0=True, column_major_scales=True)` against `per_token_cast(fmt="e4m3", 128, round_sf=True, use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True)` -- **quantised values bitwise identical, scales and dequantised values exactly zero difference**. The two `round_sf` spellings are the same function: DeepInfer does `(int + 0x7FFFFF) & 0x7F800000`, which carries into the exponent field, i.e. ceiling to a power of two; FlagGems does `exp2(ceil(log2(s)))`. Same function, so the test demands exactness rather than tolerance and reports any difference. `ab_import_side_effects.py` pins the mechanism that makes this safe: import is inert, `only_enable(include=[...])` touches exactly what it names. Recorded also that the package root pulls Triton TLE -> libflagcx, so selective imports or vendoring stay the default posture rather than reaching for `enable()`. Two questions the plan puts ahead of the work, because without answers the payoff is unknowable: nothing on this line has ever been measured for speed, and `batched_transpose_weight_sf` is still the only path operator never A/B'd -- it rewrites every weight's scale factor and is the leading suspect for the token divergence. Fixing that could be worth more than any op swap here. Also names what is deliberately off the table: `enable()` with no arguments, swapping tilelang_ops again, and churning broadcast-only micro-ops that were never the bottleneck.
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
"""What does importing `flag_gems` actually do to the process?
|
|
|
|
This exists because the answer was assumed wrongly once. The earlier belief was that importing the
|
|
package registers ATen implementations process-wide, which would make any dependency on it a
|
|
numerical risk for the whole engine. That is not what happens: import creates the `aten_lib` handle
|
|
and pulls the op modules in, but registration only occurs when `enable()` or `only_enable()` is
|
|
called.
|
|
|
|
The distinction decides the architecture, so it is measured rather than read:
|
|
|
|
* `import flag_gems` -> ATen dispatch tables must be unchanged
|
|
* `only_enable(include=[...])` -> exactly the named ops change, nothing else
|
|
* `enable()` -> everything changes (not run here, and not wanted)
|
|
|
|
Also records the import's real cost, which is the reason to stay selective: the package root
|
|
star-imports every op module and loads Triton TLE, which dlopen's `libflagcx`. That needs
|
|
`LD_LIBRARY_PATH` to find `libcudart`, and it is a native dependency the engine does not otherwise
|
|
have.
|
|
|
|
Run: PYTHONPATH=. python scripts/flaggems_port/ab_import_side_effects.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import sys
|
|
|
|
import torch
|
|
|
|
# A spread of ops: FlagGems provides all of these, so any registration would be visible here.
|
|
OPS = ["aten::rms_norm", "aten::silu", "aten::softmax", "aten::topk", "aten::lerp"]
|
|
|
|
|
|
def snapshot() -> dict[str, str]:
|
|
"""Hash each op's dispatch table. Content, not length -- registration rewrites entries."""
|
|
out = {}
|
|
for op in OPS:
|
|
try:
|
|
table = torch._C._dispatch_dump_table(op)
|
|
except Exception as exc: # noqa: BLE001 - op may not exist on this build
|
|
table = f"unavailable: {exc}"
|
|
out[op] = hashlib.md5(table.encode()).hexdigest()[:12]
|
|
return out
|
|
|
|
|
|
def report(tag: str, before: dict[str, str], after: dict[str, str]) -> bool:
|
|
changed = [op for op in OPS if before[op] != after[op]]
|
|
print(f"\n[{tag}]")
|
|
for op in OPS:
|
|
mark = "changed" if before[op] != after[op] else "same"
|
|
print(f" {op:16} {before[op]} -> {after[op]} {mark}")
|
|
return bool(changed)
|
|
|
|
|
|
def main() -> int:
|
|
failures = []
|
|
|
|
before = snapshot()
|
|
import flag_gems
|
|
|
|
print(f"flag_gems {flag_gems.__version__} on {flag_gems.vendor_name}/{flag_gems.device}")
|
|
|
|
# 1. Import alone must not touch anything.
|
|
if report("after import flag_gems (expect: nothing changed)", before, snapshot()):
|
|
failures.append("import registered ATen ops")
|
|
|
|
# 2. only_enable must touch exactly what it names.
|
|
before_narrow = snapshot()
|
|
flag_gems.only_enable(include=["rms_norm"])
|
|
after_narrow = snapshot()
|
|
changed = [op for op in OPS if before_narrow[op] != after_narrow[op]]
|
|
print("\n[only_enable(include=['rms_norm'])]")
|
|
print(f" changed: {changed or 'none'}")
|
|
if changed != ["aten::rms_norm"]:
|
|
failures.append(f"only_enable changed {changed}, expected only aten::rms_norm")
|
|
|
|
print()
|
|
if failures:
|
|
print(f"UNEXPECTED: {failures}")
|
|
return 1
|
|
print("AS_EXPECTED: import is inert; only_enable is scoped to the ops it names")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|