# vrajpatel105/paged-inference-engine architecture

> Architecture diagram of the GitHub repository vrajpatel105/paged-inference-engine, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/vrajpatel105/paged-inference-engine
- Repository: https://github.com/vrajpatel105/paged-inference-engine

## Overview

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: [`core/model_runner.py`](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/model_runner.py) (Model runner), [`demo/api.py`](https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/api.py) (Demo API), [`transformer/model.py`](https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/model.py) (Decoder model), [`core/scheduler.py`](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/scheduler.py) (Batch scheduler)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_serving["Serving loop"]
  node_engine["Engine entry<br/>[engine.py]"]
  node_runner["Model runner<br/>[model_runner.py]"]
  node_scheduler["Batch scheduler<br/>[scheduler.py]"]
  node_sequence["Sequence state<br/>[sequence.py]"]
end

subgraph group_cache["Paged execution"]
  node_allocator["Block allocator<br/>[block_manager.py]"]
  node_attention["FlashAttention kernel<br/>[flash_attention.py]"]
end

subgraph group_model["Model lifecycle"]
  node_model["Decoder model<br/>[model.py]"]
  node_tokenizer["Tokenizer<br/>[tokenizer.py]"]
  node_training["Model training<br/>[train.py]"]
  node_checkpoint["Model checkpoint"]
  node_loader["Checkpoint loader<br/>[load_checkpoint.py]"]
end

subgraph group_demo["Demo interface"]
  node_api["Demo API<br/>[api.py]"]
  node_web["Demo page<br/>[index.html]"]
  node_output["Generated results<br/>[api.py]"]
end

node_client(("Prompt user"))
node_dataset(("TinyStories"))

node_client -->|"submits prompt"| node_api
node_api -->|"encodes prompt"| node_tokenizer
node_api -->|"queues request"| node_runner
node_engine -->|"starts loop"| node_runner
node_runner -->|"schedules step"| node_scheduler
node_runner -->|"runs forward"| node_model
node_runner -->|"decodes results"| node_tokenizer
node_scheduler -->|"tracks requests"| node_sequence
node_scheduler -->|"allocates blocks"| node_allocator
node_model -->|"dispatches attention"| node_attention
node_model -->|"reads block table"| node_allocator
node_dataset -.->|"provides text"| node_training
node_training -->|"builds vocabulary"| node_tokenizer
node_training -->|"writes weights"| node_checkpoint
node_checkpoint -->|"loads weights"| node_loader
node_loader -->|"builds model"| node_model
node_api -->|"updates results"| node_output
node_client -->|"opens demo"| node_web
node_web -->|"requests output"| node_api

click node_engine "https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/engine.py"
click node_runner "https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/model_runner.py"
click node_scheduler "https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/scheduler.py"
click node_sequence "https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/sequence.py"
click node_allocator "https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/block_manager.py"
click node_model "https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/model.py"
click node_attention "https://github.com/vrajpatel105/paged-inference-engine/blob/main/kernels/flash_attention.py"
click node_tokenizer "https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/tokenizer.py"
click node_training "https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/train.py"
click node_loader "https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/load_checkpoint.py"
click node_api "https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/api.py"
click node_web "https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/static/index.html"
click node_output "https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/api.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_engine,node_runner,node_scheduler,node_sequence,node_client toneBlue
class node_allocator,node_attention toneAmber
class node_model,node_tokenizer,node_training,node_checkpoint,node_loader toneMint
class node_api,node_web,node_output toneRose
class node_dataset toneIndigo
```

## Components

### Serving loop

- [Engine entry](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/engine.py): `core/engine.py`
- [Model runner](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/model_runner.py): `core/model_runner.py`
- [Batch scheduler](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/scheduler.py): `core/scheduler.py`
- [Sequence state](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/sequence.py): `core/sequence.py`

### Paged execution

- [Block allocator](https://github.com/vrajpatel105/paged-inference-engine/blob/main/core/block_manager.py): `core/block_manager.py`
- [FlashAttention kernel](https://github.com/vrajpatel105/paged-inference-engine/blob/main/kernels/flash_attention.py): `kernels/flash_attention.py`

### Model lifecycle

- [Decoder model](https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/model.py): `transformer/model.py`
- [Tokenizer](https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/tokenizer.py): `transformer/tokenizer.py`
- [Model training](https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/train.py): `transformer/train.py`
- **Model checkpoint**
- [Checkpoint loader](https://github.com/vrajpatel105/paged-inference-engine/blob/main/transformer/load_checkpoint.py): `transformer/load_checkpoint.py`

### Demo interface

- [Demo API](https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/api.py): `demo/api.py`
- [Demo page](https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/static/index.html): `demo/static/index.html`
- [Generated results](https://github.com/vrajpatel105/paged-inference-engine/blob/main/demo/api.py): `demo/api.py`

### Other components

- **Prompt user**
- **TinyStories**

## Connections

- 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