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

save tmp

Browse Source
This commit is contained in:
Max Podkorytov
2025-12-16 21:28:53 -06:00
parent 11f7c3f136
commit 88538ca6d2
3 changed files with 21 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
example/ck_tile/17_grouped_gemm/CMakeLists.txt
View File

@@ -6,6 +6,7 @@ if(GPU_TARGETS MATCHES "gfx94|gfx95")
add_executable(tile_example_quant_grouped_gemm quant_grouped_gemm.cpp)
add_executable(tile_example_grouped_gemm_preshuffle grouped_gemm_preshuffle.cpp)
add_executable(tile_example_grouped_gemm_multi_d grouped_gemm_multi_d.cpp)
add_executable(tile_example_grouped_gemm_persistent_async grouped_gemm_persistent_async.cpp)
set(EXAMPLE_GEMM_COMPILE_OPTIONS)
if(CK_USE_OCP_FP8)
list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)

34
example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.cpp
View File

@@ -279,20 +279,18 @@ int run_grouped_gemm_persistent_async_example(ck_tile::ArgParser& arg_parser,
return m;
};
// Create GEMM descriptor
grouped_gemm_kargs desc;
desc.M = M;
desc.N = N;
desc.K = K;
desc.k_batch = kbatch;
desc.stride_A = get_stride(a_layout, M, K);
desc.stride_B = get_stride(b_layout, K, N);
desc.stride_E = get_stride(c_layout, M, N);
desc.a_ptr = a_dev_bufs[i]->GetDeviceBuffer();
desc.b_ptr = b_dev_bufs[i]->GetDeviceBuffer();
desc.e_ptr = c_dev_bufs[i]->GetDeviceBuffer();
gemm_descs.push_back(desc);
gemm_descs.push_back({a_dev_bufs[i]->GetDeviceBuffer(),
b_dev_bufs[i]->GetDeviceBuffer(),
{/*ds_ptr*/},
c_dev_bufs[i]->GetDeviceBuffer(),
kbatch,
M,
N,
K,
get_stride(a_layout, M, K),
get_stride(b_layout, K, N),
{/*stride_Ds*/},
get_stride(c_layout, M, N)});
}
// Allocate workspace for kernel arguments
@@ -303,12 +301,10 @@ int run_grouped_gemm_persistent_async_example(ck_tile::ArgParser& arg_parser,
std::vector<ck_tile::GemmTransKernelArg<>> kargs;
int cumulative_chunks = 0;
for(int i = 0; i < gemm_descs.size(); ++i)
for(const auto& desc : gemm_descs)
{
const auto& desc = gemm_descs[i];
// Calculate chunk offset for this group (if async mode enabled)
int chunk_offset = cumulative_chunks;
// int chunk_offset = cumulative_chunks;
if(async_args.enable_async && async_args.tiles_per_chunk_m > 0)
{
const int m_tiles = (desc.M + GemmConfig::M_Tile - 1) / GemmConfig::M_Tile;
@@ -331,6 +327,8 @@ int run_grouped_gemm_persistent_async_example(ck_tile::ArgParser& arg_parser,
desc.k_batch});
}
ck_tile::ignore = cumulative_chunks;
// Copy kernel args to device
const auto stream = ck_tile::stream_config{nullptr, true, 1, warmup, repeat};
void* kargs_ptr = gemm_workspace.GetDeviceBuffer();

8
example/ck_tile/17_grouped_gemm/grouped_gemm_persistent_async.hpp
View File

@@ -13,7 +13,7 @@ template <typename GroupedGemKernelParam,
typename ALayout,
typename BLayout,
typename CLayout>
void invoke_grouped_gemm_persistent(const ck_tile::stream_config& s,
float invoke_grouped_gemm_persistent(const ck_tile::stream_config& s,
const ck_tile::index_t num_groups,
void* kargs_ptr,
bool splitk)
@@ -90,7 +90,7 @@ void invoke_grouped_gemm_persistent(const ck_tile::stream_config& s,
<< blocks.x << ", " << blocks.y << ", " << blocks.z << "}" << std::endl;
}
ck_tile::launch_kernel(s,
return ck_tile::launch_kernel(s,
ck_tile::make_kernel<kBlockPerCu>(
Kernel{},
grids,
@@ -102,13 +102,13 @@ void invoke_grouped_gemm_persistent(const ck_tile::stream_config& s,
if(splitk)
{
Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::atomic_add>{});
}
else
{
Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
return Run(ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::set>{});
}
}