diff --git a/mt76x02_txrx.c b/mt76x02_txrx.c
index 301b4318..b27c389f 100644
--- a/mt76x02_txrx.c
+++ b/mt76x02_txrx.c
@@ -141,11 +141,15 @@ int mt76x02_tx_prepare_skb(struct mt76_dev *mdev, void *txwi_ptr,
 {
 	struct mt76x02_dev *dev = container_of(mdev, struct mt76x02_dev, mt76);
 	struct ieee80211_hdr *hdr = (struct ieee80211_hdr *)tx_info->skb->data;
+	struct ieee80211_tx_info *info = IEEE80211_SKB_CB(tx_info->skb);
 	struct mt76x02_txwi *txwi = txwi_ptr;
-	bool ampdu = IEEE80211_SKB_CB(tx_info->skb)->flags & IEEE80211_TX_CTL_AMPDU;
+	bool ampdu = info->flags & IEEE80211_TX_CTL_AMPDU;
 	int hdrlen, len, pid, qsel = MT_QSEL_EDCA;
 
-	if (qid == MT_TXQ_PSD && wcid && wcid->idx < 128)
+	if (wcid && wcid->idx < MT76x02_N_WCIDS &&
+	    ((info->flags & (IEEE80211_TX_CTL_NO_PS_BUFFER |
+			     IEEE80211_TX_CTL_CLEAR_PS_FILT)) ||
+	     (info->control.flags & IEEE80211_TX_CTRL_PS_RESPONSE)))
 		mt76x02_mac_wcid_set_drop(dev, wcid->idx, false);
 
 	hdrlen = ieee80211_hdrlen(hdr->frame_control);