How vrajpatel105/paged-inference-engine works
Paged Inference Engine serves tokenized prompts through a continuous-batching loop and a decoder using paged INT8 KV storage and a Triton attention kernel. Training and checkpoint conversion form a separate model-preparation branch; the FastAPI demo is an optional client-facing entry point. The README documents these capabilities, while sampled source confirms the scheduler, allocator, kernel, API, and training behavior; some wiring inside the unsampled portions of the runner and model is therefore omitted.
Main entry points: model_runner.py Model runner · api.py Demo API · model.py Decoder model · scheduler.py Batch scheduler
Components 16
Serving loop
- Engine entry
core/engine.py - Model runner
core/model_runner.py - Batch scheduler
core/scheduler.py - Sequence state
core/sequence.py
Paged execution
- Block allocator
core/block_manager.py - FlashAttention kernel
kernels/flash_attention.py
Model lifecycle
- Decoder model
transformer/model.py - Tokenizer
transformer/tokenizer.py - Model training
transformer/train.py - Model checkpoint
- Checkpoint loader
transformer/load_checkpoint.py
Demo interface
- Demo API
demo/api.py - Demo page
demo/static/index.html - Generated results
demo/api.py
Other components
- Prompt user
- TinyStories
Connections 19
- Prompt user → Demo API: submits prompt
- Demo API → Tokenizer: encodes prompt
- Demo API → Model runner: queues request
- Engine entry → Model runner: starts loop
- Model runner → Batch scheduler: schedules step
- Model runner → Decoder model: runs forward
- Model runner → Tokenizer: decodes results
- Batch scheduler → Sequence state: tracks requests
- Batch scheduler → Block allocator: allocates blocks
- Decoder model → FlashAttention kernel: dispatches attention
- Decoder model → Block allocator: reads block table
- TinyStories → Model training: provides text
- Model training → Tokenizer: builds vocabulary
- Model training → Model checkpoint: writes weights
- Model checkpoint → Checkpoint loader: loads weights
- Checkpoint loader → Decoder model: builds model
- Demo API → Generated results: updates results
- Prompt user → Demo page: opens demo
- Demo page → Demo API: requests output