docs(h20): record the torch-reference bring-up path and what it measured

Section 6 writes down the three eager imports that blocked the plan, the annotation trap that
looks like it should solve it and does not, and the comparison table showing the quantisation
chain is bitwise equal -- plus, plainly, what is still missing so the next step is not mistaken
for a finished port.
This commit is contained in:
loopyt 2026-10-22 16:08:54 +08:00
parent cbb28a8d65
commit dca7f959cf

58
TODO.md
View File

@ -586,3 +586,61 @@ import 路径上。修完之后 `import deepinfer` 在纯 Triton 环境不再需
**8 卡端到端已跑通,输出与参考镜像逐 token 相同**。剩下的回退就是 P1/P2 那几条遗留,
在两个镜像上一模一样。
- 两边共同的硬骨头只有 **C1(engram gather 不可 capture)**,其余都是收尾工作。
---
## 六、国产卡 bring-up:走 `tile_kernels.torch`(2026-10-22 实测)
目标:在没有 TileLang 的卡上先把模型跑起来,**先用纯 torch 参考实现保正确性**,再按 profile
把热路径逐个换成 Triton。§2.10b 说明了大头在哪(模型前向直接调 `tile_kernels` 的
quant/MoE/MHC/engram,25 处模块级 import / 23 个 API)。
### 6.1 第一步已通:`tile_kernels.torch` 能在无 TileLang 下导入
原本不行,因为三层 eager import 把纯 Python 的参考实现和编译栈绑在一起:
| 层 | 问题 | 处理 |
|---|---|---|
| `tile_kernels/__init__.py` | `from . import (config, quant, …)` 拖起全部编译 kernel | 改成 PEP 562 惰性 `__getattr__` |
| `tile_kernels/quant/__init__.py` | 13 个 kernel 模块 eager re-export;`import quant.common` 会先执行它 | 同样惰性化(AST 解析原文件生成导出表,带 `__all__` 保证 `import *` 仍可用) |
| `tile_kernels/quant/common.py` | `from tilelang import language as T` + 模块级 `@T.macro` | tilelang 变可选;`@T.macro` 段包进 `if T is not None:`;4 个 host 侧注解加引号;3 个返回 dtype 句柄的成员改成调用时再解析 |
| `tile_kernels/config.py` | 只为 3 个常量 `T.max_value(...)` 而 import tilelang | 常量改字面量(448.0 / 6.0),逐位校验一致 |
⚠️ **不能用 `from __future__ import annotations`**:它把 `@T.macro` 体内的注解也变成字符串,
tilelang 的 macro 解析器自己去 eval 时会报 `NameError: name 'float32' is not defined`。
必须只把 host 侧那 4 个注解单独加引号。
脚本:`scripts/enable_torch_kernel_refs.sh`(幂等,`--revert` 可逆)
+ `scripts/tile_kernels_patch/{patch_tile_kernels,patch_quant_common,lazify_package_init}.py`
+ `scripts/verify_tile_kernels_patch.py`(证明有 tilelang 时**接口面完全不变**:子模块属性、
`from pkg import sub`、config 旋钮、kernel 侧的 `from .common import *`、`dir()`、
未知属性仍报 AttributeError → `PATCH_SAFE`)
### 6.2 关键未知项已消除:参考实现数值等价 ✅
`scripts/tile_kernels_patch/compare_with_cuda.py`,用 DeepInfer 在 SM90 上的真实参数
(`fmt="e4m3", num_per_channels=128, round_sf=True, use_packed_ue8m0=False,
use_tma_aligned_col_major_sf=True`):
| op | 结论 |
|---|---|
| `quant.per_token_cast` | **bitwise 一致**(data + sf,4 组 shape 含 m=1/257、n=128/7168) |
| `quant.cast_back` | **rel 0.000e+00** |
| `quant.swiglu_forward_and_per_token_cast` | **bitwise 一致** |
| `engram.engram_hash` | **bitwise 一致**(溢出敏感,必须是 int64 池算术) |
| `transform.apply_rotary` | 65536 元素里 1 个差 1.17/1.33 ULP(bf16 末位舍入),可接受 |
**这是整条桥的前提**:量化链 bitwise 相同意味着「先跑通」不会引入 token 级偏差。
### 6.3 还没做的(下一步)
1. **路由层**:把 deepinfer 里 25 处 `import tile_kernels` 改成走一个 compat 包
(`deepinfer/tk_compat/{quant,moe,transform,mhc,engram,config}.py`),按「有无 tilelang /
`DEEPINFER_KERNELS=ref`」选 CUDA 实现还是 `_ref`。改动最小形式是只换 import 那一行。
2. **MHC 8 个 + MoE gate 未做 A/B**:`mhc_*` 与 `moe_topk_gate` 还没验证(语义更复杂),
做之前不能假设等价。
3. **性能**:纯 torch 的 `per_token_cast` 每层都过,预计是数量级退化。桥的作用是
「先正确」,之后按 profile 换 Triton。
4. **上游**:以上 3 个补丁应该提到 FlagOS,否则每次重装都要重跑脚本。
(`flag_gems` 不是替代品:903 个 ATen 级算子替换,没有这些定制融合算子。)