#define SKB_CB_ATHBUF(__skb)   (*((struct ath_buf **)__skb->cb))
 
+static inline bool ath9k_check_auto_sleep(struct ath_softc *sc)
+{
+       return sc->ps_enabled &&
+              (sc->sc_ah->caps.hw_caps & ATH9K_HW_CAP_AUTOSLEEP);
+}
+
 static struct ieee80211_hw * ath_get_virt_hw(struct ath_softc *sc,
                                             struct ieee80211_hdr *hdr)
 {
        hdr = (struct ieee80211_hdr *)skb->data;
 
        /* Process Beacon and CAB receive in PS state */
-       if ((sc->ps_flags & PS_WAIT_FOR_BEACON) &&
-           ieee80211_is_beacon(hdr->frame_control))
+       if (((sc->ps_flags & PS_WAIT_FOR_BEACON) || ath9k_check_auto_sleep(sc))
+           && ieee80211_is_beacon(hdr->frame_control))
                ath_rx_ps_beacon(sc, skb);
        else if ((sc->ps_flags & PS_WAIT_FOR_CAB) &&
                 (ieee80211_is_data(hdr->frame_control) ||
                        sc->rx.rxotherant = 0;
                }
 
-               if (unlikely(sc->ps_flags & (PS_WAIT_FOR_BEACON |
-                                            PS_WAIT_FOR_CAB |
-                                            PS_WAIT_FOR_PSPOLL_DATA)))
+               if (unlikely(ath9k_check_auto_sleep(sc) ||
+                            (sc->ps_flags & (PS_WAIT_FOR_BEACON |
+                                             PS_WAIT_FOR_CAB |
+                                             PS_WAIT_FOR_PSPOLL_DATA))))
                        ath_rx_ps(sc, skb);
 
                ath_rx_send_to_mac80211(hw, sc, skb, rxs);