Our first tour of TGI at /tgi/ told the maintenance-mode story: a server that officially stepped back and let the ecosystem bloom around it. That framing, while accurate on the marketing page, undersells the code, because the main branch is anything but dormant. The deployment model quietly inverted over the last two years. Where TGI once delegated scheduling to a Python server, today the Rust router itself owns the queue, the continuous batching loop, the KV-cache block allocator, and a radix-tree prefix cache, while each Python shard process shrinks into a thin forward-pass worker that answers gRPC calls. Reading the v3 architecture is worth the visit even if you never run TGI, because it is a working answer to the question every serving stack faces: which side of a language boundary should hold the scheduler?

The repository splits cleanly along that boundary. The launcher crate is the bootstrap binary, the router crate holds the HTTP layer and the scheduling machinery, the backends directory contains versioned gRPC glue, and the server directory is the Python side with its models, layers, and the protobuf definitions in proto. The files in play are remarkably readable. The launcher’s main.rs is about two thousand lines including argument plumbing, the v3 backend batcher sits in under six hundred lines, and the radix allocator with its trie is roughly a thousand including tests. The Python shard server is barely three hundred lines. As before, this is an educational walk through the published code of huggingface/text-generation-inference under the Apache-2.0 license.

TGI's v3 story: one launcher process spawns a Rust router and one Python gRPC shard per GPU, requests flow through validation into a budget-gated queue, the Rust batcher allocates paged KV blocks through a radix prefix cache, and every prefill or decode step broadcasts over Unix domain sockets to all shards.

Reading the overview from left to right:

• The launcher at launcher/src/main.rs enters through main(), resolves the shard count with find_num_shards, and spawns one manager thread per rank before starting the router.
• The HTTP layer lives in router/src/server.rs, with the health probe at lines 221, generate and streaming endpoints at 247 and 445, and the OpenAI chat surface at 1138.
• The Infer facade gates admission with a semaphore and hands validated requests to the Backend trait, whose only job is to schedule a request and return a stream.
• The v3 batcher at backends/v3/src/backend.rs runs batching_task, pulling budget-checked work from the queue state and driving prefill and decode calls.
• The RadixAllocator maps token prefixes to KV blocks through a RadixTrie with reference counting and LRU eviction.
• On the wire, ShardedClient discovers every shard through service_discovery and fans each call out to all of them.
• Inside each shard, the Python TextGenerationService servicer pops and stores live batches in a Cache and drives FlashCausalLM.

One launcher, three processes

The launcher is the only component that knows the whole topology, and its main() reads like a checklist that ends with a watchdog loop. After parsing arguments it derives the budget knobs that would otherwise be magic: when max_batch_prefill_tokens is not set, the block at lines 2114 through 2154 queries nvidia-smi for the card model, estimates optimal prefill size from an FLOPs table per GPU, and clamps the result against VRAM and the model’s maximum position embeddings. The same block of code special-cases VLM architectures with smaller defaults, a pragmatic admission that vision inputs blow past text budgets. Default CUDA graph sizes come out as the list 1, 2, 4, 8, 16, 32 unless quantization forbids graphs.

Weight download runs as a subprocess first, via download_convert_model, which shells out to the Python server’s own CLI so both sides agree on the cache layout and safetensors conversion. Then spawn_shards launches one thread per rank, each running shard_manager. The shard manager is where tensor parallelism is actually stitched together: it exports RANK, WORLD_SIZE, MASTER_ADDR, and MASTER_PORT into the child environment, which is precisely the contract torch.distributed expects for NCCL rendezvous, and it tells every rank to listen on its own Unix domain socket at {uds_path}-{rank}. The parent then waits for Ready signals from all shards at lines 1663 through 1687 before calling spawn_webserver, because the router’s very first act is to call warmup on the shards.

That supervision loop is honest about failure. Once everything is up, main() parks in a polling loop at lines 2333 through 2350: if any shard reports Failed, the whole container exits; if the webserver process dies, the shards are shut down too. There is no restart-and-pray. A sharded forward pass where one rank vanishes is unrecoverable state, and the launcher treats it that way.

The router is the scheduler now

The v3 handshake happens in connect_backend. The router connects over the master shard’s socket at line 111, clears any stale batch state left from a router restart at 117, and calls warmup with the configured budgets. The shard answers with the maximum total tokens its KV cache can hold, which the router then treats as a hard token budget for every future batch. This is the moment ownership transfers: from here on the Python side never decides what runs, it only executes what arrives.

Request handling starts at the Infer facade. A semaphore caps concurrent requests, validation runs on a dedicated worker pool so tokenization never blocks the HTTP runtime, and the validated request goes to backend.schedule(). One detail deserves applause: the stream wrapper at lines 164 through 188 notices when a generation stops early on the Length finish reason with tokens still owed, appends the produced text to the input, re-validates, and silently schedules a continuation. Users get their full max_new_tokens even when the backend’s single-pass limit cut things short, and no HTTP client ever learns that a seam existed.

The detail view: launcher wiring on the left, the Infer admission layer with its validation pool and continuation stream, the batching task's prefill-decode loop with filter and concat paths, the radix trie that rents shared prefix blocks, and the Python servicer that keeps live batches in a plain dictionary between gRPC calls.

Inside the batching task

The batcher is a single Tokio task that owns all in-flight state, which means no locks and no races, just message passing. After each notify it pulls a batch from the queue and enters a loop: prefill it, then repeatedly decode until the server stops returning cached batches. Between decodes it tries to grow. If waiting requests have waited longer than max_waiting_tokens, or the model supports chunked prefill, the block at lines 186 through 201 relaxes the minimum-batch-size rule that normally prevents onboarding tiny batches, and the new work is concatenated. With chunking support, the concat at lines 221 through 236 prefills the newcomer alongside the decode step in one call; without it, the batcher runs an extra prefill for the new batch and carries both forward.

The queue’s next_batch enforces the economics. Each candidate entry must fit three budgets at once: prefill tokens against max_batch_prefill_tokens, total tokens including the request’s full max_new_tokens reservation against max_batch_total_tokens, and, on the padded path, the batch-wide maximum input length multiplied across entries. The reservation trick at lines 333 through 338 shows why prefix caching needs a scheduler partner: when the radix trie holds the entire prompt, one token is subtracted back so the transformer always has at least one fresh position to compute, since a forward pass over a fully cached sequence produces nothing.

Prefill chunking gets symmetric treatment in the queue. When an oversized prompt busts the prefill budget and the model supports chunking, the block at lines 344 through 374 splits the request exactly at the remaining budget, admits the first chunk, and pushes the rest back for later steps. The Python side mirrors this with a per-batch budget in generate_token that reserves decode headroom while iterating requests in reverse age order, prioritizing older chunks. Nothing in the chunked request ever waits for the whole prompt to fit.

The radix prefix cache

Block allocation is where TGI’s Rust layer earns its keep. The RadixAllocator reserves block 0 for health checks at line 46 and keeps a free list plus a trie. Allocation walks the trie with the request’s token ids at allocate, increments the reference count on the matched node before anything else can evict it, and only then requests fresh blocks for the suffix from alloc_or_reclaim, which evicts the least recently used unreferenced prefixes when the free list runs dry. On free, the block at lines 157 through 208 decrements the refcount, inserts the newly computed prefix back into the trie, and cleverly handles the race where a partially overlapping request cached some of the same tokens mid-flight by returning the double-cached blocks to the free list.

The trie itself is a compressed radix tree keyed by token sequences hashed at block granularity. Its find returns matched block ids, evict walks nodes in access-time order, and insert splits existing nodes at arbitrary boundaries when a new prompt shares only part of a cached prefix. The unit tests in the same file read like a spec for multi-tenant chat serving: overlapping prefills, partial matches, and eviction under pressure all have dedicated cases.

On the wire this turns into block tables rather than KV payloads. ShardedClient broadcasts the same prefill or decode call to every shard and merges the generations from the master reply, while warmup takes the minimum max-total-tokens across shards because vocab sharding can leave different ranks with slightly different block counts.

The thin Python shard

Each shard process runs the TextGenerationService servicer over a Unix socket, with ServiceDiscovery answering the router’s bootstrap question, ClearCache and FilterBatch maintaining shared state, and Warmup probing peak memory. The Prefill handler concatenates an incoming chunk with the cached batch when chunking is on at lines 170 through 179, and Decode merges several cached batches into one forward pass at lines 208 through 214, which is how the router’s concat decisions land on the GPU as a single kernel launch sequence.

FlashCausalLM’s warmup does the arithmetic every serving stack must do: lines 1519 through 1521 compute per-block KV cost as block size times KV heads times head size times layers times two tensors times dtype width, a probe forward runs at 1548, and free memory divided by that cost yields the KV block count that was quoted back to the router. Tensor parallelism is visible as simple division here: attention heads and KV heads shrink by world size at lines 1281 and 1291, and each rank holds its slice of the cache. CUDA graphs are pre-captured per decode batch size at lines 1649 through 1665, and the speculative path at line 1906 threads draft token ids through the same forward call, with a drafter riding on the last rank just as in the stacks we compared at /vllm-2/ and /sglang-2/.

What to take home

Three takeaways generalize beyond this codebase. First, keeping the scheduler in the same process as the HTTP layer removes an entire serialization hop from the admission path, at the cost of re-implementing allocation logic in Rust. Second, prefix caching belongs next to the scheduler, because reuse decisions are allocation decisions; the radix trie and the batch budgets share one queue state precisely so they cannot disagree. Third, thin workers with a stateless protocol are a failure-isolation story, and the launcher’s exit-on-crash policy is the honest complement. If you want to see the alternative balance points, our tours of vLLM at /vllm/ and /vllm-2/, SGLang at /sglang/ and /sglang-2/, LMDeploy at /lmdeploy/, MLC-LLM at /mlc-llm/, and the single-file approach at /llamafile/ each moved a different piece across the same boundary. TGI’s answer, and the grpc service contract behind it in proto/generate.proto, is worth studying before you pick a side.

Enjoyed this post? Never miss out on future posts by following us