Summary
PTODSL对接RMSNorm的需求整理
Motivation / use case
场景
RMSNorm 是个 SimtVF kernel:每核串行处理多个 token,每 token 在 SimtVF 内做 float4 向量访存、行内 sum-of-squares、跨 128 lane all-reduce、rstd = rsqrt(Σx²/d + eps),再 y = x * rstd * w。其中权重 W 对一个核内的 token 循环是不变量,但当前每个 token 进 SimtVF 都要把 w 从 UB 重读进寄存器一次,每核重复 batch/N_CORES 次。
目标两层:① PTODSL 能表达这个 SimtVF kernel;② 支持跨 SimtVF 寄存器驻留优化——w 一次读进寄存器、后续 token 复用,把 UB→VRF 从每核 64 次降到 1 次。
设计文档:
A. 对 PTODSL 前端的需求
-
SimtVF 编程接口 —— 在 PTODSL 暴露完整的 SimtVF kernel 接口(launch 维度 + per-lane entry body),对应 pto.simt_launch + pto.simt_entry。现状:已有 @pto.simt / get_tid_x|y|z / store_vfsimt_info、tracing 能发 pto.simt_entry,需确定simt vf的正式接口。
-
fragment 寄存器数组 —— per-lane 可寻址的 local scratch(如 x_frag[32] / w_frag[32]),lowering 到 llvm.alloca。现状:PTODSL 只有 SSA 标量 + tile,没有可寻址的 per-lane local 数组。
-
向量化 load/store 表达 —— 用类似 Python slice 的语法描述 4-lane(float4)连续访存,lowering 到 LLVM dialect 的 <4 x float> 向量 load/store。现状: 缺少lane-local 128-bit 向量 load/store表达。
-
跨 workitem reduce / all-reduce 接口 —— RMSNorm 的行内规约是跨 128 个 SIMT workitem 的 all-reduce(AscendAllReduce<SumOp, 128>,两级:warp 内 redux + 跨 warp shuffle + syncthreads)。PTODSL 要能实现它——暴露底层原语 pto.shuffle_bfly / pto.redux_add|max|min / pto.syncthreads(PTO 已有,见 include/PTO/IR/VPTOOps.td),或提供一个高层 all-reduce 接口。现状:PTODSL 只暴露 SIMD 的 intra-vreg 规约(vcadd/vcgadd 等),没有跨 workitem 的 shuffle/redux/syncthreads。
B. 对 PTODSL 编译生成 PTO IR 的需求
-
pto.simtscope(内联 SIMT region op) —— PTODSL 把 SimtVF 生成内联 pto.simtscope(晚 outline);unroll、生成 keep/resume、cleanup 都在内联形态做完再 outline 成 simt_entry。现状:只有 outline 后的 simt_entry + simt_launch,没有内联 region。
-
pto.persistent 属性 —— 标在 fragment 的 llvm.alloca op 上(对应 TIR 的 tl.simt_persistent);驻留 pass 据此识别 def/use,在 simtscope 边界生成 keep/resume。现状: 缺少该表达，需要根据第2点接口再确定。
1/3/4属于RMSNorm基本功能对接需求，2/5/6属于寄存器驻留优化的需求。如果按优先级排序，建议1/3/4优先支持，再支持2/5/6
Summary
PTODSL对接RMSNorm的需求整理
Motivation / use case
场景
RMSNorm 是个 SimtVF kernel:每核串行处理多个 token,每 token 在 SimtVF 内做 float4 向量访存、行内 sum-of-squares、跨 128 lane all-reduce、
rstd = rsqrt(Σx²/d + eps),再y = x * rstd * w。其中权重W对一个核内的 token 循环是不变量,但当前每个 token 进 SimtVF 都要把w从 UB 重读进寄存器一次,每核重复batch/N_CORES次。目标两层:① PTODSL 能表达这个 SimtVF kernel;② 支持跨 SimtVF 寄存器驻留优化——
w一次读进寄存器、后续 token 复用,把UB→VRF从每核 64 次降到 1 次。设计文档:
SIMT_VF→simt_launch/simt_entry、float4 向量访存、AscendAllReduce):https://github.com/KurrinQu/PTOAS/blob/tilelang-rmsnorm-pto-design-doc/docs/designs/TileLang_RMSNorm_to_PTO_Design.mdpto.simtscope、pto.persistent、mem2reg):https://github.com/KurrinQu/PTOAS/blob/tilelang-rmsnorm-pto-design-doc/docs/designs/TileLang_RMSNorm_PersistentFragment_Design.mdA. 对 PTODSL 前端的需求
SimtVF 编程接口 —— 在 PTODSL 暴露完整的 SimtVF kernel 接口(launch 维度 + per-lane entry body),对应
pto.simt_launch+pto.simt_entry。现状:已有@pto.simt/get_tid_x|y|z/store_vfsimt_info、tracing 能发pto.simt_entry,需确定simt vf的正式接口。fragment 寄存器数组 —— per-lane 可寻址的 local scratch(如
x_frag[32]/w_frag[32]),lowering 到llvm.alloca。现状:PTODSL 只有 SSA 标量 + tile,没有可寻址的 per-lane local 数组。向量化 load/store 表达 —— 用类似 Python slice 的语法描述 4-lane(float4)连续访存,lowering 到 LLVM dialect 的
<4 x float>向量 load/store。现状: 缺少lane-local 128-bit 向量 load/store表达。跨 workitem reduce / all-reduce 接口 —— RMSNorm 的行内规约是跨 128 个 SIMT workitem 的 all-reduce(
AscendAllReduce<SumOp, 128>,两级:warp 内redux+ 跨 warpshuffle+syncthreads)。PTODSL 要能实现它——暴露底层原语pto.shuffle_bfly/pto.redux_add|max|min/pto.syncthreads(PTO 已有,见include/PTO/IR/VPTOOps.td),或提供一个高层 all-reduce 接口。现状:PTODSL 只暴露 SIMD 的 intra-vreg 规约(vcadd/vcgadd等),没有跨 workitem 的shuffle/redux/syncthreads。B. 对 PTODSL 编译生成 PTO IR 的需求
pto.simtscope(内联 SIMT region op) —— PTODSL 把SimtVF生成内联pto.simtscope(晚 outline);unroll、生成 keep/resume、cleanup 都在内联形态做完再 outline 成simt_entry。现状:只有 outline 后的simt_entry+simt_launch,没有内联 region。pto.persistent属性 —— 标在 fragment 的llvm.allocaop 上(对应 TIR 的tl.simt_persistent);驻留 pass 据此识别 def/use,在 simtscope 边界生成 keep/resume。现状: 缺少该表达，需要根据第2点接口再确定。