diff --git a/TODO.md b/TODO.md index 30c8518..fe78306 100644 --- a/TODO.md +++ b/TODO.md @@ -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 级算子替换,没有这些定制融合算子。)