return 0;
 }
 
+static void iwl_mac_sta_notify(struct ieee80211_hw *hw,
+                              struct ieee80211_vif *vif,
+                              enum sta_notify_cmd cmd,
+                              struct ieee80211_sta *sta)
+{
+       struct iwl_priv *priv = hw->priv;
+       struct iwl_station_priv *sta_priv = (void *)sta->drv_priv;
+       int sta_id;
+
+       /*
+        * TODO: We really should use this callback to
+        *       actually maintain the station table in
+        *       the device.
+        */
+
+       switch (cmd) {
+       case STA_NOTIFY_ADD:
+               atomic_set(&sta_priv->pending_frames, 0);
+               if (vif->type == NL80211_IFTYPE_AP)
+                       sta_priv->client = true;
+               break;
+       case STA_NOTIFY_SLEEP:
+               WARN_ON(!sta_priv->client);
+               sta_priv->asleep = true;
+               if (atomic_read(&sta_priv->pending_frames) > 0)
+                       ieee80211_sta_block_awake(hw, sta, true);
+               break;
+       case STA_NOTIFY_AWAKE:
+               WARN_ON(!sta_priv->client);
+               sta_priv->asleep = false;
+               sta_id = iwl_find_station(priv, sta->addr);
+               if (sta_id != IWL_INVALID_STATION)
+                       iwl_sta_modify_ps_wake(priv, sta_id);
+               break;
+       default:
+               break;
+       }
+}
+
 /*****************************************************************************
  *
  * sysfs attributes
        .reset_tsf = iwl_mac_reset_tsf,
        .bss_info_changed = iwl_bss_info_changed,
        .ampdu_action = iwl_mac_ampdu_action,
-       .hw_scan = iwl_mac_hw_scan
+       .hw_scan = iwl_mac_hw_scan,
+       .sta_notify = iwl_mac_sta_notify,
 };
 
 static int iwl_pci_probe(struct pci_dev *pdev, const struct pci_device_id *ent)
 
 #define        STA_MODIFY_TX_RATE_MSK          0x04
 #define STA_MODIFY_ADDBA_TID_MSK       0x08
 #define STA_MODIFY_DELBA_TID_MSK       0x10
+#define STA_MODIFY_SLEEP_TX_COUNT_MSK  0x20
 
 /* Receiver address (actually, Rx station's index into station table),
  * combined with Traffic ID (QOS priority), in format used by Tx Scheduler */
 
                        pos += scnprintf(buf + pos, bufsz - pos,
                                        "flags: 0x%x\n",
                                        station->sta.station_flags_msk);
-                       pos += scnprintf(buf + pos, bufsz - pos,
-                                       "ps_status: %u\n", station->ps_status);
                        pos += scnprintf(buf + pos, bufsz - pos, "tid data:\n");
                        pos += scnprintf(buf + pos, bufsz - pos,
                                        "seq_num\t\ttxq_id");
 
        struct iwl_qosparam_cmd def_qos_parm;
 };
 
-#define STA_PS_STATUS_WAKE             0
-#define STA_PS_STATUS_SLEEP            1
-
 
 struct iwl3945_station_entry {
        struct iwl3945_addsta_cmd sta;
        struct iwl_tid_data tid[MAX_TID_COUNT];
        u8 used;
-       u8 ps_status;
        struct iwl_hw_key keyinfo;
 };
 
        struct iwl_addsta_cmd sta;
        struct iwl_tid_data tid[MAX_TID_COUNT];
        u8 used;
-       u8 ps_status;
        struct iwl_hw_key keyinfo;
 };
 
  * When mac80211 creates a station it reserves some space (hw->sta_data_size)
  * in the structure for use by driver. This structure is places in that
  * space.
- *
- * At the moment use it for the station's rate scaling information.
  */
 struct iwl_station_priv {
        struct iwl_lq_sta lq_sta;
+       atomic_t pending_frames;
+       bool client;
+       bool asleep;
 };
 
 /* one for each uCode image (inst/data, boot/init/runtime) */
 
        struct iwl4965_rx_mpdu_res_start *amsdu;
        u32 len;
        u32 ampdu_status;
-       u16 fc;
        u32 rate_n_flags;
 
        /**
                priv->last_tsf = le64_to_cpu(phy_res->timestamp);
        }
 
-       fc = le16_to_cpu(header->frame_control);
-       switch (fc & IEEE80211_FCTL_FTYPE) {
-       case IEEE80211_FTYPE_MGMT:
-       case IEEE80211_FTYPE_DATA:
-               if (priv->iw_mode == NL80211_IFTYPE_AP)
-                       iwl_update_ps_mode(priv, fc  & IEEE80211_FCTL_PM,
-                                               header->addr2);
-               /* fall through */
-       default:
-               iwl_pass_packet_to_mac80211(priv, header, len, ampdu_status,
-                               rxb, &rx_status);
-               break;
-
-       }
+       iwl_pass_packet_to_mac80211(priv, header, len, ampdu_status,
+                                   rxb, &rx_status);
 }
 EXPORT_SYMBOL(iwl_rx_reply_rx);
 
 
 }
 EXPORT_SYMBOL(iwl_sta_rx_agg_stop);
 
-static void iwl_sta_modify_ps_wake(struct iwl_priv *priv, int sta_id)
+void iwl_sta_modify_ps_wake(struct iwl_priv *priv, int sta_id)
 {
        unsigned long flags;
 
        priv->stations[sta_id].sta.station_flags &= ~STA_FLG_PWR_SAVE_MSK;
        priv->stations[sta_id].sta.station_flags_msk = STA_FLG_PWR_SAVE_MSK;
        priv->stations[sta_id].sta.sta.modify_mask = 0;
+       priv->stations[sta_id].sta.sleep_tx_count = 0;
        priv->stations[sta_id].sta.mode = STA_CONTROL_MODIFY_MSK;
        spin_unlock_irqrestore(&priv->sta_lock, flags);
 
        iwl_send_add_sta(priv, &priv->stations[sta_id].sta, CMD_ASYNC);
 }
+EXPORT_SYMBOL(iwl_sta_modify_ps_wake);
 
-void iwl_update_ps_mode(struct iwl_priv *priv, u16 ps_bit, u8 *addr)
+void iwl_sta_modify_sleep_tx_count(struct iwl_priv *priv, int sta_id, int cnt)
 {
-       /* FIXME: need locking over ps_status ??? */
-       u8 sta_id = iwl_find_station(priv, addr);
+       unsigned long flags;
 
-       if (sta_id != IWL_INVALID_STATION) {
-               u8 sta_awake = priv->stations[sta_id].
-                               ps_status == STA_PS_STATUS_WAKE;
+       spin_lock_irqsave(&priv->sta_lock, flags);
+       priv->stations[sta_id].sta.station_flags |= STA_FLG_PWR_SAVE_MSK;
+       priv->stations[sta_id].sta.station_flags_msk = STA_FLG_PWR_SAVE_MSK;
+       priv->stations[sta_id].sta.sta.modify_mask =
+                                       STA_MODIFY_SLEEP_TX_COUNT_MSK;
+       priv->stations[sta_id].sta.sleep_tx_count = cpu_to_le16(cnt);
+       priv->stations[sta_id].sta.mode = STA_CONTROL_MODIFY_MSK;
+       spin_unlock_irqrestore(&priv->sta_lock, flags);
 
-               if (sta_awake && ps_bit)
-                       priv->stations[sta_id].ps_status = STA_PS_STATUS_SLEEP;
-               else if (!sta_awake && !ps_bit) {
-                       iwl_sta_modify_ps_wake(priv, sta_id);
-                       priv->stations[sta_id].ps_status = STA_PS_STATUS_WAKE;
-               }
-       }
+       iwl_send_add_sta(priv, &priv->stations[sta_id].sta, CMD_ASYNC);
 }
-
 
 int iwl_sta_rx_agg_start(struct iwl_priv *priv,
                         const u8 *addr, int tid, u16 ssn);
 int iwl_sta_rx_agg_stop(struct iwl_priv *priv, const u8 *addr, int tid);
-void iwl_update_ps_mode(struct iwl_priv *priv, u16 ps_bit, u8 *addr);
+void iwl_sta_modify_ps_wake(struct iwl_priv *priv, int sta_id);
+void iwl_sta_modify_sleep_tx_count(struct iwl_priv *priv, int sta_id, int cnt);
 #endif /* __iwl_sta_h__ */
 
 {
        struct ieee80211_hdr *hdr = (struct ieee80211_hdr *)skb->data;
        struct ieee80211_tx_info *info = IEEE80211_SKB_CB(skb);
+       struct ieee80211_sta *sta = info->control.sta;
+       struct iwl_station_priv *sta_priv = NULL;
        struct iwl_tx_queue *txq;
        struct iwl_queue *q;
        struct iwl_device_cmd *out_cmd;
 
        IWL_DEBUG_TX(priv, "station Id %d\n", sta_id);
 
+       if (sta)
+               sta_priv = (void *)sta->drv_priv;
+
+       if (sta_priv && sta_id != priv->hw_params.bcast_sta_id &&
+           sta_priv->asleep) {
+               WARN_ON(!(info->flags & IEEE80211_TX_CTL_PSPOLL_RESPONSE));
+               /*
+                * This sends an asynchronous command to the device,
+                * but we can rely on it being processed before the
+                * next frame is processed -- and the next frame to
+                * this station is the one that will consume this
+                * counter.
+                * For now set the counter to just 1 since we do not
+                * support uAPSD yet.
+                */
+               iwl_sta_modify_sleep_tx_count(priv, sta_id, 1);
+       }
+
        txq_id = skb_get_queue_mapping(skb);
        if (ieee80211_is_data_qos(fc)) {
                qc = ieee80211_get_qos_ctl(hdr);
        ret = iwl_txq_update_write_ptr(priv, txq);
        spin_unlock_irqrestore(&priv->lock, flags);
 
+       /*
+        * At this point the frame is "transmitted" successfully
+        * and we will get a TX status notification eventually,
+        * regardless of the value of ret. "ret" only indicates
+        * whether or not we should update the write pointer.
+        */
+
+       /* avoid atomic ops if it isn't an associated client */
+       if (sta_priv && sta_priv->client)
+               atomic_inc(&sta_priv->pending_frames);
+
        if (ret)
                return ret;
 
        return ret ? ret : idx;
 }
 
+static void iwl_tx_status(struct iwl_priv *priv, struct sk_buff *skb)
+{
+       struct ieee80211_hdr *hdr = (struct ieee80211_hdr *) skb->data;
+       struct ieee80211_sta *sta;
+       struct iwl_station_priv *sta_priv;
+
+       sta = ieee80211_find_sta(priv->vif, hdr->addr1);
+       if (sta) {
+               sta_priv = (void *)sta->drv_priv;
+               /* avoid atomic ops if this isn't a client */
+               if (sta_priv->client &&
+                   atomic_dec_return(&sta_priv->pending_frames) == 0)
+                       ieee80211_sta_block_awake(priv->hw, sta, false);
+       }
+
+       ieee80211_tx_status_irqsafe(priv->hw, skb);
+}
+
 int iwl_tx_queue_reclaim(struct iwl_priv *priv, int txq_id, int index)
 {
        struct iwl_tx_queue *txq = &priv->txq[txq_id];
             q->read_ptr = iwl_queue_inc_wrap(q->read_ptr, q->n_bd)) {
 
                tx_info = &txq->txb[txq->q.read_ptr];
-               ieee80211_tx_status_irqsafe(priv->hw, tx_info->skb[0]);
+               iwl_tx_status(priv, tx_info->skb[0]);
                tx_info->skb[0] = NULL;
 
                if (priv->cfg->ops->lib->txq_inval_byte_cnt_tbl)