CHAPTER 05

Choose the next batch

A free block answers whether a request can own more memory. It does not answer whether that request should run next. One user has submitted a long document; another wants a short answer; a third is already receiving tokens. They compete for computation as well as storage.

5.1 A smaller chunk is an opportunity to choose

Chapter 1’s batching experiment gave each active request one token per round. Now give the scheduler prompt work and arrivals. Request A arrives at time 0 with 12 prompt positions. B and C arrive at time 3 with two each. A step processes one request’s prompt, or part of it. Give a step with $q$ positions the assigned duration $1+q$: one unit of overhead plus one per position. The first output arrives when the prompt finishes. This experiment isolates prefill; it does not schedule subsequent decode tokens.

python3 labs/11_chunks.py

Compare three rules. First, finish each prompt in arrival order, also called first in, first out (FIFO). Second, allow at most four positions per step but put an unfinished request back at the front. Third, use the same four-position limit and put it at the back, behind arrivals seen during that step. Predict which rule helps B before reading the output.

Whole prompts, FIFO (budget 12)
   0->13: A q=12 waiting=['B', 'C']
  13->16: B q= 2 waiting=['C']
  16->19: C q= 2 waiting=[]
  first-token delays: A=13, B=13, C=16
  work=16 positions; steps=3; finish=19

With whole prompts, A occupies time 0–13. B finishes at 16, so its first-token delay is $16-3=13$. Now cut A into chunks, keeping its continuation at the front:

Chunks, continuation at front (budget 4)
   0-> 5: A q= 4 waiting=['A', 'B', 'C']
   5->10: A q= 4 waiting=['A', 'B', 'C']
  10->15: A q= 4 waiting=['B', 'C']
  15->18: B q= 2 waiting=['C']
  18->21: C q= 2 waiting=[]
  first-token delays: A=15, B=15, C=18
  work=16 positions; steps=5; finish=21

The chunks create boundaries at 5 and 10, but the scheduler chooses A again at both. B’s delay rises to 15 because A now pays the overhead three times. Moving the continuation to the back changes who runs at those boundaries:

Chunks, continuation at back (budget 4)
   0-> 5: A q= 4 waiting=['B', 'C', 'A']
   5-> 8: B q= 2 waiting=['C', 'A']
   8->11: C q= 2 waiting=['A']
  11->16: A q= 4 waiting=['A']
  16->21: A q= 4 waiting=[]
  first-token delays: A=21, B=5, C=8
  work=16 positions; steps=5; finish=21

The third rule uses those boundaries to change requests. At time 5 the queue is B, C, A. B finishes at 8, five units after its arrival. C finishes at 11, eight units after its arrival. A pays for their earlier answers: its own first token moves from time 13 to 21. The queue operation responsible for the difference is small:

38        admit()  # Arrivals during this chunk precede a rotated continuation.
39        if remaining[name]:
40            if rotate:
41                waiting.append(name)
42            else:
43                waiting.appendleft(name)
44        else:
45            first[name] = now

All three runs process 16 positions. The chunked runs take five steps instead of three, so their total duration is 21 rather than 19. Smaller chunks have changed latency and overhead without reducing the useful work. Chapter 2’s distinction between work and elapsed time applies here too. Figure 5.1 draws the three runs.

Chunk size and queue order are separate decisions. A short step creates a point at which other work could run. Whether it does depends on what the scheduler chooses there. The rotation rule is a teaching alternative; we will now inspect where DeepInfer actually puts an unfinished prompt.

5.2 Follow a prompt through the scheduler

In chapter 1, LLMEngine.step called the C++ scheduler and passed its metadata to the worker. Open that scheduler box. It makes three decisions in sequence: choose a kind of step, choose requests that fit, then describe the positions that the worker will evaluate. The block and state managers from chapter 4 participate in the second decision.

The queue names have a useful historical comparison. vLLM v0.6.3 keeps waiting, running and swapped queues of sequence groups (vllm/core/scheduler.py, lines 340–348). Its preemption path can either discard cached progress for recomputation or swap blocks out for later recovery [13]. DeepInfer’s C++ scheduler keeps waiting, running and prompted. In this version, only fork_seq adds requests to prompted (scheduler.hpp, line 1118); lines 581–618 handle their admission to decode. Although the SWAPPED status name remains, this scheduler has no corresponding swapped queue or swap branch in _preempt [37, 6].

We use that fixed vLLM release as a comparison, not an identified fork point. The familiar waiting/running distinction is only the starting point. Follow DeepInfer’s rank-wide step-type votes and its separate in-flight counts below: they determine what a request may do before the preceding device result is known. Figure 5.2 places these decisions on the request path.

5.2.1 Choose prefill or decode

DeepInfer’s released scheduler distinguishes three engine types. A PROMPT engine runs the prefill path; a COMPLETION engine runs the decode path. A NORMAL engine can do either (scheduler.hpp, lines 288–295). A normal engine chooses one kind for each step. It does not build a batch mixing prompt chunks and decode rows in these two paths [37].

That choice also involves other participants. A data-parallel rank is one participant with its own request batch in a cooperating group. The ranks must agree on the kind of step. Detailed distribution belongs to chapter 7; here the consequence is that a non-empty local queue cannot decide alone.

308bool is_any_point = (decode_step != 0) && (decode_step % prefill_any_attempt_interval == 0);
309bool is_all_point = (decode_step % prefill_attempt_interval == 0);
310if (!is_any_point && !is_all_point) {
311    return _schedule_decode(acquire);
312}
313bool want = (get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK) && _has_prefill_waiting();

At an ALL attempt, every rank must want prefill; at an ANY attempt, one is enough. Lines 315–319 implement the votes. A rank wants prefill only if its cache usage is at or below the watermark and the head of its waiting queue is schedulable. The default intervals are 70 and 100 decode steps respectively, from environ.py, lines 174–178. An ANY point takes precedence when the two coincide. A prefill step resets decode_step to zero, which is another ALL attempt point [30, 37].

Agreement to try is not admission. After trying, the ranks vote again to check whether any produced a prefill result. If all attempts returned no work, they fall back to decode (lines 328–340). The source comment explains why: an empty prompt step would reset the counter, invite another unsuccessful attempt and prevent running requests from making progress. The ANY interval therefore offers an opportunity for isolated prefill work; it is not a wall-clock latency guarantee or a promise that memory will be available.

5.2.2 The head of the queue must fit

_schedule_prefill examines waiting.front(). For a fresh request, it checks the prompt length and asks whether both pools can provide the storage required by get_num_blocks. Chapter 4 established that this reserves blocks for the whole prompt, even when computation will be chunked.

386auto can_allocate = _can_allocate_with_state(seq->get_num_blocks());
387if (can_allocate == AllocStatus::LATER) {
388    break;
389} else if (can_allocate == AllocStatus::NEVER) {
390    LOG_WARN("Input prompt ({} tokens) is too long and "
391             "exceeds the capacity of block_manager",
392        num_prompt_tokens);
393    seq->schedule.status = SequenceStatus::FINISHED_IGNORED;
394    ignored_seqs.push_back(seq);
395    waiting.pop_front();
396    continue;
397}

The difference between break and continue is a scheduling policy. LATER leaves the request at the head and stops this pass. A smaller request behind it is not tried. This is head-of-line blocking: work that could fit waits behind work that cannot yet fit. NEVER marks the request ignored, removes it and tries the next one. It cannot become runnable by waiting.

For example, suppose three blocks are free. The head needs four and the next request needs one. If the pool can eventually provide four, the head receives LATER and neither request is admitted in this pass. Choosing to skip the head would use the free memory now, but repeated skips could keep a large request waiting indefinitely. The released path chooses to stop; the code establishes that choice without establishing which policy is best for a service.

Two other limits follow (lines 440–449). The scheduler limits the number of requests alongside those already running, and subtracts each selected request’s whole history length from prefill_sum_seqlen_k_remain. Despite the name, this second check does not charge only the current chunk’s key length. It is separate from the query-position budget below. A request can pass the storage check and still fail one of these limits.

5.2.3 Choose how many positions to evaluate

Chunking is off by default: EngineArgs.chunked_prefill is False in deepinfer/engine/arg_utils.py, line 42. The shared example setup enables it in examples/_config.py, line 120, unless the caller overrides that setting. The chunked paths below assume it has been enabled [39, 17].

The unsatisfied prompt work is get_len() - cache.num_tokens. The scheduler first caps that work by the remaining query budget:

470if (seq_workload > (max_num_batched_tokens - num_batched_tokens) / cp_align * cp_align) {
471    new_workload = std::min(seq_workload, (max_num_batched_tokens - num_batched_tokens) / cp_align * cp_align);
472    seqlen_q = new_workload;
473    TORCH_CHECK(
474        seqlen_q < seq_workload, "_schedule: is_chunked must be true:, seqlen_q:", seqlen_q, " seq_workload: ", seq_workload);
475}

Let $T$ be max_num_batched_tokens, $U$ the positions already charged to this batch and $a$ the alignment. The available allowance here is

For example, with $T=259$, $U=252$ and $a=4$, only four of the seven remaining positions are available to the next chunk. The constructor sets cp_align to 1 for one context-parallel participant, or twice the group size otherwise (line 90). Context parallelism divides context computation across devices; Appendix B shows how its ranks are grouped and how each sequence is split between them. For this chapter, alignment means that useful positions and charged positions need not match.

The cap is not the final chunk size. Lines 477–510 respect explicit checkpoints, move boundaries away from the inside of an image, check whether chunking is enabled and enforce a minimum chunk length. Once the chunk passes those checks, the scheduler commits the decision:

515num_batched_tokens += (seqlen_q + cp_align - 1) / cp_align * cp_align;
516
517TORCH_CHECK(seq->schedule.status == SequenceStatus::WAITING, "status ", int(seq->schedule.status));
518if (seq->cache.block_table.empty()) {
519    acquire.allocate_seq(*seq);
520}
521
522if (!is_chunked) {
523    num_curr_seqs += 1;
524    seq->schedule.status = SequenceStatus::RUNNING;
525    allocate_mtp_info_slot(*seq);
526}

The first line charges $a\lceil q∕a\rceil$ positions for $q$ useful queries. Five queries at alignment four cost eight budget positions. This is an accounting rule, not a claim that every padded position produces an output token.

The fresh request now owns blocks and a state slot. It enters the selected list and leaves waiting. If this is only a partial prompt, its status stays WAITING; if it is the final chunk, its status becomes RUNNING. Status is not queue membership: a partially processed request can temporarily be outside the waiting queue while its chunk is in flight.

When process_outputs handles an unfinished chunk, it puts the request back at the front:

1016TORCH_CHECK(seq_ptr->schedule.status == SequenceStatus::WAITING,
1017    SPRINT_VARS(
1018        "In process_outputs, should be waiting", seq.cache.num_tokens, seq.input.tokens.get_len(), seq.get_nextn_len()));
1019waiting.push_front(seq_ptr);

That is the front-return rule from the experiment, not its rotation rule. Real execution adds overlap, group agreement and the possibility of selecting several prompts in one step, with a limit: a NORMAL engine stops selecting requests as soon as it selects an unfinished chunk (lines 543–545). That chunk is therefore the last request selected in the step. Returned to the front after processing, it again precedes later arrivals. The experiment’s times do not describe this engine, but these two queue operations explain why chunking alone does not let a later short prompt overtake a long one.

5.2.4 Prefill can keep decode waiting

Combine that return to the front with the step-type rule. The worker processes a prompt step’s outputs before returning, even with ZOS enabled (worker.py, lines 101–102 and 182–187). An unfinished A is therefore back at the head before the next scheduling call. Since prefill reset decode_step to zero, that call is another ALL attempt. With one rank, prefill can continue while the head remains schedulable, cache usage is at most the 0.9 watermark and the admission checks succeed. There need not be a decode step between A’s chunks [7, 37].

With several ranks, consecutive ALL attempts require every rank to want prefill. At an ANY attempt, however, some ranks may have no local work. At either kind of attempt, a rank that wanted prefill may still admit nothing, for example because its head request needs blocks that are not free. If another rank has scheduled prefill, such ranks execute an empty prompt step (scheduler.hpp, lines 332–340); their running decode requests do not advance either. During these prompt steps the engine schedules no new decode work. It may still deliver results from an earlier decode step, which the worker post-processes before sampling the current step. Once those results are drained, waiting for the next decode step can lengthen chapter 2’s inter-token intervals. This consequence follows from the inspected control flow; it is not a measured engine result.

The scope matters. NORMAL serves offline or combined prefill/decode execution. The production Q&A describes separate prefill and decode instances, the separate roles represented by PROMPT and COMPLETION in the released scheduler. The repeated choice between prefill and decode described here belongs to the normal engine, not that separated production architecture [2, 37].

5.2.5 A chunk becomes positions and addresses

Trace a text-only request with 300 prompt tokens. Use a normal engine with chunking enabled, a query budget of 128, alignment one, no prefix hit and no draft positions. Assume all other limits permit admission, enough memory is free and each previous chunk has been processed before the next is scheduled.

At its first admission, get_num_blocks asks for $\lceil 300∕128\rceil =3$ blocks. The allocator also gives it a state slot. All three chunks keep those allocations:

Scheduler::schedule builds a metadata record for each selected request. Its key length is the processed count plus in-flight positions plus this step’s query length (lines 696–707). Our trace has no in-flight positions from another step. The SequenceMetadata constructor takes the block-table reference and state ID from the sequence (sequence.hpp, lines 374–381).

One downstream consumer makes the connection to chapter 3 explicit. The prefill input builder converts those lengths and that state ID into SWA write addresses:

98auto slot_mapping_store_swa = get_or_create_tensor(result, has_out, "slot_mapping_store_swa", {total_q}, torch::kInt32);
99auto p_slot_mapping_store_swa = slot_mapping_store_swa.mutable_data_ptr<int32_t>();
100int32_t i_q = 0;
101for (auto p_seq : seq_metadata_list) {
102    for (int32_t j = 0; j < p_seq->seqlen_q; ++j, ++i_q) {
103        if (p_seq->seqlen_q - j <= swa_block_size)
104            p_slot_mapping_store_swa[i_q] = p_seq->state_id * swa_block_size + (p_seq->seqlen_k - p_seq->seqlen_q + j) % swa_block_size;
105        else
106            p_slot_mapping_store_swa[i_q] = -1;
107    }
108}
109TORCH_CHECK_EQ(i_q, total_q);

Within a chunk, query $j$ is at logical position $p=\mathtt{\text{seqlen\_k}}-\mathtt{\text{seqlen\_q}}+j$. For the second chunk, $j=0$ gives $p=128$. In state slot 2 with the default ring length 134, its destination is $2\times 134+128=396$; query $j=6$ wraps to $2\times 134+0=268$. All our chunks fit within the ring. For a chunk longer than the ring, the condition keeps only its last ring-length positions; earlier destinations receive $-1$. Chapter 3 explained why prompt attention has a temporary input and reads the old cache before writing these final rows [38].

The scheduler has therefore chosen more than request names. It has chosen a range of logical positions, a budget charge and the storage through which the worker will address those positions.

5.3 When running requests exhaust the pool

Admission is not the last memory decision. A request already producing tokens can need another block. Preemption means taking resources away from such a request so another can proceed. It may suspend work for later recovery or terminate the request. DeepInfer has both outcomes.

Run a second experiment:

python3 labs/12_preempt.py

The pool has 19 usable blocks, each with four positions. A, B and C own all of them. B and C have just generated a token beyond the end of their last block; that token is in their history but has no cache row yet.

Which request should give up its blocks? B owns the most, but it has generated the least output. These are different possible priorities. The experiment uses the inspected decode ordering, then applies each of the three branch conditions in _preempt:

stop=False, A.so_status=0
  output order: [('A', 8), ('C', 4), ('B', 2)]
  candidate=B; victim=A; blocks A=5, B=11, C=3
  A: WAITING, processed=0, blocks=[], state=-1
  waiting=['A', 'D'], ignored=[], free=[3, 2, 1]
  history retained=True; scheduled=['B', 'C']

stop=True, A.so_status=0
  output order: [('A', 8), ('C', 4), ('B', 2)]
  candidate=B; victim=A; blocks A=5, B=11, C=3
  A: FINISHED_PREEMPTED, processed=19, blocks=[], state=-1
  waiting=['D'], ignored=['A'], free=[3, 2, 1]
  history retained=True; scheduled=['B', 'C']

stop=False, A.so_status=1
  output order: [('A', 8), ('C', 4), ('B', 2)]
  candidate=B; victim=A; blocks A=5, B=11, C=3
  A: FINISHED_PREEMPTED, processed=19, blocks=[], state=-1
  waiting=['D'], ignored=['A'], free=[3, 2, 1]
  history retained=True; scheduled=['B', 'C']

A loses five blocks. B and C take one each, leaving three free. In the first run A goes ahead of D in the waiting queue and loses its processed cache count. Its 20-token history survives. In the second run a stop flag makes A finish; in the third an active structured-output status does so even though the stop flag is false. All three return A’s blocks and state slot. The nonzero processed count printed for a finished A is no permission to read its freed storage.

5.3.1 Read the comparator and both ends of the queue

After admitting ready forks from prompted, the decode scheduler sorts running and selects requests:

624std::sort(running.begin(), running.end(),
625    [](const SequencePtr &a, const SequencePtr &b) { return a->get_output_len() > b->get_output_len(); });
626scheduled.reserve(max_seqs);
627size_t l = 0, r = running.size();
628while (l < r && scheduled.size() < max_seqs) {
629    auto &seq = running[r - 1];
630    while (!block_manager->can_allocate_single_nolock() && r - l > 1) {
631        auto victim_seq = running[l];
632        _preempt(victim_seq, ignored_seqs);
633        l += 1;
634    }
635    if (block_manager->can_allocate_single_nolock()) {
636        _append_slot(*seq);
637        scheduled.push_back(std::move(seq));
638    } else {
639        _preempt(seq, ignored_seqs);
640    }

The comparator orders output lengths from greatest to least: A, C, B in the experiment. The candidate is running[r - 1], so B goes first. If no block is free, the victim is running[l], so A is preempted. Output length is already generated work, not remaining answer length or allocated memory. Nothing here can tell that B will finish soon.

Read two nearby comments with that trace in mind. The decode scheduling comment says it preempts the shortest output (line 557, author’s translation), while the body takes a victim from the longest-output end. In sequence.hpp, line 56, the comment on the preemption finish reason says that the request occupying the most KV cache is forcibly terminated (author’s translation). The example separates those quantities: A is selected although B owns more blocks. The control flow also permits recomputation. These comments do not fully describe the inspected policy.

There is another constraint in the body: even a candidate whose existing tail has space is considered only after the free-block check succeeds. _append_slot might then allocate nothing. Thus “my table is large enough” does not by itself guarantee that this loop will schedule the request. If no other victim remains and no block becomes free, it preempts the candidate itself. Pins from chapter 4 matter here: freeing one request does not necessarily free any shared block.

After choosing survivors, the scheduler sorts them again, this time by total history length, before building per-request query lengths (lines 648–657). The first sort determines who survives memory pressure; the second determines the batch order passed onwards. Ties in these comparators do not establish an arrival-order guarantee [37].

5.3.2 Termination and recomputation are different contracts

The default depends on the entry point. Unless the caller supplies a value, offline LLM sets stop_preempted_request to False (llm.py, line 68), giving the recovery policy used in the lab’s first run. Construction paths that retain the EngineArgs default instead use True (arg_utils.py, line 40), which requests termination. The structured-output guard still applies in either case [4, 39].

The preemption branch first asks whether it is allowed to recover this request:

770const bool so_blocks_recompute = seq->structured_output.status != 0;
771if (stop_preempted_request || so_blocks_recompute) {
772    if (so_blocks_recompute && !stop_preempted_request) {
773        LOG_WARN("Request {} is preempted and stops: structured output (so_status {}) can not be recomputed.",
774            seq->request.request_id, seq->structured_output.status);
775    } else {
776        LOG_WARN("Request {} is preempted and stops.", seq->request.request_id);
777    }
778    seq->schedule.status = SequenceStatus::FINISHED_PREEMPTED;
779    free_seq_nolock(*seq);
780    ignored_seqs.push_back(seq);
781} else {
782    LOG_WARN("Request {} is preempted and will be recomputed.", seq->request.request_id);

Structured output constrains generation to an allowed form, such as a grammar, and carries state beyond token IDs. The actual guard is structured_output.status != 0; it is not a test of whether the original request merely supplied a grammar. The log message says this state cannot be recomputed by the present path. Sequence::reset_for_recompute helps us see the boundary:

307void reset_for_recompute() {
308    cache.reset_for_recompute();
309    schedule.reset_for_recompute();
310    mtp.reset_for_recompute();
311}

The helpers reset the current and previous processed counts and the SWA valid start; they clear both in-flight counters and selected draft state. They do not erase the token history, sampling configuration or structured-output matcher (sequence.hpp, lines 134–142, 223–227 and 258–261). A recovery path would have to establish a compatible matcher state too. That is an inference from the state separation, not a claim that replaying a grammar is impossible [6].

In the recovery branch, line 783 calls this reset. Its Chinese comment points to the reset helpers as the inventory of rolling state (author’s translation). The remaining operations put the request back and return its resources:

785    waiting.push_front(seq);
786    seq->schedule.status = SequenceStatus::WAITING;
787    free_seq_nolock(*seq);
788}

When A returns to prefill in our case, its history still has 20 tokens but its processed count is zero. The prompt path therefore sees 20 positions of work, including tokens already generated. It rebuilds state for that history; it does not erase the eight output tokens and start a new answer. This recovery adds computation and delay. It does not by itself prove bitwise identity after replay; chapter 10 separates the relevant correctness contracts.

In the vLLM comparison, an unspecified preemption mode selects recomputation when get_max_num_running_seqs()==1 and swapping otherwise; a configured mode can override this. Its recomputation path sets the sequence to WAITING, frees its blocks and resets computed progress (vllm/core/scheduler.py, lines 1472–1510) [13]. DeepInfer also has a rebuild-from-history path, but here the choice is between recomputation and termination, with engine-type and structured-output guards. Those are different recovery contracts despite the shared term preempt.

Recomputation also requires a route back to prefill. The scheduler constructor rejects a COMPLETION engine with stop_preempted_request=false (lines 102–104): that engine has no prefill scheduling path. Termination is reported as insufficient_system_resource by get_finished_reason. A larger service could retry elsewhere, but that is not an action performed by this branch.

Preemption may happen while a previous step is still in flight. In process_outputs, the comment and check at lines 917–924 describe skipping results for a decode request already reset to WAITING. Otherwise stale output could advance the newly reset count or return the request to the wrong queue. This is where scheduling meets chapter 6’s subject: work can finish after the host has decided that its result is no longer usable.

5.4 A position budget is not a time budget

The released prefill path caps query positions, request counts, full lengths and allocations. These are useful constraints, but they are not equivalent to a maximum step duration. Chapter 2 showed that compute and memory costs depend on different quantities. Equal query counts can still read different amounts of context or route work to different experts.

The deployment Q&A describes an additional production policy: a fitted quadratic $\text{cost}(q,c)$, where $q$ is the number of new positions and $c$ the cached prefix length. A nanosecond budget named compute_quota_limit and a 32,768-token cap constrain the selected chunk. The policy also reduces the budget across a data-parallel group [2]. The inspected _schedule_prefill does not implement that quadratic solver. Its printed budget branch uses remaining token capacity and alignment. The timed model in the Q&A belongs to the documented service; we should not read it into these source lines. The runtime model used for draft verification is a separate path, which chapter 8 will examine.

We can now follow a batch decision from a waiting queue to a worker address, and follow memory pressure back to a reset or a reported failure. The remaining question is when those decisions can safely be made while an earlier batch is still executing.

5.5 Experiments to try next

1. In labs/11_chunks.py, remove the per-step overhead. Predict which first-token delays become equal and which still depend on queue order. Update the expected results before running the altered experiment.
2. In the memory experiment, give A and B equal output lengths. What additional rule would make the victim deterministic? Python’s sort is stable; the C++ call shown here is std::sort. Do not infer the engine’s tie behaviour from the lab.
3. Trace a candidate whose existing block table is large enough when the free list is empty. Locate the check that still causes a preemption. Propose a change, then list the ownership and in-flight cases it would need to preserve. Label the change as a proposal, not an existing DeepInfer policy.
4. In _schedule_prefill, compare a first admission, an unfinished chunk and a forked prefix. Which can own a state slot while its status is WAITING? Use this to revisit chapter 4’s extra state-slot exercise.