KernelSwift: Optimize gated RMSNorm execution

[CearX]

KernelSwift 算子创新大赛

• 队伍：六轩岛
• 赛题：T3 — Gated RMSNorm
• 队长：贝阿朵利切
• 队员：Madeline
• 关联的 ninetoothed PR：KernelSwift: Reuse prepared Triton launch plans ninetoothed#216

修改内容

本 PR 是 T3 的算子部分：

• 新增 Gated RMSNorm 的 NineToothed 实现和 PyTorch 接口；
• 支持 norm before/after gate、SiLU/Swish、Sigmoid、无 gate、显式 group size 和非连续输入；
• 归约和激活使用 float32 计算，输出转换回输入 dtype；
• 减少 wrapper 中重复的 shape、group 和 activation 处理；
• 增加正确性测试、smoke test 和 benchmark 脚本。

编译器和运行时优化位于关联的 ninetoothed PR。

测试结果

• 海光 BW：4 个公开 case 均通过正确性检查；
• 天数智芯 BI-V150：4 个公开 case 均通过正确性检查；
• Gated RMSNorm 测试：135 passed；
• Ruff 和 git diff --check：通过。

公开 4-case 本地端到端测试：

平台	优化前	优化后	加速比
海光 BW	457.512 µs	228.287 µs	2.0041×
天数智芯 BI-V150	422.3785 µs	214.2483 µs	1.9714×

正式成绩以组委会评测环境为准。

pytest output:

135 passed in 124.17s (0:02:04)

技术报告与复现

方案

T3 的目标是实现 Gated RMSNorm，并降低端到端调用开销。作品由两个 PR 组成：

• ntops KernelSwift: Optimize gated RMSNorm execution #107 实现算子、PyTorch 接口和算子测试；
• ninetoothed #216 优化 Triton launch plan 的准备与复用。

算子将最后一维按 group_size 分组，每个逻辑分组由一个 program 处理，在同一个 application 中完成归约、归一化、gate 激活和 weight 乘法。中间计算使用 float32，结果转换回输入 dtype。gate 可以放在归一化之前或之后，并支持 SiLU/Swish、Sigmoid 和无 gate 模式。

运行时优化针对重复调用时的 host 开销。第一次调用完成正常的参数准备，后续参数布局一致时复用 launch plan；缓存键包含 tensor 布局和调用参数，缓存数量有上限。该优化位于通用 Triton materializer/runtime，不包含平台专用分支。

实验方法

题面未公开 T3 的正式评测脚本。本地验证使用 PR 中的 smoke test、pytest 和 benchmark 脚本；海光 BW 和天数智芯 BI-V150 使用相同的 4 个 case。每个 case 先与 PyTorch reference 做正确性比较，再测完整 forward 的 wall time；计时前 warmup，计时后同步设备。表格中的结果取多轮测量中位数，优化前后使用独立进程和缓存。

版本以以下两个 Draft PR 的当前提交为准：

• KernelSwift: Optimize gated RMSNorm execution #107
• KernelSwift: Reuse prepared Triton launch plans ninetoothed#216

构建与复现

git clone -b t3-b-submission-ready https://github.com/CearX/ninetoothed.git
git clone -b t3-a-submission-ready https://github.com/CearX/ntops.git

python3 -m venv .venv
. .venv/bin/activate
python -m pip install -U pip
python -m pip install -e './ninetoothed[all]'
python -m pip install -e './ntops[testing]'

# 算子正确性
python ntops/scripts/run_gated_rms_norm_smoke.py --device cuda
python -m pytest -q ntops/tests/test_gated_rms_norm.py

# ninetoothed runtime 测试
python -m pytest -q ninetoothed/tests/test_triton_runtime_auto_tuning.py

# 快速性能测试；完整测试去掉 --quick
python ntops/scripts/benchmark_gated_rms_norm.py \
  --device cuda --quick --output t3-results.jsonl

在当前 CoreX 环境中，runtime 测试中的两个 device compilation case 需要排除，命令如下：

python -m pytest -q ninetoothed/tests/test_triton_runtime_auto_tuning.py \
  -k 'not triton_tuple_configurations_are_benchmarked_and_cached and not triton_prepared_cache_releases_gpu_tensor_storage'

工程说明

A 部分共新增约 1700 行，但 Gated RMSNorm 的 kernel 和 PyTorch 接口合计约 200 行，算子测试约 130 行，其余主要是 benchmark、launch profiler 和 smoke test。benchmark 负责固定 case、同步计时和结果输出；profiler 用于拆分 wrapper、launch preparation 与 kernel 时间。这些脚本只用于复现和定位性能，不参与算子运行。

B 部分约一半是 runtime 测试，生产代码主要分布在通用 runtime 和 Triton materializer。runtime 管理调用参数、缓存和 tensor 生命周期，materializer 负责 Triton ABI、执行计划重绑定以及 auto-tuning 路径。为了让新建的 input/output tensor 也能复用已经准备好的 launch plan，缓存键需要覆盖 shape、stride、dtype、device、storage offset、标量参数和调用形式；普通 launch 与 auto-tuned launch 都需要处理这一过程。

当前实现的主要维护问题是两条 launch 路径中仍有相似的缓存查找和 plan promotion 逻辑，同时保留了通用 structural key 与 Triton observer 两种 key 构造方式。observer 用于降低重复遍历 ABI 的 host 开销，通用 key 负责其他调用形式，但两者的职责还可以进一步收拢。现有测试覆盖了布局、标量、alias、tensor 生命周期、缓存淘汰和 auto-tuning；后续重构方向是抽取共用的缓存控制流程，同时保留轻量 observer。

接口和数学定义参考 vLLM RMSNormGated，本次改动未引入第三方源码。ntops、ninetoothed 和 vLLM 均采用 Apache-2.0 许可证。