diff --git a/dma.c b/dma.c
index e9ef2858..8d91a8ef 100644
--- a/dma.c
+++ b/dma.c
@@ -100,6 +100,7 @@ mt76_dma_cleanup(struct mt76_dev *dev, struct mt76_queue *q, bool flush,
 	struct mt76_queue_entry entry;
 	int last;
 
+	spin_lock_bh(&q->lock);
 	if (flush)
 		last = -1;
 	else
@@ -120,6 +121,9 @@ mt76_dma_cleanup(struct mt76_dev *dev, struct mt76_queue *q, bool flush,
 		if (q->tail == last)
 		    last = ioread32(&q->regs->dma_idx);
 	}
+	if (!flush)
+		mt76_txq_schedule(dev, q);
+	spin_unlock_bh(&q->lock);
 }
 
 static void *
diff --git a/mt76x2_dma.c b/mt76x2_dma.c
index f52d5dee..a7ec6c8b 100644
--- a/mt76x2_dma.c
+++ b/mt76x2_dma.c
@@ -117,11 +117,7 @@ mt76_tx_cleanup(struct mt76x2_dev *dev, struct mt76_queue *q, bool flush)
 	if (!q->desc)
 		return;
 
-	spin_lock_bh(&q->lock);
 	mt76_queue_cleanup(dev, q, flush, mt76x2_tx_cleanup_entry);
-	if (!flush)
-		mt76_txq_schedule(&dev->mt76, q);
-	spin_unlock_bh(&q->lock);
 }
 
 static int