执行摘要
新增 CuTeDSL KDA 解码内核,为 KDA 架构模型提供约 1.05x 性能提升。
PR body中说明,为支持Kimi-Linear/Kimi-2.5和其他KDA架构模型,引入CuTeDSL解码内核以提升性能。基准测试在H800上进行,batch_size=1时获得约1.05x性能提升,但当前仅支持解码,未集成到e2e后端,因为预填充内核布局不匹配(VK vs KV),计划后续支持CuTeDSL预填充内核或调整Triton内核。
建议技术管理者和工程师精读此PR,重点关注:
- 内核实现文件
cutedsl_kda.py中的设计决策,如K维度门控处理和VK布局适配,以理解性能优化技巧。 - review中的线程安全讨论和布局统一权衡,这些揭示了长期架构演进方向。
- 基准测试脚本
bench_cutedsl_kda_decode.py,学习正确性验证和性能测量方法。
review中的核心讨论包括:
- 线程安全问题:gemini-code-assist[bot]指出
_compiled_kernels和_cu_seqlens_cache等全局缓存缺乏同步,可能在高并发下引发竞争条件,但未在review中解决。 - 设计权衡:在Issue评论中,kaixih建议统一VK布局以避免复杂性,而非添加KV内核;yuan-luo回复修改到VK非小事,需实现CuTeDSL预填充内核,揭示了架构演进方向。
- 代码风格:BBuf建议重命名文件
cutedsl_kda.py为kda_cutedsl.py,但yuan-luo回复遵循GDN风格,保留原名;kaixih要求移除调试打印语句和中文注释,已执行。 - 性能测量:gemini-code-assist[bot]指出基准测试中的回退计时机制不准确,但yuan-luo忽略,因为它在异常分支中。
参与讨论