mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 17:26:25 +08:00
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. |
||
|---|---|---|
| .. | ||
| ab_import_side_effects.py | ||
| ab_quant.py | ||