@pradheepraop: implemented the top-k kernel from the kernel design section in the msa paper. https://github.com/Mantissagithub/learn_c…

X AI KOLs Timeline Tools

Summary

Implemented a top-k kernel from the kernel design section of the MSA paper, using exp-free comparison and warp-level tree merging with CUDA shuffles. The code is available on GitHub.

implemented the top-k kernel from the kernel design section in the msa paper. https://github.com/Mantissagithub/learn_cuda/blob/msa/07_projects/msa/top_k.cu… it comes with two ideas: - exp-free comparison: no need to compute softmax, since softmax preserves ordering - each warp lane scans a 1/32 stride, keeps a small local top-k, and tree-merges the results with shuffles ended up revising cuda the whole night and also cleaned up my learn_cuda repo, so any feedback/optimizations are welcome.
Original Article
View Cached Full Text

Cached at: 06/15/26, 01:04 PM

implemented the top-k kernel from the kernel design section in the msa paper.

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

it comes with two ideas:

  • exp-free comparison: no need to compute softmax, since softmax preserves ordering
  • each warp lane scans a 1/32 stride, keeps a small local top-k, and tree-merges the results with shuffles

ended up revising cuda the whole night and also cleaned up my learn_cuda repo, so any feedback/optimizations are welcome.


Mantissagithub/learn_cuda

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

learn_cuda

Learning cuda (aka gpu programming)

Similar Articles

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.