ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-02 04:31:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

[CK_TILE] Top-K with Sigmoid kernel (#3062)

Browse Source 
* Add sigmoid option to topk_softmax

* fix formatting

* add to changelog

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Use else if

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com>
This commit is contained in:
Sami Remes
2025-10-28 17:54:06 +00:00
committed by GitHub
parent 6f58d6e457
commit 515e283091
10 changed files with 319 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
include/ck_tile/ops/topk_softmax/kernel/topk_softmax_kernel.hpp
View File

@@ -21,7 +21,7 @@ struct TopkSoftmaxHostArgs
index_t num_experts;
index_t topk;
index_t stride_input; // row stride for input, at least experts
index_t stride_output; // row stride for output/indices, at least tpok
index_t stride_output; // row stride for output/indices, at least topk
};
template <typename Pipeline_>
@@ -45,7 +45,7 @@ struct TopkSoftmaxKernel
index_t num_experts;
index_t topk;
index_t stride_input; // row stride for input, at least experts
index_t stride_output; // row stride for output/indices, at least tpok
index_t stride_output; // row stride for output/indices, at least topk
};
using Kargs = TopkSoftmaxKargs;