How forge-infer works

Why. A contiguous KV buffer per sequence wastes memory two ways: internal waste from oversized buffers, and external fragmentation as different-sized sequences come and go. Paging eliminates external fragmentation entirely and bounds internal waste to block_size minus one slots per sequence.

How. Memory is split into fixed-size blocks. Each sequence carries a BlockTable of the physical blocks holding its tokens. Blocks are allocated lazily, one at a time. The free list is a stack so freshly freed blocks are reused first. The default cache is 512 blocks of 16 token slots each.

Why. The scheduler needs to ask whether a step will fit and act on the answer atomically. If append could half-succeed it would leak blocks under pressure, exactly when the engine can least afford it.

How. append pre-checks the free count against blocks_needed_for. If the cache cannot grow the sequence it returns CacheError::OutOfBlocks { needed, free } without touching state. The sequence is unchanged, no blocks have moved, and the scheduler can preempt and retry safely. out_of_blocks_is_reported_and_leaves_state_intact pins this behaviour.

Why. Pulling a block per token would waste allocator work. Pulling a block ahead would waste memory on sequences that finish early. Lazy growth pulls a new block only when the existing slack runs out.

How. slack = capacity - num_tokens, where capacity is blocks_held * block_size. If extra_tokens fits in slack, no new block is needed. Otherwise the additional blocks count is ceil((extra_tokens - slack) / block_size). The first token of a sequence pulls one block, the next fifteen reuse it, and the seventeenth pulls another.

Why. A continuous-batching engine wants to keep the running batch full. The admission loop is what lets a freshly arrived request join the very next decode iteration rather than waiting for a batch boundary.

How. While running.len() < max_batch_size and the next waiting request has tokens that fit, pull it in. Admission reserves blocks for the full current length, including any output a resumed sequence already produced. On a reservation failure the admit is rolled back via cache.free and admission stops for this step.

Why. Every running sequence needs one more block this step. The scheduler has to decide before the forward pass whether the whole batch fits, so it can preempt deterministically and not partway through.

How. Sum blocks_needed_for(seq, 1) across the running batch. Compare to the free block count. If the batch fits, reserve one decode block per sequence and add it to plan.decode_batch. If it does not, drop into preemption.

Why. When blocks run out, something has to give. The choice is which sequence to drop. Preempting the least-progressed sequence minimises the recompute cost when it resumes, the same intuition as shortest-remaining-time scheduling.

How. Pick the running sequence with the fewest output tokens, breaking ties towards the higher id for determinism. Free its blocks entirely via cache.free, set its state to Waiting, and push it back to the front of the waiting queue so it resumes before fresh requests. Repeat until the remaining batch fits. preempts_when_blocks_run_out and preempted_sequence_resumes_later pin this exact behaviour.

Why. vLLM offers two preemption paths: recompute and KV swap to host memory. forge-infer picks recompute. The reasoning is simple, the code is short, and the worst case is provably bounded.

How. When a sequence is preempted, its prompt plus output is re-prefilled when the cache can fit it again. No host memory pool, no copy path, no second eviction policy. Freeing any running sequence releases at least as many blocks as the next step needs, so the loop cannot deadlock. The cost is wasted recompute proportional to the preempted output length, which is why the policy preempts the least-progressed sequence.

Why. Sequences end in two ways: they hit max_new_tokens, or the model emits eos. Both paths need to free blocks promptly so the next admission has room.

How. After the decode batch is committed, walk the running set and retire anything past its limit. Free its blocks via cache.free and list it in plan.finished. push_token(id, token, true) caps the sequence limit to current output length so the next schedule call retires it through the same path eos_retires_a_sequence_early checks.

Why. A tick() that schedules and runs the forward pass together is shorter to write but impossible to test without a model. The whole policy belongs to the scheduler; the model belongs below the seam.

How. schedule() returns a StepPlan describing what was admitted, what to decode, and what was retired. It mutates only the cache and the queues. The engine consumes the plan: for each id in plan.decode_batch it calls model.forward, takes argmax, and pushes the token back via push_token. admission_blocks_when_prompt_does_not_fit asserts the policy with no model in the test at all.

Why. The engine is the thin layer that joins the scheduler to the model. It is intentionally boring; all the interesting decisions are above and below it.

How. Engine::step calls scheduler.schedule, iterates plan.decode_batch, calls model.forward over prompt plus output, takes argmax, checks against eos_token, and routes the token back through scheduler.push_token. The engine never decides what to run, and it never knows how the cache is laid out. It just connects them.

Why. Autoregressive generation is latency bound: each token needs one target forward pass, and you cannot start token n+1 before you have token n. If a cheap draft model can guess the next few tokens, the expensive target can verify them in one shot.

How. Pair the target with a small draft model. The draft proposes a run of k tokens via draft_run. The target verifies the whole run in one forward pass. Every accepted draft token is a free token. A fully accepted round of k drafts emits k + 1 tokens for one target step: the k accepted plus the bonus token the target sampled after the run.

Why. The subtlety is keeping the output exact. Speculative decoding must be indistinguishable from plain sampling from the target. Without the rejection-sampling test it would be an approximation, not an optimisation.

How. For each drafted token t with draft probability q(t) and target probability p(t): accept with probability min(1, p(t)/q(t)). On the first rejection, stop, resample that one position from the target distribution, and discard the rest of the draft. The acceptance draw is a seeded function of the context, not a thread RNG, so the test speculative_output_matches_plain_target_decoding can assert the speculatively decoded stream equals plain greedy target decoding token for token.

Why. The loop still has to make progress when the draft is wrong. Three failure modes need explicit handling: a rejected proposal, a draft of probability zero, and an early eos token mid-run.

How. On rejection, the rejected position is resampled from the target distribution and the rest of the draft is discarded, so the round still emits at least one token (rejection_falls_back_to_a_target_token). On q(t) <= 0 the ratio is undefined; the code clamps accept_prob to 1, the correct limit, which keeps the exactness test passing. On accepted or bonus eos the round returns immediately and generate stops the stream.

Why. A serving stack needs a tokeniser. The forge-infer one is byte-level on purpose: every input is encodable, every output decodable, no out-of-vocabulary path, no dependency on a model-specific BPE.

How. encode maps bytes to TokenIds. decode maps TokenIds back. The codec is a pure function of its input, so the same prompt produces the same token stream and the same output bytes every time, which is what the determinism guarantees depend on.

Why. A serving binary has to be callable from the same client code people are already using. The OpenAI completions shape is the de facto interface; matching it lets existing OpenAI client code point base_url at forge-infer and call /v1/completions unchanged.

How. axum exposes /generate (native shape) and /v1/completions (OpenAI-compatible). Both accept prompt and max_tokens. /v1/completions adds stream true, which switches the response to Server-Sent Events: one data: chunk per emitted token until the final data: [DONE]. The address is configurable through FORGE_ADDR.