How infini-ai-lab/magicpig works
MagicPIG accelerates long-context LLM generation by selecting attention work with locality-sensitive hashing, with CPU sparse attention complementing GPU attention. Its documented user workflow is to provide a model and prompt or benchmark data, generate or evaluate predictions, and compare accuracy and throughput. The tree includes RULER data preparation, prediction backends, and evaluation, plus the LSH source and vendored FBGEMM components; it does not include the documented generation entry point or the main sparse-attention implementation, so those connections are left out. The README documents evaluation modes, but sampled source does not establish their internal call wiring.
Main entry points: prepare.py Data preparation · evaluate.py RULER evaluator · llama_ref.py Single-GPU Llama · llama_dist.py Distributed Llama
Components 20
LSH Sampling
- LSH kernel
library/lsh/lsh.cc - LSH interface
library/lsh/lsh.h
Benchmark Data
- Data preparation
evaluations/RULER/data/prepare.py - Prompt templates
evaluations/RULER/data/template.py - Tokenizer
evaluations/RULER/data/tokenizer.py - Needle tasks
evaluations/RULER/data/synthetic/niah.py - Synthetic QA
evaluations/RULER/data/synthetic/qa.py - Variable tracking
evaluations/RULER/data/synthetic/variable_tracking.py
Prediction Backends
- Single-GPU Llama
evaluations/RULER/pred/llama_ref.py - Distributed Llama
evaluations/RULER/pred/llama_dist.py - Mistral backend
evaluations/RULER/pred/mistral_ref.py - Quest baseline
evaluations/RULER/pred/llama_quest.py - TopK baseline
evaluations/RULER/pred/llama_topk.py - Oracle sampling
evaluations/RULER/pred/oraclesampling_cache.py - API backend
evaluations/RULER/pred/call_api.py - Model clients
evaluations/RULER/pred/client_wrappers.py - Model wrappers
evaluations/RULER/pred/model_wrappers.py
Evaluation
- RULER evaluator
evaluations/RULER/eval/evaluate.py
Other components
- Engineer
- Model providers
Connections 25
- Engineer → Data preparation: selects tasks
- Data preparation → Needle tasks: builds examples
- Data preparation → Synthetic QA: builds examples
- Data preparation → Variable tracking: builds examples
- Data preparation → Prompt templates: formats prompts
- Data preparation → Tokenizer: tokenizes data
- Engineer → Single-GPU Llama: runs benchmark
- Engineer → Distributed Llama: runs benchmark
- Engineer → Mistral backend: runs benchmark
- Engineer → Quest baseline: selects baseline
- Engineer → TopK baseline: selects baseline
- Engineer → Oracle sampling: selects baseline
- Engineer → API backend: selects backend
- Single-GPU Llama → Model wrappers: uses wrappers
- Distributed Llama → Model wrappers: uses wrappers
- API backend → Model clients: uses clients
- Model clients → Model providers: requests inference
- Single-GPU Llama → RULER evaluator: produces predictions
- Distributed Llama → RULER evaluator: produces predictions
- Mistral backend → RULER evaluator: produces predictions
- Quest baseline → RULER evaluator: produces predictions
- TopK baseline → RULER evaluator: produces predictions
- Oracle sampling → RULER evaluator: produces predictions
- API backend → RULER evaluator: produces predictions
- LSH interface → LSH kernel: exposes kernel