add type check

include/ck/tensor_operation/gpu/device/impl/device_batched_mha_bwd_xdl_cshuffle_qloop_light_v1.hpp

@@ -1293,6 +1293,16 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
             return false;
         }
 
+        // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+        if constexpr(is_same<OutputDataType, half_t>::value ||
+                     is_same<OutputDataType, bhalf_t>::value)
+        {
+            if(KzRaw % 2 != 0)
+            {
+                std::cout << "K_q must be a multiple of 2" << std::endl;
+                return false;
+            }
+        }
         // Check vector load/store requirement
         const auto a_stride_lowest =
             ABlockTransferSrcVectorDim == 2 ? arg.a_mz_kz_strides_[1] : arg.a_mz_kz_strides_[0];

include/ck/tensor_operation/gpu/device/impl/device_batched_mha_bwd_xdl_cshuffle_qloop_light_v2.hpp

@@ -1325,6 +1325,16 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
             return false;
         }
 
+        // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+        if constexpr(is_same<OutputDataType, half_t>::value ||
+                     is_same<OutputDataType, bhalf_t>::value)
+        {
+            if(KzRaw % 2 != 0)
+            {
+                std::cout << "K_q must be a multiple of 2" << std::endl;
+                return false;
+            }
+        }
         // Check vector load/store requirement
         const auto a_stride_lowest =
             ABlockTransferSrcVectorDim == 2 ? arg.a_mz_kz_strides_[1] : arg.a_mz_kz_strides_[0];

include/ck/tensor_operation/gpu/device/impl/device_batched_mha_bwd_xdl_cshuffle_qloop_v1.hpp

@@ -1153,11 +1153,15 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
         }
 
         // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+        if constexpr(is_same<OutputDataType, half_t>::value ||
+                     is_same<OutputDataType, bhalf_t>::value)
+        {
             if(KzRaw % 2 != 0)
             {
                 std::cout << "K_q must be a multiple of 2" << std::endl;
                 return false;
             }
+        }
 
         // Check vector load/store requirement
         const auto a_stride_lowest =

include/ck/tensor_operation/gpu/device/impl/device_batched_mha_bwd_xdl_cshuffle_qloop_v2.hpp

@@ -1190,11 +1190,15 @@ struct DeviceBatchedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
         }
 
         // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+        if constexpr(is_same<OutputDataType, half_t>::value ||
+                     is_same<OutputDataType, bhalf_t>::value)
+        {
             if(KzRaw % 2 != 0)
             {
                 std::cout << "K_q must be a multiple of 2" << std::endl;
                 return false;
             }
+        }
 
         // Check vector load/store requirement
         const auto a_stride_lowest =

include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_bwd_xdl_cshuffle_qloop_light_v1.hpp

@@ -1336,11 +1336,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V1
             }
 
             // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
                 if(KzRaw % 2 != 0)
                 {
                     std::cout << "K_q must be a multiple of 2" << std::endl;
                     return false;
                 }
+            }
 
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2

include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_bwd_xdl_cshuffle_qloop_light_v2.hpp

@@ -1408,11 +1408,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_Light_V2
             }
 
             // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
                 if(KzRaw % 2 != 0)
                 {
                     std::cout << "K_q must be a multiple of 2" << std::endl;
                     return false;
                 }
+            }
 
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2

include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_bwd_xdl_cshuffle_qloop_v1.hpp

@@ -1182,11 +1182,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V1
             }
 
             // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
                 if(KzRaw % 2 != 0)
                 {
                     std::cout << "K_q must be a multiple of 2" << std::endl;
                     return false;
                 }
+            }
 
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2

include/ck/tensor_operation/gpu/device/impl/device_grouped_mha_bwd_xdl_cshuffle_qloop_v2.hpp

@@ -1253,11 +1253,15 @@ struct DeviceGroupedMultiheadAttentionBackward_Qloop_Xdl_CShuffle_V2
             }
 
             // saving dQ data with atomic_add instruction, so KzRaw must be a multiple of 2
+            if constexpr(is_same<OutputDataType, half_t>::value ||
+                         is_same<OutputDataType, bhalf_t>::value)
+            {
                 if(KzRaw % 2 != 0)
                 {
                     std::cout << "K_q must be a multiple of 2" << std::endl;
                     return false;
                 }
+            }
 
             // Check vector load/store requirement
             const auto a_stride_lowest = ABlockTransferSrcVectorDim == 2