From 6f552781215f542f36dde611b8bf3922c02a214a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tadej=20Ciglari=C4=8D?= Date: Sat, 1 Feb 2025 01:05:35 +0100 Subject: [PATCH] bugfix generic-k code in top-k with softmax (#1993) * 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> --- .../61_hopper_gemm_with_topk_and_softmax.cu | 7 +++++-- .../epilogue/fusion/sm90_visitor_topk_softmax.hpp | 9 ++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu b/examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu index ac21697e..a5f59ca2 100644 --- a/examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu +++ b/examples/61_hopper_gemm_with_topk_and_softmax/61_hopper_gemm_with_topk_and_softmax.cu @@ -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. diff --git a/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp b/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp index 5ac64423..330e1fde 100644 --- a/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp +++ b/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp @@ -209,13 +209,14 @@ void add_element_to_desc_sorted_array(cutlass::Array& a, Element b) // slower generic path with branching, slower, and can cause register spill CUTLASS_PRAGMA_UNROLL for (int k = 0; k < N; ++k) { - if (a[k] <= b) { + if (a[k] < b) { // Shift down CUTLASS_PRAGMA_UNROLL for (int l = N - 1; l > k; --l) { a[l] = a[l-1]; } a[k] = b; + break; } } } @@ -237,7 +238,7 @@ void merge_desc_sorted_arrays(cutlass::Array& a, const cutlass::Arra int j = 0; CUTLASS_PRAGMA_UNROLL for (int k = 0; k < N; ++k) { - if (a[k] <= b[j]) { + if (a[k] < b[j]) { // Shift down CUTLASS_PRAGMA_UNROLL for (int l = N - 1; l > k; --l) { @@ -334,7 +335,9 @@ template < struct Sm90TopKSoftmaxColReduction { private: static_assert(is_same_v, "Fused Top-K + Softmax reduction requires FP32 accumulation."); - static_assert(TopK == 2 || TopK == 4, "Fused Top-K + Softmax reduction only supports K=2 and K=4."); + static_assert(TopK == 2 || TopK == 4, + "Fused Top-K + Softmax reduction only allows K=2 and K=4, because those cases have been performance-optimized. Other values of K can be enabled by removing this assertion, but they may come with serious performance implications." + ); static_assert(Alignment * sizeof_bits_v % 128 == 0, "sub-16B alignment not supported yet"); // Reduction tensors