diff --git a/TODO.md b/TODO.md index b0ba604..30c8518 100644 --- a/TODO.md +++ b/TODO.md @@ -218,7 +218,52 @@ SM90 强制 **per-128 + fp32 scale**,checkpoint 是 **per-32 + e8m0**。 「技术栈换成 Triton」的隐含目标是「只装 Triton 也能跑」,但当前做不到,因为 **TileLang 概念漏到了接口层**。已在第四章验证并修掉第一批: `ops.py` 加 `DEEPINFER_OPS_STRICT`、`tile_kernels` / `tilelang` 的 import 移进参考实现与单测函数。 -剩下的问题见第四章 §4.4。 + +### 2.10b ⚠️ 「没有 TileLang」的真实阻塞面比 16 个 kernel 大得多(2026-10-22 实测) + +做了一次「假装这两个包不存在」的实测:用 `sitecustomize.py` 在 meta_path 上拦 +`tilelang` / `tile_kernels` / `tvm_ffi`(必须走 sitecustomize,因为 8 个 DP rank 是 spawn +出来的,父进程的拦截不继承),然后 `import deepinfer`。 + +**结果:import 就失败了,根本走不到那 16 个 kernel。** + +``` +hugging_face/weights.py:38 → model_executor/layers/linear.py:4 → import tile_kernels +ModuleNotFoundError: No module named 'tile_kernels' +``` + +**为什么**:`tilelang_ops` 那层抽象只包住了「kernel 库」里的 44 个 kernel,但模型前向 +**直接调** `tile_kernels`(它是 FlagOS 的另一个包,`tile_kernels/__init__.py:1` 就是 +`import tilelang`,140 个 py 里 166 处引用 → **留着 tile_kernels 等于留着 tilelang**)。 + +AST 扫描结果(排除 `tilelang_ops/cuda/` 与 `asc/`): + +| 项 | 数量 | 分布 | +|---|---|---| +| **模块级 `import tile_kernels`** | **25 处 / 13 个文件** | `linear` `rotary_embedding` `deepseek` `v4_modules/{attn,engram,engram_metadata,hc,mhc_ops,mhc_modeling,mlp}` `vision_encoder/vit` `parallel_utils/userbuffer` | +| **直接调用的 API** | **23 个** | 见下 | + +按移植成本分三类: + +| 类别 | API | 说明 | +|---|---|---| +| **琐碎** | `set_num_sms` `set_pdl` `set_token_alignment` `is_ascend` `utils.dtype_to_str` `use_deterministic_algorithms` | 配置开关与字符串工具,几行代码替代 | +| **热路径真 kernel** | `quant.per_token_cast`(FP8 per-token 量化,linear/mlp/attn/hc 四处都用)
`quant.swiglu_forward_and_per_token_cast`(融合 SwiGLU+量化)
`moe.moe_topk_gate`
`transform.apply_rotary`
`quant.swiglu_forward`
`quant.batched_transpose_weight_sf` | 每 token 每层都过,必须真写 Triton,纯 torch 版本会慢到不可用 | +| **DeepSeek 专有重 kernel** | 9 个 `mhc.*`(expand / head_compute_mix / norm_fn / post / pre_apply_mix / pre_big_fuse / pre_split_mixes / sinkhorn)+ `engram.engram_hash` | MHC 注意力与 Engram,语义复杂,工作量最大 | + +**唯一现成的退路**:`tile_kernels/torch/` 是一整套**纯 torch 参考实现**(`cast` `engram` +`mhc` `moe` `norm` `reduce_fused` `rope` `swiglu` `topk`,**0 处 tilelang**),导出 +`engram_hash_ref`、8 个 `mhc_*_ref`、`moe_topk_gate_forward`、`swiglu_forward`、 +`rotary_embedding_ref`、`cast`/`cast_back`、`batched_transpose_weight_sf`、 +`norm_forward_ref` 等,**几乎覆盖上表所有重 API**。但 `tile_kernels/__init__.py` 第一行 +`import tilelang` 会挡住 `import tile_kernels.torch`,需要把 `__init__` 改成 lazy +(或把 `torch/` 抽成独立包)。 + +→ 这条能给国产卡一个「**正确但慢**」的 bring-up 桥:先用 `*_ref` 跑通数值,再按 profile +把热路径逐个换成 Triton。`flag_gems` 不能直接替代(903 个 ATen 级算子,没有这些定制融合算子)。 + +**结论:不要再把「国产卡上不能有 tilelang」说成「只差 16 个 kernel」。** +16 个 kernel 是小头,模型前向对 `tile_kernels` 的直接依赖是大头。 ### 2.2 MTP 4 个 kernel(✅ 2026-10-22 全部完成,见下表) ### 2.3 采样器 verify 变体