Building the Batch

The idea in one minute#

A textbook batch is a rectangle: batch × sequence_length, short rows padded. vLLM’s batch is not a rectangle. One step may hold a request computing a single token, another computing 5,000 tokens of its prompt and a third computing three. Padding them to a common length would waste almost all of the work. Instead the model runner flattens everything into one long row of tokens and builds a few small index tensors that say which request each token belongs to, what position it has, where in the KV cache its result must be written, and where its context must be read from. The transformer’s dense layers never notice; only attention needs the indices. Seven tensors describe the whole step.

A picture#

flowchart TB
  subgraph SCH["From the scheduler"]
    direction LR
    A["A: 1 token<br/><small>decoding</small>"]
    B["B: 5 tokens<br/><small>mid-prompt</small>"]
    C["C: 3 tokens<br/><small>new</small>"]
  end
  SCH --> FLAT[":i-layers: <b>One flat row of 9 tokens</b><br/><small>input_ids, positions</small>"]
  FLAT --> DENSE[":nvidia: <b>Embedding, MLP, norms</b><br/><small>treat it as 9 independent rows</small>"]
  FLAT --> IDX[":i-list-checks: <b>Index tensors</b><br/><small>query_start_loc, seq_lens,<br/>block_table, slot_mapping</small>"]
  IDX --> ATT[":nvidia: <b>Attention</b><br/><small>write new K/V at slot_mapping<br/>read context through block_table</small>"]
  DENSE --> ATT
  ATT --> LAST[":i-funnel: <b>logits_indices</b><br/><small>keep 3 rows: each request's last token</small>"]
  LAST --> S[":i-zap: <b>Sampler</b>"]
  class A,B,C neutral
  class FLAT,IDX memory
  class DENSE,ATT compute
  class LAST,S queue

How it really works#

Why flatten#

Most of a transformer treats each token position independently. Embedding lookup, the feed-forward network, normalisation and the final projection are all applied row by row. For those layers the only thing that matters is the total number of tokens; which request each belongs to is irrelevant.

Attention is the single place where tokens interact, and it interacts only within a request. So the batch can be one flat list, provided attention is told where each request’s rows start and stop. A step with one 5,000-token chunk and 200 single-token decodes is then exactly 5,200 rows of work, with no padding at all.

The seven tensors#

For a step with N requests and T scheduled tokens in total:

How they are computed#

The older runner’s _prepare_inputs documents itself with a worked example in its comments. For num_scheduled_tokens = [2, 5, 3]:

# Get request indices.
# E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)

# cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
# self.query_pos.np[:10]: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
cu_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np)

# Get positions.
positions_np = (
    self.input_batch.num_computed_tokens_cpu[req_indices]
    + self.query_pos.np[: cu_num_tokens[-1]]
)

A token’s position is its request’s num_computed_tokens plus its offset within this step’s chunk. That is all “chunked prefill” means at the tensor level: the second chunk of a prompt starts at a position other than zero.

Token IDs are then gathered from the persistent batch, which stores every request’s full token list as one row of a max_num_reqs × max_model_len array:

# Get token indices.
# E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
# -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
# where M is the max_model_len.
token_indices = positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]

# NOTE(woosuk): We use torch.index_select instead of np.take here
# because torch.index_select is much faster than np.take for large tensors.
torch.index_select(
    self.input_batch.token_ids_cpu_tensor.flatten(), 0, token_indices_tensor,
    out=self.input_ids.cpu[:total_num_scheduled_tokens],
)

One vectorised gather, no Python loop over requests. The comment about index_select versus np.take is typical of this file: each line was chosen by measurement.

And the boundaries:

self.query_start_loc.np[0] = 0
self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
# Note: pad query_start_loc to be non-decreasing, as kernels
# like FlashAttention requires that
self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])

The buffer is sized for the maximum number of requests; the unused tail is filled with the total so that every “request” beyond the real ones has zero length.

The first line of the function#

# OPTIMIZATION: Start copying the block table first.
# This way, we can overlap the copy with the following CPU operations.
self.input_batch.block_table.commit_block_table(num_reqs)

The block table is the largest tensor in the step. Its copy to the GPU is started first, asynchronously, so that it runs while the CPU computes everything above. Ordering work so that transfers overlap computation is a recurring pattern in the runner.

Reading and writing the KV cache#

Recall that the KV cache is one large tensor per layer, cut into blocks, and that a request’s tokens live in blocks scattered anywhere in it.

Writing. Each token processed this step produces a key and a value in every layer. They go to a specific physical slot:

slot = block_table[request][position // block_size] * block_size + position % block_size

slot_mapping holds that number for all T tokens. In the current code a small GPU kernel computes it from positions, query_start_loc and the block table (_COMPUTE_SLOT_MAPPING_KERNEL in vllm/v1/worker/block_table.py), so the mapping is never materialised in Python.

Reading. To compute attention for a token, the kernel needs the keys and values of all earlier tokens of the same request. It does not get a contiguous slice. It gets the request’s row of the block_table and its seq_len, and follows the block IDs (Attention Backends).

The block table on the worker is updated by appending. When the scheduler grants a request a new block, only that one ID crosses to the worker, and it is written into the next column of the request’s row.

Only one row per request reaches the vocabulary#

The final projection from a hidden state to a score for every token in the vocabulary is the largest matrix multiplication per row: hidden size times 100,000 or more. Computing it for all T rows would be pure waste, because a next-token distribution is only needed at the last position of each request:

logits_indices = query_start_loc[1:] - 1

The hidden states are indexed with logits_indices before the projection. A step with 5,200 tokens from 201 requests projects 201 rows, not 5,200.

A request whose chunk did not reach the end of its prompt still gets a row and a sampled token, because it is cheaper to compute a fixed shape than to branch. The source notes it: “We will ignore the sampled tokens from the partial requests.” The scheduler knows the request is still prefilling and discards the result.

With speculative decoding, several positions per request need distributions (one per draft token plus one), and logits_indices is longer.

Multimodal inputs take a different door#

For a text token, input_ids is all the model needs. For an image, the prompt contains a run of placeholder tokens, and the real input is a tensor of embeddings produced by a vision encoder.

When the scheduler grants a chunk that overlaps an image’s placeholders, the runner runs the encoder for that image (once; the result is kept in an encoder cache until the request has consumed it) and then, after the ordinary embedding lookup, overwrites the embeddings at the placeholder positions with the encoder’s output. From there on the flat batch is processed as usual. See Multimodal.

Padding, after all that#

One kind of padding does survive: to the nearest CUDA graph size. A recorded graph works for one exact token count, so a batch of 37 tokens is padded to the next captured size, 40 (CUDA Graphs and Compilation). The padding rows are marked as belonging to no request (their query_start_loc range is empty) and their outputs are dropped. This is a few percent of waste in exchange for a much faster launch.

The same job in Model Runner V2#

The newer runner computes the same tensors on the GPU with Triton kernels. Its design document gives the reasons:

1. Better async behavior: GPU can derive values (for example with speculative decoding) that CPU may not know yet.
2. Lower CPU overhead: input prep is very cheap on GPU and avoids Python bottlenecks.

The first reason connects to Async Scheduling: the input token for a decoding request is the previous step’s sample, which exists only on the GPU. Assembling input_ids on the GPU means that value never has to travel to the CPU and back.

Large per-request state such as full prompt token lists stays in CPU memory and is read by GPU kernels directly through unified virtual addressing, rather than being duplicated on the GPU.

Code#

Build all seven tensors for a step with three requests in different phases. The block size is 4 so that the numbers stay small.

package main

import "fmt"

const blockSize = 4 // small, so the numbers are readable (vLLM's default is 16)

type request struct {
	name       string
	tokenIDs   []int // prompt + output so far
	computed   int   // num_computed_tokens before this step
	scheduled  int   // tokens the scheduler granted this step
	blockTable []int // physical block IDs, in order
}

func main() {
	// Three requests in one step: num_scheduled_tokens = [1, 5, 3].
	reqs := []request{
		// A is decoding: one token, the one it sampled last step, at position 10.
		{"A", seq(100, 11), 10, 1, []int{7, 2, 12}},
		// B is in the middle of its prompt: 5 tokens at positions 4..8.
		{"B", seq(200, 20), 4, 5, []int{3, 9, 5}},
		// C is new: its first 3 tokens.
		{"C", seq(300, 3), 0, 3, []int{6}},
	}

	var reqIndices, positions, inputIDs, slotMapping, seqLens []int
	queryStartLoc := []int{0}

	for i, r := range reqs {
		for q := 0; q < r.scheduled; q++ {
			pos := r.computed + q
			reqIndices = append(reqIndices, i)
			positions = append(positions, pos)
			inputIDs = append(inputIDs, r.tokenIDs[pos])
			// The token at logical position pos lives in physical slot:
			//   block_table[pos / block_size] * block_size + pos % block_size
			slot := r.blockTable[pos/blockSize]*blockSize + pos%blockSize
			slotMapping = append(slotMapping, slot)
		}
		queryStartLoc = append(queryStartLoc, len(positions))
		seqLens = append(seqLens, r.computed+r.scheduled)
	}

	// One row of logits is needed per request: its LAST scheduled token.
	var logitsIndices []int
	for i := range reqs {
		logitsIndices = append(logitsIndices, queryStartLoc[i+1]-1)
	}

	fmt.Println("flat batch of", len(inputIDs), "tokens from", len(reqs), "requests")
	fmt.Println()
	fmt.Println("req_indices     ", reqIndices)
	fmt.Println("positions       ", positions)
	fmt.Println("input_ids       ", inputIDs)
	fmt.Println("slot_mapping    ", slotMapping, "  <- where each new K/V is WRITTEN")
	fmt.Println()
	fmt.Println("query_start_loc ", queryStartLoc, "  <- request i owns rows [loc[i], loc[i+1])")
	fmt.Println("seq_lens        ", seqLens, "  <- how much context each request attends over")
	fmt.Println("logits_indices  ", logitsIndices, "  <- rows that get a next-token distribution")
	fmt.Println()
	fmt.Println("block_table (one row per request; where its context is READ from):")
	for _, r := range reqs {
		fmt.Printf("  %s %v\n", r.name, r.blockTable)
	}

	fmt.Println()
	for i, r := range reqs {
		kind := "prefill chunk: sampled token is discarded"
		if r.computed+r.scheduled == len(r.tokenIDs) {
			kind = "caught up: a token is sampled from row " + fmt.Sprint(logitsIndices[i])
		}
		fmt.Printf("  %s computed %d + scheduled %d of %d known tokens -> %s\n",
			r.name, r.computed, r.scheduled, len(r.tokenIDs), kind)
	}
}

func seq(base, n int) []int {
	out := make([]int, n)
	for i := range out {
		out[i] = base + i
	}
	return out
}

Check one slot by hand. Request B’s token at position 8 is in its third block (8 / 4 = 2), which is physical block 5, at offset 8 % 4 = 0: slot 5 × 4 + 0 = 20. Its four neighbours at positions 4 to 7 are in a different physical block entirely (block 9, slots 36 to 39). Two adjacent tokens of the same request are nineteen slots apart in memory, and nothing in the model cares.

Remember this#

• The batch is flat: all scheduled tokens of all requests in one row, with no padding between requests.
• query_start_loc and seq_lens tell attention where each request’s rows and context are.
• A token’s position is num_computed_tokens + offset in this chunk.
• slot_mapping says where new keys and values are written; block_table says where context is read.
• Only each request’s last row is projected to the vocabulary.
• The block table copy starts first so it overlaps with the rest of input preparation.
• The newer runner builds these tensors on the GPU, which is what async scheduling needs.

Try it#

1. Add a fourth request that is new and has a 9-token prompt with block table [1, 4, 8]. Predict slot_mapping for its tokens before running.
2. Change blockSize to 16 and give request B the block table [3]. Which of the seven tensors change, and which do not?
3. In the vLLM source, find _prepare_inputs and locate the line that computes logits_indices for the speculative-decoding case. How does its length differ?

Check yourself#

1. Why does flattening the batch not confuse the feed-forward layers?
2. A token is at position 37 of a request whose block table is [12, 4, 30] with block size 16. Which physical slot holds its key and value?
3. Why is the vocabulary projection applied to only one row per request?

Sources#

Checked on 5 October 2026 against main at commit 0c16eee.

• vllm/v1/worker/gpu_model_runner.py — _prepare_inputs and its worked comments
• vllm/v1/worker/block_table.py — the block table and the slot-mapping kernel
• vllm/v1/worker/gpu_input_batch.py — the persistent batch
• Model Runner V2 design