From 50ee0ac283091bcd17f5a466a55571f9faafc26a Mon Sep 17 00:00:00 2001
From: jakpiase <jakub.piasecki@amd.com>
Date: Tue, 26 Nov 2024 13:56:32 +0100
Subject: [PATCH] Add check for bf16 splitk support for grouped gemm splitk
 (#1673)

* add check for bf16 splitk support for grouped gemm splitk

* Update if condition

---------

Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com>

[ROCm/composable_kernel commit: b70f367f8051e0c66071a25ab95a77e076762808]
---
 .../device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp  | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp
index 6d9d1459c8..cb0afbb08d 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp
@@ -538,6 +538,11 @@ struct DeviceGroupedGemmXdlSplitKCShuffle : public DeviceGroupedGemmSplitK<ALayo
             return false;
         }
 
+        if(std::is_same_v<EDataType, ck::bhalf_t> && arg.K_BATCH > 1 && !is_bf16_atomic_supported())
+        {
+            return false;
+        }
+
         bool supported = true;
         for(std::size_t i = 0; i < arg.gemm_kernel_args_.size(); ++i)
         {