bool more_data)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
-       struct iwl_mvm_sta *mvmsta = (void *)sta->drv_priv;
 
        /* TODO: how do we tell the fw to send frames for a specific TID */
 
         * The fw will send EOSP notification when the last frame will be
         * transmitted.
         */
-       iwl_mvm_sta_modify_sleep_tx_count(mvm, mvmsta->sta_id, reason,
-                                         num_frames);
+       iwl_mvm_sta_modify_sleep_tx_count(mvm, sta, reason, num_frames);
 }
 
 static void iwl_mvm_mac_sta_notify(struct ieee80211_hw *hw,
        case STA_NOTIFY_AWAKE:
                if (WARN_ON(mvmsta->sta_id == IWL_INVALID_STATION))
                        break;
-               iwl_mvm_sta_modify_ps_wake(mvm, mvmsta->sta_id);
+               iwl_mvm_sta_modify_ps_wake(mvm, sta);
                break;
        default:
                break;
 
        rcu_read_unlock();
 }
 
-void iwl_mvm_sta_modify_ps_wake(struct iwl_mvm *mvm, int sta_id)
+void iwl_mvm_sta_modify_ps_wake(struct iwl_mvm *mvm,
+                               struct ieee80211_sta *sta)
 {
+       struct iwl_mvm_sta *mvmsta = (void *)sta->drv_priv;
        struct iwl_mvm_add_sta_cmd cmd = {
                .add_modify = STA_MODE_MODIFY,
-               .sta_id = sta_id,
+               .sta_id = mvmsta->sta_id,
                .modify_mask = STA_MODIFY_SLEEPING_STA_TX_COUNT,
                .sleep_state_flags = cpu_to_le16(STA_SLEEP_STATE_AWAKE),
+               .mac_id_n_color = cpu_to_le32(mvmsta->mac_id_n_color),
        };
        int ret;
 
                IWL_ERR(mvm, "Failed to send ADD_STA command (%d)\n", ret);
 }
 
-void iwl_mvm_sta_modify_sleep_tx_count(struct iwl_mvm *mvm, int sta_id,
+void iwl_mvm_sta_modify_sleep_tx_count(struct iwl_mvm *mvm,
+                                      struct ieee80211_sta *sta,
                                       enum ieee80211_frame_release_type reason,
                                       u16 cnt)
 {
        u16 sleep_state_flags =
                (reason == IEEE80211_FRAME_RELEASE_UAPSD) ?
                        STA_SLEEP_STATE_UAPSD : STA_SLEEP_STATE_PS_POLL;
+       struct iwl_mvm_sta *mvmsta = (void *)sta->drv_priv;
        struct iwl_mvm_add_sta_cmd cmd = {
                .add_modify = STA_MODE_MODIFY,
-               .sta_id = sta_id,
+               .sta_id = mvmsta->sta_id,
                .modify_mask = STA_MODIFY_SLEEPING_STA_TX_COUNT,
                .sleep_tx_count = cpu_to_le16(cnt),
+               .mac_id_n_color = cpu_to_le32(mvmsta->mac_id_n_color),
                /*
                 * Same modify mask for sleep_tx_count and sleep_state_flags so
                 * we must set the sleep_state_flags too.
 
                          struct iwl_mvm_int_sta *bsta);
 int iwl_mvm_rm_bcast_sta(struct iwl_mvm *mvm, struct iwl_mvm_int_sta *bsta);
 void iwl_mvm_sta_drained_wk(struct work_struct *wk);
-void iwl_mvm_sta_modify_ps_wake(struct iwl_mvm *mvm, int sta_id);
-void iwl_mvm_sta_modify_sleep_tx_count(struct iwl_mvm *mvm, int sta_id,
+void iwl_mvm_sta_modify_ps_wake(struct iwl_mvm *mvm,
+                               struct ieee80211_sta *sta);
+void iwl_mvm_sta_modify_sleep_tx_count(struct iwl_mvm *mvm,
+                                      struct ieee80211_sta *sta,
                                       enum ieee80211_frame_release_type reason,
                                       u16 cnt);
 int iwl_mvm_drain_sta(struct iwl_mvm *mvm, struct iwl_mvm_sta *mvmsta,