ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-20 14:59:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add KBatch support for gemm_ab_scale (#2740)

Browse Source 
* Add KBatch support for gemm_ab_scale

* Revert kernel parameters change

* Remove printing

* fix formatting

* fix check

* Use {} in if

---------

Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com>
This commit is contained in:
Sami Remes
2025-10-09 07:33:16 +01:00
committed by GitHub
parent e99356dabc
commit 9d4bfe3932
5 changed files with 34 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
include/ck/tensor_operation/gpu/device/device_gemm_multiple_d_ab_scale.hpp
View File

@@ -58,6 +58,8 @@ struct DeviceGemmMultipleD_ABScale : public BaseOperator
CDEElementwiseOperation cde_element_op) = 0;
virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
virtual void SetKBatch(BaseArgument* arg, int KBatch) const = 0;
};
template <typename ALayout,

6
include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3_ab_scale.hpp
View File

@@ -311,6 +311,12 @@ struct DeviceGemmMultiD_ABScale_Xdl_CShuffle_V3
}
};
void SetKBatch(BaseArgument* base_arg, int KBatch) const override
{
auto& arg = *dynamic_cast<Argument*>(base_arg);
arg.KBatch = KBatch;
}
static constexpr bool IsValidCompilationParameter()
{
// TODO: properly implement this check