@pradheepraop: 实现了 MSA 论文中内核设计部分的 top-k 内核。https://github.com/Mantissagithub/learn_c…

X AI KOLs Timeline 工具
top-k kernel-design cuda gpu-programming attention open-source

摘要

从 MSA 论文内核设计部分实现了 top-k 内核，使用免指数比较和通过 CUDA shuffle 进行的 warp 级树合并。代码已发布于 GitHub。

实现了 MSA 论文内核设计部分的 top-k 内核。 https://github.com/Mantissagithub/learn_cuda/blob/msa/07_projects/msa/top_k.cu… 其中包含两个思路： - 免指数比较：无需计算 softmax，因为 softmax 保持顺序不变。 - 每个 warp 线程（lane）扫描 1/32 的步幅，维持一个小的本地 top-k，然后通过 shuffle 指令进行树合并。 整晚都在修改 CUDA，并清理了我的 learn_cuda 仓库，欢迎任何反馈或优化建议。
查看原文
查看缓存全文

缓存时间: 2026/06/15 13:04

实现了 msa 论文中内核设计部分的 top-k 内核。

https://github.com/Mantissagithub/learn_cuda/blob/msa/07_projects/msa/top_k.cu…

它基于两个思路：

  • 免指数比较：无需计算 softmax，因为 softmax 不改变大小顺序
  • 每个 warp 线程扫描 1/32 步长，维护一个小的本地 top-k，并通过洗牌操作树形合并结果

结果熬夜修改了 CUDA 代码，也清理了我的 learn_cuda 仓库，因此欢迎任何反馈或优化建议。

Mantissagithub/learn_cuda

Source: https://github.com/Mantissagithub/learn_cuda

learn_cuda

学习 CUDA（即 GPU 编程）

相似文章

MSLK kernel reference (Website)

TLDR AI

Documentation reference for MSLK 1.3.0, a library of fused GPU kernels for transformer workloads including attention, quantization, GEMM, and MoE routing, supporting CUDA and ROCm with PyTorch integration.

AdaExplore：基于失败驱动的自适应与多样性保留搜索的高效内核生成

arXiv cs.CL

来自卡内基梅隆大学、华盛顿大学和Arm的研究人员提出了AdaExplore，这是一种用于GPU内核代码生成的LLM智能体框架。该框架通过失败驱动自适应与多样性保留搜索技术，在不进行额外微调的情况下，在KernelBench Level-2和Level-3基准测试中分别实现了3.12倍和1.72倍的加速。