Skip to content

Latest commit

 

History

2 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

paged-infer-engine

A from-scratch, high-throughput LLM inference engine implementing:

  • Continuous (iteration-level) batching — new requests admitted / completed requests evicted every decode step without waiting for the whole batch.
  • Paged KV-cache — fixed-size physical block pool with a per-sequence logical→physical page table, copy-on-write reference counting, and O(1) alloc/free.
  • PagedAttention kernel — Triton (primary) and optional raw CUDA variant; gathers KV data from non-contiguous physical blocks in-kernel, no compaction copy.
  • Speculative decoding — draft model generates K candidate tokens; target model scores them in one batched forward pass; rejection sampling maintains distributional correctness; KV-cache rollback on rejection uses Phase 1 copy-on-write.
  • gRPC + WebSocket ingress — server-streaming gRPC for production clients, WebSocket gateway for browser demos.
  • Prometheus telemetry — per-request latency histograms, tokens/sec, KV-cache occupancy, fragmentation ratio, speculative acceptance rate.

Tech Stack

Layer Technology
Core runtime C++17/20
Kernels Triton (primary), CUDA (optional benchmark variant)
Model execution Python + PyTorch
C++↔Python bridge pybind11
Serving gRPC (server-streaming), WebSocket
Telemetry Prometheus client
Build (C++) CMake ≥ 3.20
Build (Python) Poetry / pyproject.toml
Dev environment Docker + docker-compose

Target Hardware

Single-node, 1–2 NVIDIA GPUs (Ampere or newer). The memory abstraction layer is designed so multi-GPU tensor parallelism is a Phase 6+ extension, not a rewrite.

Phase Roadmap

Phase Module Status
0 Repository scaffolding + dev environment 🔄 In Progress
1 Block allocator & virtual page table ⏳ Pending
2 PagedAttention kernel (Triton/CUDA) ⏳ Pending
3 Continuous batching engine & priority scheduler ⏳ Pending
4 Speculative decoding pipeline ⏳ Pending
5 Ingress server, telemetry & benchmarking suite ⏳ Pending

Quick Start

# Build the dev container
docker-compose -f docker/docker-compose.yml up --build -d

# Enter the container
docker exec -it paged-infer-engine-dev bash

# Verify GPU + framework stack
nvidia-smi
python -c "import torch, triton; print(f'CUDA device: {torch.cuda.get_device_name(0)}')"

# Configure and build C++ core (inside container or with CUDA toolkit installed)
bash scripts/build.sh

# Run all unit tests
bash scripts/run_unit_tests.sh

# Run benchmarks
bash scripts/run_benchmarks.sh

Directory Structure

paged-infer-engine/
├── CMakeLists.txt          # C++/CUDA build
├── pyproject.toml          # Python package + dependencies
├── docker/                 # Dev container
├── proto/                  # gRPC proto definitions
├── src/
│   ├── memory/             # Phase 1: Block allocator + page table
│   ├── kernels/            # Phase 2: PagedAttention Triton/CUDA kernels
│   ├── scheduler/          # Phase 3: Continuous batcher + priority queue
│   ├── speculative/        # Phase 4: Draft/target/rejection sampler
│   ├── serving/            # Phase 5: gRPC + WebSocket servers
│   ├── telemetry/          # Prometheus metrics exporters
│   ├── bindings/           # pybind11 C++↔Python bridge
│   └── common/             # Shared config, logging, types
├── bench/                  # Standalone benchmarks + reports
├── tests/
│   ├── unit/
│   └── integration/
└── scripts/                # build.sh, run_unit_tests.sh, run_benchmarks.sh

License

MIT

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages