KernelSwift: Reuse prepared Triton launch plans

[CearX]

KernelSwift 算子创新大赛

• 队伍：六轩岛
• 赛题：T3 — Gated RMSNorm
• 队长：贝阿朵利切
• 队员：Madeline
• 关联的 ntops PR：KernelSwift: Optimize gated RMSNorm execution ntops#107

修改内容

本 PR 是 T3 的编译器和运行时部分，主要减少重复调用中的 host 端准备开销：

• 缓存并复用 Triton launch plan；
• 使用 tensor 的 shape、stride、dtype、device、offset 以及调用参数区分不同执行计划；
• 对连续使用同一 tensor 的调用提供更轻量的复用路径；
• 限制缓存大小，并在 tensor 生命周期结束后释放对应记录；
• 增加 launch plan 缓存、tensor 生命周期和参数变化相关测试。

Gated RMSNorm 算子实现位于关联的 ntops PR。

测试结果

• 海光 BW：4 个公开 case 均通过正确性检查；
• 天数智芯 BI-V150：4 个公开 case 均通过正确性检查；
• Runtime 测试：61 passed，2 deselected；
• Ruff、贡献规范检查和 git diff --check：通过。

两项 deselected 测试是当前 CoreX 环境不支持的 device compilation case：

• test_triton_tuple_configurations_are_benchmarked_and_cached[cuda]
• test_triton_prepared_cache_releases_gpu_tensor_storage[cuda]

与关联 ntops PR 组合后的公开 4-case 本地端到端测试：

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

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

pytest output:

61 passed, 2 deselected in 3.03s

技术报告与复现

方案

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

• ntops Generate slices only for the innermost indices #107 实现算子、PyTorch 接口和算子测试；
• ninetoothed KernelSwift: Reuse prepared Triton launch plans #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 ntops#107
• KernelSwift: Reuse prepared Triton launch plans #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 许可证。