diff --git a/drivers/net/wireless/ath/ath9k/ath9k.h b/drivers/net/wireless/ath/ath9k/ath9k.h
index 94d887b65e6993a778c12dd3ca47d55426a013dd..1e86147831812de4f5b20ed8490dc5cd286dbac6 100644
--- a/drivers/net/wireless/ath/ath9k/ath9k.h
+++ b/drivers/net/wireless/ath/ath9k/ath9k.h
@@ -340,7 +340,8 @@ void ath_tx_aggr_stop(struct ath_softc *sc, struct ieee80211_sta *sta, u16 tid);
 void ath_tx_aggr_resume(struct ath_softc *sc, struct ieee80211_sta *sta, u16 tid);
 
 void ath_tx_aggr_wakeup(struct ath_softc *sc, struct ath_node *an);
-bool ath_tx_aggr_sleep(struct ath_softc *sc, struct ath_node *an);
+void ath_tx_aggr_sleep(struct ieee80211_sta *sta, struct ath_softc *sc,
+		       struct ath_node *an);
 
 /********/
 /* VIFs */
diff --git a/drivers/net/wireless/ath/ath9k/main.c b/drivers/net/wireless/ath/ath9k/main.c
index edaa7843bf4c7767f275e7ce1a7553f600cb7038..0ebf7321df123dcbd1f496907f2ea13047ca542b 100644
--- a/drivers/net/wireless/ath/ath9k/main.c
+++ b/drivers/net/wireless/ath/ath9k/main.c
@@ -1833,8 +1833,7 @@ static void ath9k_sta_notify(struct ieee80211_hw *hw,
 	switch (cmd) {
 	case STA_NOTIFY_SLEEP:
 		an->sleeping = true;
-		if (ath_tx_aggr_sleep(sc, an))
-			ieee80211_sta_set_tim(sta);
+		ath_tx_aggr_sleep(sta, sc, an);
 		break;
 	case STA_NOTIFY_AWAKE:
 		an->sleeping = false;
diff --git a/drivers/net/wireless/ath/ath9k/xmit.c b/drivers/net/wireless/ath/ath9k/xmit.c
index f5d4764888b910ce83a1104a9d6c3b2e74c555c9..c2bfc57958d82a2737f405167f5f4217f5266f03 100644
--- a/drivers/net/wireless/ath/ath9k/xmit.c
+++ b/drivers/net/wireless/ath/ath9k/xmit.c
@@ -542,7 +542,7 @@ static void ath_tx_complete_aggr(struct ath_softc *sc, struct ath_txq *txq,
 	/* prepend un-acked frames to the beginning of the pending frame queue */
 	if (!skb_queue_empty(&bf_pending)) {
 		if (an->sleeping)
-			ieee80211_sta_set_tim(sta);
+			ieee80211_sta_set_buffered(sta, tid->tidno, true);
 
 		spin_lock_bh(&txq->axq_lock);
 		if (clear_filter)
@@ -1153,12 +1153,13 @@ void ath_tx_aggr_stop(struct ath_softc *sc, struct ieee80211_sta *sta, u16 tid)
 	ath_tx_flush_tid(sc, txtid);
 }
 
-bool ath_tx_aggr_sleep(struct ath_softc *sc, struct ath_node *an)
+void ath_tx_aggr_sleep(struct ieee80211_sta *sta, struct ath_softc *sc,
+		       struct ath_node *an)
 {
 	struct ath_atx_tid *tid;
 	struct ath_atx_ac *ac;
 	struct ath_txq *txq;
-	bool buffered = false;
+	bool buffered;
 	int tidno;
 
 	for (tidno = 0, tid = &an->tid[tidno];
@@ -1172,8 +1173,7 @@ bool ath_tx_aggr_sleep(struct ath_softc *sc, struct ath_node *an)
 
 		spin_lock_bh(&txq->axq_lock);
 
-		if (!skb_queue_empty(&tid->buf_q))
-			buffered = true;
+		buffered = !skb_queue_empty(&tid->buf_q);
 
 		tid->sched = false;
 		list_del(&tid->list);
@@ -1184,9 +1184,9 @@ bool ath_tx_aggr_sleep(struct ath_softc *sc, struct ath_node *an)
 		}
 
 		spin_unlock_bh(&txq->axq_lock);
-	}
 
-	return buffered;
+		ieee80211_sta_set_buffered(sta, tidno, buffered);
+	}
 }
 
 void ath_tx_aggr_wakeup(struct ath_softc *sc, struct ath_node *an)
diff --git a/include/net/mac80211.h b/include/net/mac80211.h
index b5f7ada2f87bff460e9fa80c67d5100057b1b133..e66638e749c6400636f487f0c6faa83c202cb6c7 100644
--- a/include/net/mac80211.h
+++ b/include/net/mac80211.h
@@ -2361,17 +2361,35 @@ static inline int ieee80211_sta_ps_transition_ni(struct ieee80211_sta *sta,
 #define IEEE80211_TX_STATUS_HEADROOM	13
 
 /**
- * ieee80211_sta_set_tim - set the TIM bit for a sleeping station
+ * ieee80211_sta_set_buffered - inform mac80211 about driver-buffered frames
  * @sta: &struct ieee80211_sta pointer for the sleeping station
+ * @tid: the TID that has buffered frames
+ * @buffered: indicates whether or not frames are buffered for this TID
  *
  * If a driver buffers frames for a powersave station instead of passing
- * them back to mac80211 for retransmission, the station needs to be told
- * to wake up using the TIM bitmap in the beacon.
- *
- * This function sets the station's TIM bit - it will be cleared when the
- * station wakes up.
- */
-void ieee80211_sta_set_tim(struct ieee80211_sta *sta);
+ * them back to mac80211 for retransmission, the station may still need
+ * to be told that there are buffered frames via the TIM bit.
+ *
+ * This function informs mac80211 whether or not there are frames that are
+ * buffered in the driver for a given TID; mac80211 can then use this data
+ * to set the TIM bit (NOTE: This may call back into the driver's set_tim
+ * call! Beware of the locking!)
+ *
+ * If all frames are released to the station (due to PS-poll or uAPSD)
+ * then the driver needs to inform mac80211 that there no longer are
+ * frames buffered. However, when the station wakes up mac80211 assumes
+ * that all buffered frames will be transmitted and clears this data,
+ * drivers need to make sure they inform mac80211 about all buffered
+ * frames on the sleep transition (sta_notify() with %STA_NOTIFY_SLEEP).
+ *
+ * Note that technically mac80211 only needs to know this per AC, not per
+ * TID, but since driver buffering will inevitably happen per TID (since
+ * it is related to aggregation) it is easier to make mac80211 map the
+ * TID to the AC as required instead of keeping track in all drivers that
+ * use this API.
+ */
+void ieee80211_sta_set_buffered(struct ieee80211_sta *sta,
+				u8 tid, bool buffered);
 
 /**
  * ieee80211_tx_status - transmit status callback
diff --git a/net/mac80211/sta_info.c b/net/mac80211/sta_info.c
index c52e58c0a979a757574334f53061d2e610246fba..016742d4c48e0e61e5bded35947c78afef3eb53a 100644
--- a/net/mac80211/sta_info.c
+++ b/net/mac80211/sta_info.c
@@ -1117,11 +1117,15 @@ void ieee80211_sta_block_awake(struct ieee80211_hw *hw,
 }
 EXPORT_SYMBOL(ieee80211_sta_block_awake);
 
-void ieee80211_sta_set_tim(struct ieee80211_sta *pubsta)
+void ieee80211_sta_set_buffered(struct ieee80211_sta *pubsta,
+				u8 tid, bool buffered)
 {
 	struct sta_info *sta = container_of(pubsta, struct sta_info, sta);
 
+	if (!buffered)
+		return;
+
 	set_sta_flags(sta, WLAN_STA_PS_DRIVER_BUF);
 	sta_info_set_tim_bit(sta);
 }
-EXPORT_SYMBOL(ieee80211_sta_set_tim);
+EXPORT_SYMBOL(ieee80211_sta_set_buffered);