# infini-ai-lab/magicpig architecture

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

- Interactive diagram: https://gitdiagram.com/infini-ai-lab/magicpig
- Repository: https://github.com/infini-ai-lab/magicpig

## Overview

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: [`evaluations/RULER/data/prepare.py`](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/prepare.py) (Data preparation), [`evaluations/RULER/eval/evaluate.py`](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/eval/evaluate.py) (RULER evaluator), [`evaluations/RULER/pred/llama_ref.py`](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_ref.py) (Single-GPU Llama), [`evaluations/RULER/pred/llama_dist.py`](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_dist.py) (Distributed Llama)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_g_lsh["LSH Sampling"]
  node_lsh["LSH kernel<br/>[lsh.cc]"]
  node_lsh_api["LSH interface<br/>[lsh.h]"]
end

subgraph group_g_ruler_data["Benchmark Data"]
  node_prepare["Data preparation<br/>[prepare.py]"]
  node_template["Prompt templates<br/>[template.py]"]
  node_tokenizer["Tokenizer<br/>[tokenizer.py]"]
  node_niah["Needle tasks<br/>[niah.py]"]
  node_qa["Synthetic QA<br/>[qa.py]"]
  node_variable["Variable tracking"]
end

subgraph group_g_predict["Prediction Backends"]
  node_llama_ref["Single-GPU Llama<br/>[llama_ref.py]"]
  node_llama_dist["Distributed Llama<br/>[llama_dist.py]"]
  node_mistral["Mistral backend<br/>[mistral_ref.py]"]
  node_quest["Quest baseline<br/>[llama_quest.py]"]
  node_topk["TopK baseline<br/>[llama_topk.py]"]
  node_oracle["Oracle sampling"]
  node_api["API backend<br/>[call_api.py]"]
  node_client["Model clients<br/>[client_wrappers.py]"]
  node_model_wrappers["Model wrappers<br/>[model_wrappers.py]"]
end

subgraph group_g_evaluate["Evaluation"]
  node_eval["RULER evaluator<br/>[evaluate.py]"]
end

node_user(("Engineer"))
node_external_model{{"Model providers"}}

node_user -->|"selects tasks"| node_prepare
node_prepare -->|"builds examples"| node_niah
node_prepare -->|"builds examples"| node_qa
node_prepare -->|"builds examples"| node_variable
node_prepare -->|"formats prompts"| node_template
node_prepare -->|"tokenizes data"| node_tokenizer
node_user -->|"runs benchmark"| node_llama_ref
node_user -->|"runs benchmark"| node_llama_dist
node_user -->|"runs benchmark"| node_mistral
node_user -->|"selects baseline"| node_quest
node_user -->|"selects baseline"| node_topk
node_user -->|"selects baseline"| node_oracle
node_user -->|"selects backend"| node_api
node_llama_ref -->|"uses wrappers"| node_model_wrappers
node_llama_dist -->|"uses wrappers"| node_model_wrappers
node_api -->|"uses clients"| node_client
node_client -.->|"requests inference"| node_external_model
node_llama_ref -->|"produces predictions"| node_eval
node_llama_dist -->|"produces predictions"| node_eval
node_mistral -->|"produces predictions"| node_eval
node_quest -->|"produces predictions"| node_eval
node_topk -->|"produces predictions"| node_eval
node_oracle -->|"produces predictions"| node_eval
node_api -->|"produces predictions"| node_eval
node_lsh_api -->|"exposes kernel"| node_lsh

click node_lsh "https://github.com/infini-ai-lab/magicpig/blob/v0.2/library/lsh/lsh.cc"
click node_lsh_api "https://github.com/infini-ai-lab/magicpig/blob/v0.2/library/lsh/lsh.h"
click node_prepare "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/prepare.py"
click node_template "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/template.py"
click node_tokenizer "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/tokenizer.py"
click node_niah "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/synthetic/niah.py"
click node_qa "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/synthetic/qa.py"
click node_variable "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/synthetic/variable_tracking.py"
click node_eval "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/eval/evaluate.py"
click node_llama_ref "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_ref.py"
click node_llama_dist "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_dist.py"
click node_mistral "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/mistral_ref.py"
click node_quest "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_quest.py"
click node_topk "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_topk.py"
click node_oracle "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/oraclesampling_cache.py"
click node_api "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/call_api.py"
click node_client "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/client_wrappers.py"
click node_model_wrappers "https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/model_wrappers.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_lsh,node_lsh_api toneBlue
class node_prepare,node_template,node_tokenizer,node_niah,node_qa,node_variable toneAmber
class node_llama_ref,node_llama_dist,node_mistral,node_quest,node_topk,node_oracle,node_api,node_client,node_model_wrappers toneMint
class node_eval toneRose
class node_user,node_external_model toneIndigo
```

## Components

### LSH Sampling

- [LSH kernel](https://github.com/infini-ai-lab/magicpig/blob/v0.2/library/lsh/lsh.cc): `library/lsh/lsh.cc`
- [LSH interface](https://github.com/infini-ai-lab/magicpig/blob/v0.2/library/lsh/lsh.h): `library/lsh/lsh.h`

### Benchmark Data

- [Data preparation](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/prepare.py): `evaluations/RULER/data/prepare.py`
- [Prompt templates](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/template.py): `evaluations/RULER/data/template.py`
- [Tokenizer](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/tokenizer.py): `evaluations/RULER/data/tokenizer.py`
- [Needle tasks](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/synthetic/niah.py): `evaluations/RULER/data/synthetic/niah.py`
- [Synthetic QA](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/synthetic/qa.py): `evaluations/RULER/data/synthetic/qa.py`
- [Variable tracking](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/data/synthetic/variable_tracking.py): `evaluations/RULER/data/synthetic/variable_tracking.py`

### Prediction Backends

- [Single-GPU Llama](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_ref.py): `evaluations/RULER/pred/llama_ref.py`
- [Distributed Llama](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_dist.py): `evaluations/RULER/pred/llama_dist.py`
- [Mistral backend](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/mistral_ref.py): `evaluations/RULER/pred/mistral_ref.py`
- [Quest baseline](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_quest.py): `evaluations/RULER/pred/llama_quest.py`
- [TopK baseline](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/llama_topk.py): `evaluations/RULER/pred/llama_topk.py`
- [Oracle sampling](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/oraclesampling_cache.py): `evaluations/RULER/pred/oraclesampling_cache.py`
- [API backend](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/call_api.py): `evaluations/RULER/pred/call_api.py`
- [Model clients](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/client_wrappers.py): `evaluations/RULER/pred/client_wrappers.py`
- [Model wrappers](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/pred/model_wrappers.py): `evaluations/RULER/pred/model_wrappers.py`

### Evaluation

- [RULER evaluator](https://github.com/infini-ai-lab/magicpig/blob/v0.2/evaluations/RULER/eval/evaluate.py): `evaluations/RULER/eval/evaluate.py`

### Other components

- **Engineer**
- **Model providers**

## Connections

- 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