SuperAPI AI 工具集

cutedsl-kernel-integration

编程与开发

查看项目

更新于 10-16 09:46

技能介绍

用途:在 cuDNN Frontend 中新增或更新一个仅前端的 CuTeDSL API,要求完成公开函数、导出、文档和测试的完整集成。

编辑前准备:检查仓库当前状态,避免覆盖无关改动;确认集成所需的每个原始源文件都存在,缺失时应报告而不是从相近内核推断契约;记录源出处(上游 URL、本地路径、commit,以及哪些文件对应公开 API 模块、哪些是私有 helper)。选择模板前先对内核分类:内核家族(dense GEMM、GEMM 融合、grouped GEMM、discrete grouped GEMM、MoE、attention、sparse attention 等)、执行拓扑(单内核、成对的前向/反向 API、多内核编排、helper 内核设置、分布式/运行时协调、内部调度器)、公开接口(<op_name>_<direction> 与 _jax 函数约定,返回张量、可选输出、workspace 归属、导入导出命名空间)、内部支持(helper 模块、调度器、元数据工具、生成的描述符应保持私有)、架构变体(是否需为新 GPU 透明分发到备用 CuTeDSL 模块,如 Rubin sm107 与默认 SM100;分发在内部进行时公开函数保持不变)。实现前须阅读 references/integration-pattern.md。

集成流程:在最近的既有家族下添加或更新操作包,如 python/cudnn/<operation>/、python/cudnn/gemm/cutedsl/dense/<operation>/、grouped/、discrete_grouped/ 或 python/cudnn/sdpa/<direction>/;在 api.py 实现公开函数,做配置校验、输出分配、按模板缓存或复用已编译内核,并返回 PyTorch 的 TupleDict 或 JAX 的元组/字典;通过操作或家族 init.py 以及 python/cudnn/init.py 的 _LAZY_OPTIONAL_IMPORTS 导出;除非确需新包,复用现有 CuTeDSL 依赖(cutedsl extra 目前只含 cuda-python);添加 FE OSS 文档并更新概览或操作索引链接;在 test/python/<operation>/cutedsl/ 添加测试,包含支持校验,可执行时补充数值/参考覆盖;对 grouped/discrete/MoE/SDPA 内核保留源 helper 与调度器拓扑,共享 helper 模块应作为包内文件而非公开 cudnn 导出;当 SM100 内核需 Rubin(sm107) 变体时,按 integration-pattern.md 的架构分发模式处理,不新增公开 API(现有示例:grouped_gemm_quant、grouped_gemm_glu、grouped_gemm_dglu)。

验证:对改动文件运行针对性格式化或测试;仅改技能时至少校验 SKILL.md frontmatter 合法且引用的路径都存在;内核集成在具备所需 GPU 与可选依赖的环境运行 pytest test/python/<operation>/cutedsl/test_<operation>.py,否则明确报告跳过验证;架构分发工作还需运行 pytest test/python/gemm/cutedsl/test_rubin_kernel_dispatch.py,在 Rubin 硬件上受影响操作的既有 FE API 端到端测试应在无 API 变更下仍通过。

限制:资料的官方正文未提供,具体内核版本、依赖版本与硬件可用性暂未提供;缺失源文件或环境不满足时应如实报告,不做推断。