diff --git a/net/mac80211/sta_info.c b/net/mac80211/sta_info.c
index 016742d..863d59f 100644
--- a/net/mac80211/sta_info.c
+++ b/net/mac80211/sta_info.c
@@ -641,54 +641,42 @@
 	bss->tim[aid / 8] &= ~(1 << (aid % 8));
 }
 
-static void __sta_info_set_tim_bit(struct ieee80211_if_ap *bss,
-				   struct sta_info *sta)
+void sta_info_recalc_tim(struct sta_info *sta)
 {
-	BUG_ON(!bss);
-
-	__bss_tim_set(bss, sta->sta.aid);
-
-	if (sta->local->ops->set_tim) {
-		sta->local->tim_in_locked_section = true;
-		drv_set_tim(sta->local, &sta->sta, true);
-		sta->local->tim_in_locked_section = false;
-	}
-}
-
-void sta_info_set_tim_bit(struct sta_info *sta)
-{
+	struct ieee80211_local *local = sta->local;
+	struct ieee80211_if_ap *bss = sta->sdata->bss;
 	unsigned long flags;
+	bool have_data = false;
 
-	BUG_ON(!sta->sdata->bss);
+	if (WARN_ON_ONCE(!sta->sdata->bss))
+		return;
 
-	spin_lock_irqsave(&sta->local->sta_lock, flags);
-	__sta_info_set_tim_bit(sta->sdata->bss, sta);
-	spin_unlock_irqrestore(&sta->local->sta_lock, flags);
-}
+	/* No need to do anything if the driver does all */
+	if (local->hw.flags & IEEE80211_HW_AP_LINK_PS)
+		return;
 
-static void __sta_info_clear_tim_bit(struct ieee80211_if_ap *bss,
-				     struct sta_info *sta)
-{
-	BUG_ON(!bss);
+	if (sta->dead)
+		goto done;
 
-	__bss_tim_clear(bss, sta->sta.aid);
+	have_data = test_sta_flags(sta, WLAN_STA_PS_DRIVER_BUF) ||
+		    !skb_queue_empty(&sta->tx_filtered) ||
+		    !skb_queue_empty(&sta->ps_tx_buf);
 
-	if (sta->local->ops->set_tim) {
-		sta->local->tim_in_locked_section = true;
-		drv_set_tim(sta->local, &sta->sta, false);
-		sta->local->tim_in_locked_section = false;
+ done:
+	spin_lock_irqsave(&local->sta_lock, flags);
+
+	if (have_data)
+		__bss_tim_set(bss, sta->sta.aid);
+	else
+		__bss_tim_clear(bss, sta->sta.aid);
+
+	if (local->ops->set_tim) {
+		local->tim_in_locked_section = true;
+		drv_set_tim(local, &sta->sta, have_data);
+		local->tim_in_locked_section = false;
 	}
-}
 
-void sta_info_clear_tim_bit(struct sta_info *sta)
-{
-	unsigned long flags;
-
-	BUG_ON(!sta->sdata->bss);
-
-	spin_lock_irqsave(&sta->local->sta_lock, flags);
-	__sta_info_clear_tim_bit(sta->sdata->bss, sta);
-	spin_unlock_irqrestore(&sta->local->sta_lock, flags);
+	spin_unlock_irqrestore(&local->sta_lock, flags);
 }
 
 static bool sta_info_buffer_expired(struct sta_info *sta, struct sk_buff *skb)
@@ -717,6 +705,10 @@
 	unsigned long flags;
 	struct sk_buff *skb;
 
+	/* This is only necessary for stations on BSS interfaces */
+	if (!sta->sdata->bss)
+		return false;
+
 	for (;;) {
 		spin_lock_irqsave(&sta->ps_tx_buf.lock, flags);
 		skb = skb_peek(&sta->ps_tx_buf);
@@ -736,9 +728,9 @@
 #endif
 		dev_kfree_skb(skb);
 
-		if (skb_queue_empty(&sta->ps_tx_buf) &&
-		    !test_sta_flags(sta, WLAN_STA_PS_DRIVER_BUF))
-			sta_info_clear_tim_bit(sta);
+		/* if the queue is now empty recalc TIM bit */
+		if (skb_queue_empty(&sta->ps_tx_buf))
+			sta_info_recalc_tim(sta);
 	}
 
 	return !skb_queue_empty(&sta->ps_tx_buf);
@@ -748,7 +740,6 @@
 {
 	struct ieee80211_local *local;
 	struct ieee80211_sub_if_data *sdata;
-	struct sk_buff *skb;
 	unsigned long flags;
 	int ret, i;
 
@@ -792,7 +783,7 @@
 		BUG_ON(!sdata->bss);
 
 		atomic_dec(&sdata->bss->num_sta_ps);
-		sta_info_clear_tim_bit(sta);
+		sta_info_recalc_tim(sta);
 	}
 
 	local->num_sta--;
@@ -818,6 +809,10 @@
 	 */
 	synchronize_rcu();
 
+	local->total_ps_buffered -= skb_queue_len(&sta->ps_tx_buf);
+	__skb_queue_purge(&sta->ps_tx_buf);
+	__skb_queue_purge(&sta->tx_filtered);
+
 #ifdef CONFIG_MAC80211_MESH
 	if (ieee80211_vif_is_mesh(&sdata->vif))
 		mesh_accept_plinks_update(sdata);
@@ -840,14 +835,6 @@
 	}
 #endif
 
-	while ((skb = skb_dequeue(&sta->ps_tx_buf)) != NULL) {
-		local->total_ps_buffered--;
-		dev_kfree_skb_any(skb);
-	}
-
-	while ((skb = skb_dequeue(&sta->tx_filtered)) != NULL)
-		dev_kfree_skb_any(skb);
-
 	__sta_info_free(local, sta);
 
 	return 0;
@@ -1027,9 +1014,6 @@
 	if (!(local->hw.flags & IEEE80211_HW_AP_LINK_PS))
 		drv_sta_notify(local, sdata, STA_NOTIFY_AWAKE, &sta->sta);
 
-	if (!skb_queue_empty(&sta->ps_tx_buf))
-		sta_info_clear_tim_bit(sta);
-
 	/* Send all buffered frames to the station */
 	sent = ieee80211_add_pending_skbs(local, &sta->tx_filtered);
 	buffered = ieee80211_add_pending_skbs_fn(local, &sta->ps_tx_buf,
@@ -1037,6 +1021,8 @@
 	sent += buffered;
 	local->total_ps_buffered -= buffered;
 
+	sta_info_recalc_tim(sta);
+
 #ifdef CONFIG_MAC80211_VERBOSE_PS_DEBUG
 	printk(KERN_DEBUG "%s: STA %pM aid %d sending %d filtered/%d PS frames "
 	       "since STA not sleeping anymore\n", sdata->name,
@@ -1086,8 +1072,7 @@
 
 		ieee80211_add_pending_skb(local, skb);
 
-		if (no_pending_pkts)
-			sta_info_clear_tim_bit(sta);
+		sta_info_recalc_tim(sta);
 #ifdef CONFIG_MAC80211_VERBOSE_PS_DEBUG
 	} else {
 		/*
@@ -1126,6 +1111,6 @@
 		return;
 
 	set_sta_flags(sta, WLAN_STA_PS_DRIVER_BUF);
-	sta_info_set_tim_bit(sta);
+	sta_info_recalc_tim(sta);
 }
 EXPORT_SYMBOL(ieee80211_sta_set_buffered);
diff --git a/net/mac80211/sta_info.h b/net/mac80211/sta_info.h
index c10e2e8..c9ffb7c 100644
--- a/net/mac80211/sta_info.h
+++ b/net/mac80211/sta_info.h
@@ -528,8 +528,7 @@
 int sta_info_destroy_addr_bss(struct ieee80211_sub_if_data *sdata,
 			      const u8 *addr);
 
-void sta_info_set_tim_bit(struct sta_info *sta);
-void sta_info_clear_tim_bit(struct sta_info *sta);
+void sta_info_recalc_tim(struct sta_info *sta);
 
 void sta_info_init(struct ieee80211_local *local);
 void sta_info_stop(struct ieee80211_local *local);
diff --git a/net/mac80211/status.c b/net/mac80211/status.c
index d50358c..8354dcb 100644
--- a/net/mac80211/status.c
+++ b/net/mac80211/status.c
@@ -106,6 +106,7 @@
 	if (test_sta_flags(sta, WLAN_STA_PS_STA) &&
 	    skb_queue_len(&sta->tx_filtered) < STA_MAX_TX_BUFFER) {
 		skb_queue_tail(&sta->tx_filtered, skb);
+		sta_info_recalc_tim(sta);
 		return;
 	}
 
diff --git a/net/mac80211/tx.c b/net/mac80211/tx.c
index 0ca1688..d675490 100644
--- a/net/mac80211/tx.c
+++ b/net/mac80211/tx.c
@@ -469,15 +469,6 @@
 		} else
 			tx->local->total_ps_buffered++;
 
-		/*
-		 * Queue frame to be sent after STA wakes up/polls,
-		 * but don't set the TIM bit if the driver is blocking
-		 * wakeup or poll response transmissions anyway.
-		 */
-		if (skb_queue_empty(&sta->ps_tx_buf) &&
-		    !(staflags & WLAN_STA_PS_DRIVER))
-			sta_info_set_tim_bit(sta);
-
 		info->control.jiffies = jiffies;
 		info->control.vif = &tx->sdata->vif;
 		info->flags |= IEEE80211_TX_INTFL_NEED_TXPROCESSING;
@@ -488,6 +479,12 @@
 				  round_jiffies(jiffies +
 						STA_INFO_CLEANUP_INTERVAL));
 
+		/*
+		 * We queued up some frames, so the TIM bit might
+		 * need to be set, recalculate it.
+		 */
+		sta_info_recalc_tim(sta);
+
 		return TX_QUEUED;
 	}
 #ifdef CONFIG_MAC80211_VERBOSE_PS_DEBUG