Optimize the add_rms_norm_dynamic_quant operator and extract performance optimization points into tilelang-perf-optimization skills

[JasonLee-03]

PR内容：Skill 优化&新增算子

Skill 优化

问题简述

背景： add_rms_norm_dynamic_quant算子实际优化过程经过R0 - R7多轮迭代，之前遇到优化瓶颈等相关问题，PR基于实际开发中的相关问题，对于skill提出相应的优化建议。其中，R0 - R7的优化流程如下：

轮次	优化类型	核心变化	性能得分	提升	状态	关键问题/经验
R0	基线实现	3-pass 串行架构 GM 读取 6x	4.94	-	✅ 基线	每个 pass 重复读 x1/x2 并重算 h
R1	算法层	3-pass → 2-pass 数学变换：max(|h * inv_rms * gamma|) = inv_rms * max(|h * gamma|) GM 读取 6x → 4x	-	-	-	(注：原文本缺失 R1 得分，据前文约贡献 47% 性能提升)
R2	内存带宽层	Double Buffer 三阶段流水 gamma 预加载 向量化广播	11.41	+6.1%	✅ 成功	贡献 12% MTE2/V/MTE3 流水线重叠
R3	Tiling 层	block_M: 4→16 自适应 block_N（H<256 时 block_N=H）	15.61	+36.8%	✅ 成功	贡献 34% UB 利用率 10%→42% 消除尾块
R4	架构层	block_M: 16→32 条件性 x_out readback（双 kernel）	16.40	+5.1%	⚠️ 部分成功	贡献 6% 问题：小 shape 退化 15~47% 原因：未评估分发开销（.item() 约 5~15μs）
R5	指令层	mul_add_dst 融合	16.40	±1%	❌ 失败	浪费 1 轮 问题：bandwidth-bound 上尝试 compute-bound 优化 原因：缺少瓶颈预判
R5a	Tiling 层探索	单维度递增 block_M （4→8→16→32）	-	-	❌ 失败	浪费 1 轮 问题：block_M=16 在双 kernel 下退化 15~28% 原因：未做交叉实验，错过全局最优组合
R6	架构层探索	AUTO_SYNC=False Fixed Core	-	-	❌ 失败	浪费 1 轮 问题 1：AUTO_SYNC=False 导致 17/20 精度失败（V pipe 队列排空问题） 问题 2：Fixed Core 编译失败（TVM StmtSimplifier bug） 原因：编译器限制未记录 意外发现：AUTO_SYNC=True 下需要交替 buffer（hw_a/hw_b）消除 RAW hazard
R7	架构层优化	混合策略 M<1024: 单 kernel + block_M=16 M≥1024: 双 kernel + block_M=32	17.19	+4.8%	✅ 成功	新最高分 67.19 解决 R4 小 shape 退化 分发开销评估

一、优化角度（5个维度）

1. 工作流程优化（新增3个强制门禁步骤）

优化内容：

• Step 0: 瓶颈预判（强制门禁）
  • 计算理论最小耗时（基于 GM 带宽）
  • 判断瓶颈类型（bandwidth/compute/launch-bound）
  • 根据瓶颈类型推荐/禁止优化方向
• Step 3.5: 多维度交叉实验矩阵（禁止单维度递增）
  • 识别独立维度（Tiling、架构、同步模式、Pass 数量）
  • 构建 2×2 或 3×3 交叉矩阵
  • 快速验证所有组合，选出最优配置
• Step 4.5: 组合优化（将各轮最优配置合并验证）
  • 汇总各轮最优配置
  • 检查配置冲突（Tiling vs 架构、同步 vs 指令）
  • 实现组合版本，完整验证精度 + 性能

优化理由：

• R5 在 bandwidth-bound 上尝试 compute-bound 优化（浪费 1 轮）
• R5a 单维度递增 block_M（浪费 1 轮）
• R7 的混合策略事后才发现（缺少组合验证）

2. 参考文档完善（新增2个关键文档）

优化内容：

• 新增 references/compiler-limitations.md
  • 记录 5 个已知编译器限制：
    1. AUTO_SYNC=False + V pipe 指令队列排空问题
    2. Fixed Core + 空循环 range 约束冲突
    3. mul_add_dst 在 bandwidth-bound 上无收益
    4. 双 kernel 分发开销对小 shape 的影响
    5. AUTO_SYNC=True 下连续 tile 指令的 V pipe RAW hazard
  • 提供安全特性清单（已验证可用）
• 新增 references/best-practices/vector_fused_operator_optimization.md
  • 完整的 R0-R7 优化经验总结
  • 6 个优化手段详解（Pass 缩减、Readback、Tiling、Double Buffer、向量化广播、混合策略）
  • 关键代码模式（Newton-Raphson rsqrt、交替 buffer）
  • 5 个常见陷阱
  • 最佳实践建议和检查点

优化理由：

• R6 尝试 AUTO_SYNC=False 和 Fixed Core 都失败（编译器限制未记录）
• 多 pass Vector 融合算子缺少完整的优化案例

3. 决策指导强化（新增3个关键章节）

优化内容：

• 新增 §一.五 算法层优化（最高优先级，先于核内优化）
  • Pass 数量缩减（GM 读取 -33%~50%）
  • Readback 模式（避免 Pass 2 重算）
  • 自适应 Tiling 参数（消除尾块）
  • 关键数学变换：max(|h * inv_rms * gamma|) = inv_rms * max(|h * gamma|)
• 新增 §2.2.2 交替 buffer 消除 V pipe RAW hazard
  • 问题：AUTO_SYNC=True 下连续 tile 指令存在 RAW 依赖
  • 解决：使用 hw_a/hw_b 交替 buffer
  • 代码示例：Pass 1 abs_max 计算链、Pass 2 量化链
• 新增 Step 3.5.4 分发开销评估（双 kernel 前必做）
  • 评估 .item() host 同步开销（约 5~15μs）
  • 计算开销占比：overhead_ratio = dispatch_overhead_us / kernel_us
  • 决策：占比 > 10% → 单 kernel；< 5% → 双 kernel；5~10% → 混合策略

优化理由：

• R1 的 3-pass → 2-pass 贡献 47% 性能提升，但原 skill 未强调算法层优先级
• R6 发现 AUTO_SYNC=True 下存在 RAW hazard，需要交替 buffer
• R4 引入双 kernel 未评估分发开销，小 shape 退化 15~47%

4. 反模式识别（新增2个反模式）

优化内容：

• 新增“单维度递增搜索（未做交叉实验）”反模式
  • 识别特征：逐轮递增某个参数（如 block_M: 4→8→16→32）
  • 性能原因：不同维度紧耦合，单维度递增无法发现交互效应
  • 替代写法：构建交叉实验矩阵，一轮覆盖所有组合
• 新增“正交轴串行化（Scalar Scan on Parallelizable Axis）”反模式
  • 识别特征：两层嵌套循环，外层遍历正交轴，内层是标量操作
  • 性能原因：串行处理正交轴浪费向量化并行能力
  • 替代写法：transpose 使并行轴内存连续，折叠为向量操作

优化理由：

• R5a 单维度递增 block_M，未与 kernel 架构做交叉实验（浪费 1 轮）
• 某些算子（如 cummin）存在正交轴串行化问题，但原反模式清单未覆盖

5. 代码示例更新（更新同步模式和Vector核示例）

优化内容：

• 更新 §2.2.0 同步模式决策表
  • 原推荐：纯 AIV Vector 算子 → AUTO_SYNC=False
  • 新推荐：纯 AIV Vector 算子 → AUTO_SYNC=True（当前编译器推荐）
  • 添加编译器限制说明和替代方案
• 更新 Vector 核示例代码
  • 原示例：使用 AUTO_SYNC=False
  • 新示例：使用 AUTO_SYNC=True + 手动三路 flag（mte3→mte2、mte2→v、v→mte3）
  • 添加详细的 flag 时序说明和易错点提示

优化理由：

• R6 发现 AUTO_SYNC=False 存在 V pipe 队列排空问题（17/20 精度失败）
• 原示例代码使用 AUTO_SYNC=False，会误导后续 Vector 算子

---

二、之前 Skill 的漏洞和优化逻辑问题（8个问题）

漏洞1：缺少瓶颈预判机制

• 问题描述：原 skill 没有强制在优化前做瓶颈类型判断，导致在 bandwidth-bound 算子上尝试 compute-bound 优化。
• 具体表现：

  # R5 的错误做法
  # 算子当前耗时 1558μs / 理论最优 299μs = 5.2x > 3x → bandwidth-bound
  # 在 bandwidth-bound 上尝试 mul_add_dst 融合 → 性能变化在噪声范围内（±1%）
  T.tile.mul_add_dst(sum_sq_acc, h_fp32, h_fp32)  # 无效优化

• 后果：浪费 1 轮迭代。
• 修复方案：新增 Step 0 瓶颈预判（强制门禁）

漏洞2：允许单维度递增搜索

• 问题描述：原 skill 没有要求做交叉实验，允许逐轮单维度递增搜索，导致 block_M 与 kernel 架构的紧耦合关系未被发现。
• 具体表现：

  # R5a 的错误做法
  # 逐轮递增 block_M: 4→8→16→32
  # block_M=16 在双 kernel 下退化 15~28%（与 kernel 架构紧耦合）
  # 错过全局最优组合："单 kernel + block_M=16" 和 "双 kernel + block_M=32"

• 后果：浪费 1 轮迭代，错过全局最优。
• 修复方案：新增 Step 3.5 多维度交叉实验矩阵

漏洞3：缺少组合优化验证

• 问题描述：原 skill 没有强制在所有单维度优化后做组合验证，导致混合策略事后才被发现。
• 具体表现：

  # R4 的问题
  # 所有 shape 使用同一 kernel（block_M=32 + 双 kernel）
  # 小 shape (M=256): 8.9μs ← 分发开销占比 > 10%
  # R7 事后才发现需要混合策略

• 后果：小 shape 退化 15~47%，R7 才修复。
• 修复方案：新增 Step 4.5 组合优化

漏洞4：编译器限制未记录

• 问题描述：原 skill 没有记录已知的编译器限制，导致尝试已知失败的特性。
• 具体表现：

  # R6 的错误做法 1: AUTO_SYNC=False
  # barrier_all() 无法排空 V pipe 指令队列
  # 当 n_num > 1 时，后续 scalar 操作读到旧值
  # 结果：17/20 精度失败
  
  # R6 的错误做法 2: Fixed Core
  # m_num=8 < 24 cores → TVM StmtSimplifier InternalError
  # 空循环 range 约束冲突
  # 结果：编译失败

• 后果：浪费 1 轮迭代（AUTO_SYNC=False + Fixed Core 都失败）。
• 修复方案：新增 references/compiler-limitations.md

漏洞5：算法层优化优先级不明确

• 问题描述：原 skill 没有强调算法层优化（pass 数量缩减）的最高优先级，可能导致直接进入核内优化，错过最大收益。
• 具体表现：

  # 可能的错误做法
  # 直接进入 Double Buffer / Tiling 调优
  # 错过 3-pass → 2-pass 的算法层优化（GM 读取 -33%，性能 +117%）

• 后果：可能错过 47% 的性能提升（R1 贡献）。
• 修复方案：新增 §一.五 算法层优化（最高优先级）

漏洞6：分发开销评估缺失

• 问题描述：原 skill 没有要求在引入双 kernel 前评估分发开销，导致小 shape 退化。
• 具体表现：

  # R4 的问题
  # 引入双 kernel（readback/recompute）但未评估 .item() 分发开销
  # 小 shape (M=256, kernel 4.7μs)：分发开销 10μs 占比 > 100% → 退化 -89%

• 后果：小 shape 退化 15~47%，R7 才通过混合策略修复。
• 修复方案：新增 Step 3.5.4 分发开销评估（双 kernel 前必做）

---

新增算子

新增 add_rms_norm_dynamic_quant 算子，以下为CANN-BENCH测试报告，在A3上进行测试，对比baseline为pytorch npu下实现的实测结果，加速比为0.64x，以下为精度与性能对比测试报告：

算子评测报告

评测代号: cann_final_eval_20260717_191555
评测时间: 2026-07-17T19:15:55.319301
设备: npu:0
框架版本: V0.4.0
评测集版本: tasks-v0.4.0

概览

指标	数值
评测算子数	1
总用例数	20
通过用例数	20
失败用例数	0
通过率	100.00%
综合得分	67.19

算子详情

AddRmsNormDynamicQuant（level3/add_rms_norm_dynamic_quant）

指标	数值
用例数	20
通过数	20
失败数	0
通过率	100.00%
平均加速比	0.64x
编译/运行得分	20.00
精度得分	30.00
性能得分	17.19
得分	67.19

用例ID	状态	耗时(μs)	加速比	精度误差
level3/add_rms_norm_dynamic_quant_1	✅	1541.89	0.37x	1.000000
level3/add_rms_norm_dynamic_quant_2	✅	1575.25	0.44x	1.000000
level3/add_rms_norm_dynamic_quant_3	✅	1542.77	0.37x	1.000000
level3/add_rms_norm_dynamic_quant_4	✅	1453.83	0.52x	1.000000
level3/add_rms_norm_dynamic_quant_5	✅	1645.39	0.38x	1.000000
level3/add_rms_norm_dynamic_quant_6	✅	2761.10	0.57x	1.000000
level3/add_rms_norm_dynamic_quant_7	✅	360.85	0.44x	1.000000
level3/add_rms_norm_dynamic_quant_8	✅	414.49	0.36x	1.000000
level3/add_rms_norm_dynamic_quant_9	✅	345.27	0.52x	1.000000
level3/add_rms_norm_dynamic_quant_10	✅	376.95	0.85x	1.000000
level3/add_rms_norm_dynamic_quant_11	✅	433.91	0.52x	1.000000
level3/add_rms_norm_dynamic_quant_12	✅	454.63	0.38x	1.000000
level3/add_rms_norm_dynamic_quant_13	✅	246.11	0.44x	1.000000
level3/add_rms_norm_dynamic_quant_14	✅	352.87	0.51x	1.000000
level3/add_rms_norm_dynamic_quant_15	✅	10.94	0.91x	0.000000
level3/add_rms_norm_dynamic_quant_16	✅	45.56	0.54x	1.000000
level3/add_rms_norm_dynamic_quant_17	✅	4.80	2.08x	0.000000
level3/add_rms_norm_dynamic_quant_18	✅	7.66	1.31x	0.000000
level3/add_rms_norm_dynamic_quant_19	✅	12.22	0.82x	0.000000
level3/add_rms_norm_dynamic_quant_20	✅	219.50	0.39x	1.000000

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀

[JasonLee-03]

[gemini-code-assist]

Code Review

This pull request introduces comprehensive performance optimization guidelines, compiler limitation documentation, and a highly optimized 2-pass implementation of the fused AddRmsNormDynamicQuant operator for Ascend NPU. The review feedback highlights critical issues in the new operator implementation, including severe precision loss from using CAST_RINT instead of CAST_NONE during float32 to float16/bfloat16 casting, potential V pipe RAW hazards in the dual-kernel implementations under AUTO_SYNC=True, and out-of-bounds memory access risks when the dimension H is not a multiple of block_N due to a lack of partial block handling.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

在将 h_fp32（残差相加结果）转换为 out_dtype_ub 并写入 x_out 时，代码使用了 "CAST_RINT" 模式。"CAST_RINT" 会将浮点数四舍五入为最接近的整数。由于 h_fp32 是连续的浮点值（非整数），使用 "CAST_RINT" 会导致 x_out 丢失所有小数部分，造成严重的精度损失。应该使用 "CAST_NONE" 进行标准的 float32 到 float16/bfloat16 的类型转换，以保留小数精度。此问题同样存在于第 182 行、第 419 行、第 443 行、第 679 行和第 703 行，请一并修改。

Suggested change

	T.tile.cast(out_dtype_ub, h_fp32, "CAST_RINT", tile_elements)
	T.tile.cast(out_dtype_ub, h_fp32, "CAST_NONE", tile_elements)

[gemini-code-assist]

在 _kernel_readback 中，hw_fp32 被用于连续的 in-place 操作（例如 Pass 1 中的 mul 之后紧跟 abs，以及 Pass 2 中的连续 in-place 链）。根据您在 compiler-limitations.md 和 optimization-guide.md 中记录的编译器限制，在 AUTO_SYNC=True 下，这种连续的 dst -> src 依赖链极易触发 V pipe RAW 冒险（RAW hazard），导致非确定性的精度测试失败。建议像在 _kernel_single 中那样，引入交替 buffer hw_a 和 hw_b 来消除这种冒险。此问题同样存在于 _kernel_recompute（第 361 行），请一并修改。

Suggested change

	hw_fp32 = T.alloc_ub([ROWS, block_N], "float32")
	hw_a = T.alloc_ub([ROWS, block_N], "float32")
	hw_b = T.alloc_ub([ROWS, block_N], "float32")