leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

bugfix generic-k code in top-k with softmax (#1993)

Browse Source 
* bugfix generic-k code in top-k with softmax

* Update include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp

Co-authored-by: Ali Hassani <68103095+alihassanijr@users.noreply.github.com>

* Update examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu

Co-authored-by: Ali Hassani <68103095+alihassanijr@users.noreply.github.com>

---------

Co-authored-by: Ali Hassani <68103095+alihassanijr@users.noreply.github.com>
This commit is contained in:
Tadej Ciglarič
2025-01-31 19:05:35 -05:00
committed by GitHub
co-authored by Ali Hassani
parent 3c28697b9f
commit 6f55278121
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu
+5 -2
View File
@@ -37,8 +37,11 @@
Those assumptions are as:
1. Fusion is over the N dimension.
2. Top-K is either 2 or 4 elements, and the value is static (meaning two kernels have to be
compiled to support both.)
2. Top-K value is static (meaning multiple kernels have to be compiled to support
different values.)
* NOTE: Only K=2 and K=4 cases are performance-optimized and enabled by default.
There is also a generic sort that supports all K values greater than 1, but it can lead to serious performance implications to the underlying kernel.
If necessary, users can simply remove the K==2 || K ==4 assertion under cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp, and the generic sort will automatically be used for all other Ks.
3. The GEMM tile shape along N is greater than or equal to problem size
along N.