leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 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.
include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp
+6 -3
View File
@@ -209,13 +209,14 @@ void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& 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<Element, N>& 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<ElementCompute, float>, "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<ElementOutput> % 128 == 0, "sub-16B alignment not supported yet");
// Reduction tensors