Skip to main content
ModelRunner runs on every GPU rank. The rank-0 instance is owned directly by LLMEngine; ranks 1..N-1 run in separate worker processes and communicate via shared memory + multiprocessing.Event. It handles weight loading, KV cache allocation, prefill/decode input preparation, and optional CUDA graph capture for fast decode.

​
Constructor

During __init__ the following steps happen in order:
  1. dist.init_process_group("nccl", ...) — collective barrier; all ranks must call this.
  2. Model construction and weight loading on the assigned GPU.
  3. warmup_model() — dry-run forward pass to measure peak memory.
  4. allocate_kv_cache() — allocates the KV cache pool.
  5. capture_cudagraph() — captures decode graphs (skipped if enforce_eager=True).
  6. Shared-memory setup for IPC (skipped when world_size == 1).
​
dict
required
The same config dict passed to LLMEngine. Must include all model architecture keys (vocab_size, hidden_size, num_layers, etc.) as well as the following runtime keys consumed directly by ModelRunner:
​
int
required
CUDA device index for this worker. Rank 0 is the primary process; ranks 1..world_size-1 are spawned workers.
​
Event | list[Event]
required
For rank 0: a list of multiprocessing.Event objects, one per worker rank, used to signal that new IPC data has been written to shared memory. For rank != 0: a single Event used to wait for commands from rank 0.

​
Methods

​
warmup_model

Runs a synthetic prefill forward pass at the maximum batch size (max_num_batch_tokens // max_model_length sequences of length max_model_length) to force all CUDA kernels to JIT-compile and to record the peak GPU memory footprint. The result is used by allocate_kv_cache() to determine how much memory is available for the KV cache pool.

​
allocate_kv_cache

Allocates a single global KV cache tensor of shape (2, num_layers, max_cached_blocks, block_size, num_kv_heads_per_rank, head_dim) and assigns slices of it to each attention layer’s k_cache / v_cache attributes. The number of blocks is derived from:
When world_size > 1, an all_reduce(MIN) synchronises the block count across ranks so the scheduler never allocates more blocks than the most memory-constrained rank can hold.
max_cached_blocks in config is overwritten with the computed value after this method returns.

​
prepare_prefill

Builds the input tensors for a varlen prefill forward pass and stores them in the thread-local attention context via set_context(is_prefill=True, ...): Prefix-cached tokens are skipped in input_ids and slot_mapping — only uncached tokens require a new forward pass.
​
list[Sequence]
required
Sequences scheduled for prefill. Each sequence must have an allocated block_table.

​
Return value

​
torch.Tensor
1-D long tensor of concatenated (non-cached) token IDs, on the current CUDA device.

​
prepare_decode

Builds input tensors for a decode step and stores them in the attention context via set_context(is_prefill=False, ...):
​
list[Sequence]
required
Sequences scheduled for decode.

​
Return value

​
torch.Tensor
1-D long tensor of last-token IDs, shape (batch_size,), on the current CUDA device.

​
run_model

Executes the model forward pass and returns the logit tensor.
  • Prefill / eager mode: calls self.model(input_ids) directly, then model.compute_logits(hidden_states).
  • Decode (CUDA graph): finds the smallest captured graph whose batch size ≥ len(seqs), copies the current tensors into the pre-allocated graph variables, replays the graph, and computes logits from graph_vars["outputs"].
​
torch.Tensor
required
Token IDs prepared by prepare_prefill or prepare_decode.
​
bool
required
True to use the eager path; False to use CUDA graph replay.

​
Return value

​
torch.Tensor
Logit tensor. Shape (num_tokens, vocab_size) for prefill; (batch_size, vocab_size) for decode.

​
run

Main inference entry point. Calls prepare_prefill or prepare_decode, then run_model, then the sampler. Only rank 0 samples tokens; all other ranks return None.
​
list[Sequence]
required
Sequences to process.
​
bool
required
Whether to run a prefill or decode step.

​
Return value

​
torch.Tensor | None
1-D tensor of sampled token IDs on rank 0 (one per sequence). None on worker ranks.

​
capture_cudagraph

Pre-allocates tensors at maximum sizes and captures CUDA graphs for decode batch sizes [1, 2, 4, 8] plus every multiple of 16 up to max_num_seqs. At decode time run_model selects the smallest captured graph whose size is ≥ batch_size, so the overhead of padding is minimised. Captured graphs are stored in self.graphs (dict mapping batch size → CUDAGraph) and input/output tensors are stored in self.graph_vars.
CUDA graph capture is skipped when config["enforce_eager"] = True.
Captured batch sizes:

​
call

IPC dispatch method used by LLMEngine to invoke model operations.
  • On rank 0: serialises (method_name, *args) into shared memory via write_shm(), then executes the method locally.
  • On all ranks: looks up method_name on self and calls it with args.
​
str
required
Name of the ModelRunner method to invoke (e.g. "run", "exit").
​
Any
Positional arguments forwarded to the method.

​
loop

Blocking event loop for worker ranks (rank != 0). Continuously reads IPC commands from shared memory via read_shm(), dispatches them via call(), and exits cleanly when "exit" is received.
This method asserts world_size > 1 and rank != 0. It must not be called on the rank-0 process.

​
exit

Cleans up GPU resources: closes and unlinks shared memory (rank 0 unlinks), deletes CUDA graphs, synchronises the device, and destroys the NCCL process group.

​
CUDA graph batch sizes

The table below shows which graphs are captured for a given max_num_seqs. At runtime the smallest graph bs_ >= actual_batch_size is replayed, so at most bs_ - actual_batch_size dummy tokens are computed.