* in blk_mq_dispatch_rq_list().
                 */
                list_add(&rq->queuelist, &rq_list);
-       } while (blk_mq_dispatch_rq_list(rq->mq_hctx, &rq_list, true));
+       } while (blk_mq_dispatch_rq_list(rq->mq_hctx, &rq_list, 1));
 
        return ret;
 }
                /* round robin for fair dispatch */
                ctx = blk_mq_next_ctx(hctx, rq->mq_ctx);
 
-       } while (blk_mq_dispatch_rq_list(rq->mq_hctx, &rq_list, true));
+       } while (blk_mq_dispatch_rq_list(rq->mq_hctx, &rq_list, 1));
 
        WRITE_ONCE(hctx->dispatch_from, ctx);
        return ret;
         */
        if (!list_empty(&rq_list)) {
                blk_mq_sched_mark_restart_hctx(hctx);
-               if (blk_mq_dispatch_rq_list(hctx, &rq_list, false)) {
+               if (blk_mq_dispatch_rq_list(hctx, &rq_list, 0)) {
                        if (has_sched_dispatch)
                                ret = blk_mq_do_dispatch_sched(hctx);
                        else
                ret = blk_mq_do_dispatch_ctx(hctx);
        } else {
                blk_mq_flush_busy_ctxs(hctx, &rq_list);
-               blk_mq_dispatch_rq_list(hctx, &rq_list, false);
+               blk_mq_dispatch_rq_list(hctx, &rq_list, 0);
        }
 
        return ret;
 
                 * we'll re-run it below.
                 */
                if (!blk_mq_mark_tag_wait(hctx, rq)) {
-                       blk_mq_put_dispatch_budget(rq->q);
+                       /*
+                        * All budgets not got from this function will be put
+                        * together during handling partial dispatch
+                        */
+                       if (need_budget)
+                               blk_mq_put_dispatch_budget(rq->q);
                        return PREP_DISPATCH_NO_TAG;
                }
        }
        return PREP_DISPATCH_OK;
 }
 
+/* release all allocated budgets before calling to blk_mq_dispatch_rq_list */
+static void blk_mq_release_budgets(struct request_queue *q,
+               unsigned int nr_budgets)
+{
+       int i;
+
+       for (i = 0; i < nr_budgets; i++)
+               blk_mq_put_dispatch_budget(q);
+}
+
 /*
  * Returns true if we did some work AND can potentially do more.
  */
 bool blk_mq_dispatch_rq_list(struct blk_mq_hw_ctx *hctx, struct list_head *list,
-                            bool got_budget)
+                            unsigned int nr_budgets)
 {
        enum prep_dispatch prep;
        struct request_queue *q = hctx->queue;
        if (list_empty(list))
                return false;
 
-       WARN_ON(!list_is_singular(list) && got_budget);
+       WARN_ON(!list_is_singular(list) && nr_budgets);
 
        /*
         * Now process all the entries, sending them to the driver.
                rq = list_first_entry(list, struct request, queuelist);
 
                WARN_ON_ONCE(hctx != rq->mq_hctx);
-               prep = blk_mq_prep_dispatch_rq(rq, !got_budget);
+               prep = blk_mq_prep_dispatch_rq(rq, !nr_budgets);
                if (prep != PREP_DISPATCH_OK)
                        break;
 
                        bd.last = !blk_mq_get_driver_tag(nxt);
                }
 
+               /*
+                * once the request is queued to lld, no need to cover the
+                * budget any more
+                */
+               if (nr_budgets)
+                       nr_budgets--;
                ret = q->mq_ops->queue_rq(hctx, &bd);
                if (ret == BLK_STS_RESOURCE || ret == BLK_STS_DEV_RESOURCE) {
                        blk_mq_handle_dev_resource(rq, list);
                         (hctx->flags & BLK_MQ_F_TAG_SHARED);
                bool no_budget_avail = prep == PREP_DISPATCH_NO_BUDGET;
 
+               blk_mq_release_budgets(q, nr_budgets);
+
                /*
                 * If we didn't flush the entire list, we could have told
                 * the driver there was more coming, but that turned out to
 
 void blk_mq_exit_queue(struct request_queue *q);
 int blk_mq_update_nr_requests(struct request_queue *q, unsigned int nr);
 void blk_mq_wake_waiters(struct request_queue *q);
-bool blk_mq_dispatch_rq_list(struct blk_mq_hw_ctx *hctx, struct list_head *, bool);
+bool blk_mq_dispatch_rq_list(struct blk_mq_hw_ctx *hctx, struct list_head *,
+                            unsigned int);
 void blk_mq_add_to_requeue_list(struct request *rq, bool at_head,
                                bool kick_requeue_list);
 void blk_mq_flush_busy_ctxs(struct blk_mq_hw_ctx *hctx, struct list_head *list);