How brave30/mini-inference-engine works
minfer is a PyTorch and Triton Llama inference engine. Its principal workflow accepts prompts, schedules continuous batches of decode tokens and prefill chunks, runs a flattened model forward over paged KV state, samples tokens, and returns generated outputs. The graph also shows weight/tokenizer acquisition, memory allocation and preemption, optimized attention/normalization, and CUDA-graph replay. The README documents these capabilities; sampled source excerpts are partial, so uncertain internal wiring is omitted.
Main entry points: engine.py LLM engine · model.py Llama forward · scheduler.py Batch scheduler · weights.py Weight loader
Components 17
Request API
- LLM engine
minfer/engine.py - Sampling params
minfer/sampling.py - Request output
minfer/engine.py
Batch Runtime
- Batch scheduler
minfer/scheduler.py - Sequence state
minfer/sequence.py - Page allocator
minfer/block_manager.py - CUDA graph replay
minfer/engine.py
Model Compute
- Llama forward
minfer/model.py - Prefill attention
minfer/model.py - Paged decode attention
minfer/kernels/decode_attention.py - Fused RMSNorm
minfer/kernels/rmsnorm.py
Weights And State
- Paged KV cache
minfer/model.py - Engine/model config
minfer/config.py - Weight loader
minfer/weights.py
Other components
- Inference caller
- HF tokenizer
- Model hub
Connections 21
- Inference caller → LLM engine: submits prompts
- Inference caller → Sampling params: sets parameters
- LLM engine → HF tokenizer: encodes and decodes
- LLM engine → Batch scheduler: adds and schedules
- Batch scheduler → Sequence state: tracks work
- Batch scheduler → Page allocator: reserves and frees
- LLM engine → Page allocator: allocates cache
- LLM engine → Engine/model config: reads settings
- LLM engine → Weight loader: loads model
- Weight loader → Model hub: downloads snapshots
- LLM engine → Llama forward: constructs and calls
- LLM engine → CUDA graph replay: captures or replays
- CUDA graph replay → Llama forward: runs decode forward
- Llama forward → Paged KV cache: writes and reads
- Llama forward → Paged decode attention: dispatches decode
- Llama forward → Prefill attention: dispatches prefill
- Llama forward → Fused RMSNorm: applies normalization
- LLM engine → Sampling params: samples logits
- LLM engine → Sequence state: updates tokens
- LLM engine → Request output: builds results
- LLM engine → Inference caller: returns outputs