ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-01 04:07:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update moe_flatmm_kernel to manage OOB

Browse Source
This commit is contained in:
Mohsen Saffari
2025-11-18 16:39:09 +00:00
parent 0f894662a1
commit 42d013d007
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/ck_tile/ops/flatmm/kernel/moe_flatmm_kernel.hpp
View File

@@ -1250,6 +1250,8 @@ struct MoeFlatmmKernel
constexpr int MPerThread = TileEncodingPattern::Y2;
statically_indexed_array<statically_indexed_array<index_t, MPerThread>, NumMEpiTile>
c_scatter_offsets;
statically_indexed_array<statically_indexed_array<bool, MPerThread>, NumMEpiTile>
c_scatter_valids;
auto c_coord = dram_tile_distribution.calculate_index();
static_for<0, NumMEpiTile, 1>{}([&](auto mIter) {
static_for<0, MPerThread, 1>{}([&](auto m0) {
@@ -1262,6 +1264,7 @@ struct MoeFlatmmKernel
scatter_token_id =
scatter_token_id * kargs.TopK + (fused_token >> token_id_offset);
c_scatter_offsets[mIter][m0] = scatter_token_id * kargs.stride_C;
c_scatter_valids[mIter][m0] = (scatter_token_id < kargs.NumTokens);
});
});
@@ -1302,7 +1305,8 @@ struct MoeFlatmmKernel
c_block_window.get_window_lengths(),
c_block_window.get_window_origin(),
dram_tile_distribution,
c_scatter_offsets[mIter]);
c_scatter_offsets[mIter],
c_scatter_valids[mIter]);
if constexpr(!IsInputGemm ||
EpiloguePipeline::MemoryOperation == memory_operation_enum::atomic_add)