static u32 mlx4_en_free_tx_desc(struct mlx4_en_priv *priv,
                                struct mlx4_en_tx_ring *ring,
-                               int index, u8 owner, u64 timestamp)
+                               int index, u8 owner, u64 timestamp,
+                               int napi_mode)
 {
        struct mlx4_en_tx_info *tx_info = &ring->tx_info[index];
        struct mlx4_en_tx_desc *tx_desc = ring->buf + index * TXBB_SIZE;
                        }
                }
        }
-       dev_consume_skb_any(skb);
+       napi_consume_skb(skb, napi_mode);
+
        return tx_info->nr_txbb;
 }
 
        while (ring->cons != ring->prod) {
                ring->last_nr_txbb = mlx4_en_free_tx_desc(priv, ring,
                                                ring->cons & ring->size_mask,
-                                               !!(ring->cons & ring->size), 0);
+                                               !!(ring->cons & ring->size), 0,
+                                               0 /* Non-NAPI caller */);
                ring->cons += ring->last_nr_txbb;
                cnt++;
        }
 }
 
 static bool mlx4_en_process_tx_cq(struct net_device *dev,
-                                struct mlx4_en_cq *cq)
+                                 struct mlx4_en_cq *cq, int napi_budget)
 {
        struct mlx4_en_priv *priv = netdev_priv(dev);
        struct mlx4_cq *mcq = &cq->mcq;
                        last_nr_txbb = mlx4_en_free_tx_desc(
                                        priv, ring, ring_index,
                                        !!((ring_cons + txbbs_skipped) &
-                                       ring->size), timestamp);
+                                       ring->size), timestamp, napi_budget);
 
                        mlx4_en_stamp_wqe(priv, ring, stamp_index,
                                          !!((ring_cons + txbbs_stamp) &
        struct mlx4_en_priv *priv = netdev_priv(dev);
        int clean_complete;
 
-       clean_complete = mlx4_en_process_tx_cq(dev, cq);
+       clean_complete = mlx4_en_process_tx_cq(dev, cq, budget);
        if (!clean_complete)
                return budget;