arvif->vdev_id, ret);
 }
 
+static int ath10k_mac_vif_recalc_ps_wake_threshold(struct ath10k_vif *arvif)
+{
+       struct ath10k *ar = arvif->ar;
+       u32 param;
+       u32 value;
+       int ret;
+
+       lockdep_assert_held(&arvif->ar->conf_mutex);
+
+       if (arvif->u.sta.uapsd)
+               value = WMI_STA_PS_TX_WAKE_THRESHOLD_NEVER;
+       else
+               value = WMI_STA_PS_TX_WAKE_THRESHOLD_ALWAYS;
+
+       param = WMI_STA_PS_PARAM_TX_WAKE_THRESHOLD;
+       ret = ath10k_wmi_set_sta_ps_param(ar, arvif->vdev_id, param, value);
+       if (ret) {
+               ath10k_warn(ar, "failed to submit ps wake threshold %u on vdev %i: %d\n",
+                           value, arvif->vdev_id, ret);
+               return ret;
+       }
+
+       return 0;
+}
+
+static int ath10k_mac_vif_recalc_ps_poll_count(struct ath10k_vif *arvif)
+{
+       struct ath10k *ar = arvif->ar;
+       u32 param;
+       u32 value;
+       int ret;
+
+       lockdep_assert_held(&arvif->ar->conf_mutex);
+
+       if (arvif->u.sta.uapsd)
+               value = WMI_STA_PS_PSPOLL_COUNT_UAPSD;
+       else
+               value = WMI_STA_PS_PSPOLL_COUNT_NO_MAX;
+
+       param = WMI_STA_PS_PARAM_PSPOLL_COUNT;
+       ret = ath10k_wmi_set_sta_ps_param(ar, arvif->vdev_id,
+                                         param, value);
+       if (ret) {
+               ath10k_warn(ar, "failed to submit ps poll count %u on vdev %i: %d\n",
+                           value, arvif->vdev_id, ret);
+               return ret;
+       }
+
+       return 0;
+}
+
 /*
  * Review this when mac80211 gains per-interface powersave support.
  */
                        goto err_peer_delete;
                }
 
-               param = WMI_STA_PS_PARAM_TX_WAKE_THRESHOLD;
-               value = WMI_STA_PS_TX_WAKE_THRESHOLD_ALWAYS;
-               ret = ath10k_wmi_set_sta_ps_param(ar, arvif->vdev_id,
-                                                 param, value);
+               ret = ath10k_mac_vif_recalc_ps_wake_threshold(arvif);
                if (ret) {
-                       ath10k_warn(ar, "failed to set vdev %i TX wake thresh: %d\n",
+                       ath10k_warn(ar, "failed to recalc ps wake threshold on vdev %i: %d\n",
                                    arvif->vdev_id, ret);
                        goto err_peer_delete;
                }
 
-               param = WMI_STA_PS_PARAM_PSPOLL_COUNT;
-               value = WMI_STA_PS_PSPOLL_COUNT_NO_MAX;
-               ret = ath10k_wmi_set_sta_ps_param(ar, arvif->vdev_id,
-                                                 param, value);
+               ret = ath10k_mac_vif_recalc_ps_poll_count(arvif);
                if (ret) {
-                       ath10k_warn(ar, "failed to set vdev %i PSPOLL count: %d\n",
+                       ath10k_warn(ar, "failed to recalc ps poll count on vdev %i: %d\n",
                                    arvif->vdev_id, ret);
                        goto err_peer_delete;
                }
        if (ret)
                ath10k_warn(ar, "failed to set rx wake param: %d\n", ret);
 
+       ret = ath10k_mac_vif_recalc_ps_wake_threshold(arvif);
+       if (ret) {
+               ath10k_warn(ar, "failed to recalc ps wake threshold on vdev %i: %d\n",
+                           arvif->vdev_id, ret);
+               return ret;
+       }
+
+       ret = ath10k_mac_vif_recalc_ps_poll_count(arvif);
+       if (ret) {
+               ath10k_warn(ar, "failed to recalc ps poll count on vdev %i: %d\n",
+                           arvif->vdev_id, ret);
+               return ret;
+       }
+
 exit:
        return ret;
 }