A GPU inference batching scheduler, written entirely by hand, with reference to the SGLang repository.
A fixed memory pool is provided to the scheduler. The scheduler checks for new requests and determines if memory is available for the new request. Each request first undergoes a prefill stage (FIFO order) which builds the KV cache for the input prompt, and generates the first output token. If there are no requests at the prefill stage, the requests then undergo a decode stage where the KV cached values are used to generate the next output token.
Allocating contiguous chunks sized for the maximum sequence length causes both internal fragmentation (reserved-but-unused space when a request finishes early) and external fragmentation (scattered free gaps too small for new contiguous allocations). We use PagedAttention which divides the memory pool into per-token slots, addressed via a per-request KV index array, eliminating both external and internal fragmentation.
Additionally, as more requests with long input prompts and output responses build, the memory pool may not be sufficient. In my implementation, I remove requests from the memory pool in the order of 1. shortest output length (least sunk cost), and 2. longest input length (to free the most memory). If any single request exceeds the total memory available, it is aborted.
Attention is a deterministic formula - the same tokens with the same preceding context always produce the same KV values. Maintaining a prefix cache enables us to reuse KV values that were generated by previously seen prompts. I implement this prefix cache using a Trie structure in trie.py.
I implement a flash-decoding kernel where multiple requests are batched into a single kernel launch, with each (request, head, KV-split) computed by an independent GPU program in parallel, then combined via a log-sum-exp merge across splits. Within each split, Flash Attention's online softmax avoids ever materialising the full attention score matrix, computing the running max and sum in a single pass instead. Requests with different sequence lengths are batched without padding — every request's KV context is concatenated into one flat buffer, with a per-(request, head) offset array telling each program where its own request's data begins.