Metadata-Version: 2.5
Name: economize
Version: 0.3.1
Summary: Composable efficiency techniques (cross-layer jump drafting, early exit, speculative decoding, ...) for Hugging Face causal LMs.
Project-URL: Homepage, https://github.com/Thiggel/economize
Project-URL: Repository, https://github.com/Thiggel/economize
Project-URL: Issues, https://github.com/Thiggel/economize/issues
Project-URL: Documentation, https://github.com/Thiggel/economize/tree/main/docs
Author-email: Filipe Laitenberger <f.laitenberger@gmail.com>
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: early-exit,efficiency,huggingface,inference,kv-cache,layer-skipping,llm,pruning,quantization,speculative-decoding,transformers,vllm
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.10
Requires-Dist: huggingface-hub>=1.0
Requires-Dist: safetensors>=0.4
Requires-Dist: torch>=2.4
Requires-Dist: transformers>=5.0
Provides-Extra: all
Requires-Dist: accelerate>=1.14; extra == 'all'
Requires-Dist: compressed-tensors>=0.19; extra == 'all'
Requires-Dist: datasets>=5.0; extra == 'all'
Requires-Dist: kernels>=0.17; extra == 'all'
Requires-Dist: lm-eval[hf]>=0.4.13; extra == 'all'
Requires-Dist: peft>=0.21; extra == 'all'
Requires-Dist: torchao>=0.18; extra == 'all'
Requires-Dist: trl>=1.14; extra == 'all'
Provides-Extra: bench
Requires-Dist: datasets>=5.0; extra == 'bench'
Provides-Extra: dev
Requires-Dist: mypy==2.4.0; extra == 'dev'
Requires-Dist: pytest>=8; extra == 'dev'
Requires-Dist: ruff==0.16.10; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-material<10,>=9.6; extra == 'docs'
Provides-Extra: eval
Requires-Dist: lm-eval[hf]>=0.4.13; extra == 'eval'
Provides-Extra: kernels
Requires-Dist: kernels>=0.17; extra == 'kernels'
Provides-Extra: peft
Requires-Dist: peft>=0.21; extra == 'peft'
Provides-Extra: quant
Requires-Dist: compressed-tensors>=0.19; extra == 'quant'
Requires-Dist: torchao>=0.18; extra == 'quant'
Provides-Extra: sglang
Requires-Dist: sglang; extra == 'sglang'
Provides-Extra: test
Requires-Dist: pytest>=8; extra == 'test'
Provides-Extra: train
Requires-Dist: accelerate>=1.14; extra == 'train'
Requires-Dist: datasets>=5.0; extra == 'train'
Requires-Dist: trl>=1.14; extra == 'train'
Provides-Extra: vllm
Requires-Dist: vllm>=0.30; extra == 'vllm'
Description-Content-Type: text/markdown

# economize

Composable efficiency techniques for Hugging Face causal language models.

economize attaches techniques to an unmodified `transformers` model. The base model keeps its module tree and its
`forward`. Drafters and other added components bring their own parameters and act only while decoding; weight
transforms (pruning, low rank, quantization) rewrite the base weights in place, once, when they attach. Any set of
registered techniques can be combined in one call, and every architecture adapter, technique, verifier, cache and
checkpoint format is added by registration.

Version 0.3 ships these techniques, on 46 model types (dense, MoE, and hybrid attention plus Mamba /
linear-attention models), with batched decoding of ragged prompts, draft trees and lossless speculative sampling:

| Family | Techniques (registered name) |
|---|---|
| Drafters for speculative decoding | cross-layer jump drafting (`jump`), EAGLE-3 (`eagle3`), Medusa heads and trees (`medusa`), early exit / LayerSkip (`early_exit`), training-free layer skipping (`layer_skip`), a second draft model (`draft_model`), prompt lookup (`lookup`), FR-Spec draft vocabulary (`frspec`), lookup with a learned fallback (`hybrid`), the MTP heads that DeepSeek-V3, GLM-4.5, MiMo and Qwen3.5 / Qwen3-Next ship (`vendor_mtp`) |
| Dynamic depth | CALM (`calm`), GateSkip (`gateskip`), Mixture-of-Depths (`mod`), FlexiDepth (`flexidepth`), SkipDecode (`skipdecode`), LoRA-Drop (`loradrop`) |
| Weights and activations | ShortGPT / SLEB depth pruning (`depth_prune`), magnitude / Wanda / SparseGPT pruning incl. 2:4 (`prune`), SVD / SVD-LLM low rank (`lowrank`), RTN / GPTQ / AWQ quantization to int2-int8 or fp8 with fake, torchao, Marlin or compressed-tensors kernels (`quant`), draft-only quantization (`draft_quant`), TEAL activation sparsity (`teal`) |
| KV cache | KIVI (`kivi`), FP8 KV (`kv_fp8`), StreamingLLM / H2O / SnapKV eviction (`kv_evict`) |

[docs/examples.md](https://github.com/Thiggel/economize/blob/main/docs/examples.md) has a runnable example for
every technique and feature, and [examples/](https://github.com/Thiggel/economize/tree/main/examples) has complete
scripts that prepare data, train, compress, combine, benchmark and serve; `mkdocs serve` in a clone shows all of
`docs/` as a site.

## Install

```bash
pip install economize                       # torch>=2.4, transformers>=5.0, safetensors, huggingface_hub
pip install "economize[all]"                # every optional integration below except vLLM and SGLang
```

From source:

```bash
git clone https://github.com/Thiggel/economize.git
pip install -e "./economize[dev]"
```

| Extra | Installs | For |
|---|---|---|
| `test` | pytest | the test suite |
| `peft` | peft | PEFT-wrapped base models and LoRA training of the base |
| `quant` | torchao, compressed-tensors | quantized kernels, compressed-tensors export |
| `kernels` | kernels | Hugging Face `kernels` hub kernels (`economize.integrations.kernels`) |
| `train` | accelerate, datasets, trl | drafter training and online RL rollouts |
| `eval` | lm-eval[hf] | `economize-lm-eval`, lm-evaluation-harness with economize's decoding loop |
| `bench` | datasets | prompt suites of `economize.bench` (Spec-Bench, GSM8K, HumanEval) |
| `vllm`, `sglang` | vllm, sglang | serving backends (`economize.backends`; drafters are written by `economize.backends.export_drafter`) |

Installing the package registers the vLLM plugin (`vllm.general_plugins`), the compressed-tensors quantization
backend and the `economize-lm-eval` command through entry points.

## Quick start

```python
import torch
import economize
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B", dtype=torch.bfloat16, device_map="cuda")
model = economize.load(base, "path/to/drafter")            # a trained drafter: a folder or a Hub repo id
# or in one step, loading the base model named in the checkpoint:
model = economize.from_pretrained("path/to/drafter", dtype=torch.bfloat16, device_map="cuda")

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
ids = tokenizer("The water cycle describes", return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=128)             # lossless: the same tokens as base.generate
```

`model.generate` accepts the usual `transformers` arguments (logits processors such as `repetition_penalty`,
`eos_token_id`, `return_dict_in_generate`, ...) and routes them through economize's decoding loop. `model.decode`
returns the tokens together with a log of every model call and the acceptance of each cycle:

```python
result = model.decode(ids[0], 128, decoding=economize.DecodeConfig(cycle_length=4))
result.cost_summary()      # exact passes, draft steps, layer applications relative to plain decoding
result.mean_accepted       # accepted drafts per cycle
```

## Results

Every speedup below is lossless: greedy decoding returns the same tokens as plain greedy decoding (up to the
rounding noise of batched bf16 kernels; the decoders are exactly lossless in fp32 tests), and sampling follows the
model's own distribution. All runs use batch size 1, and a speedup divides the decoding throughput by that of plain
decoding in the same runtime.

**Cross-layer jump drafting against other drafters.** Qwen2.5-0.5B in bf16 on one NVIDIA A40, static KV cache and
CUDA graphs. Every learned drafter is trained on the same 64M-token mix of the model's own answers to task prompts
for 12000 steps, and every drafter drafts from the 32768 most frequent tokens with a confidence-based draft length,
tuned per prompt set. These numbers come from the research decoder of the accompanying paper (`tskip`), which
economize's oracle tests match token for token. Mean ± standard deviation over three training seeds.

| Drafter | Seeds | Spec-Bench | HumanEval | GSM8K |
|---|---|---|---|---|
| Jump + exit (`jump`: m=16, u=17, r=17, linear predictor, full copy of block 17) | 3 | **2.05 ± 0.02** | **1.81 ± 0.01** | **2.21 ± 0.04** |
| EAGLE-3 (`eagle3`), same data and budget | 3 | 1.89 ± 0.07 | 1.63 ± 0.01 | 1.89 ± 0.02 |
| Medusa, 5 heads with their own LM heads (`medusa`, `head="own"`) | 3 | 1.28 ± 0.01 | 1.18 ± 0.02 | 1.58 ± 0.02 |
| Medusa, 5 heads sharing the LM head (`medusa`) | 3 | 1.27 ± 0.02 | 1.11 ± 0.01 | 1.39 ± 0.01 |
| Early exit after block 3, frozen blocks (`early_exit`) | 1 | 1.19 | 0.94 | 1.22 |
| Prompt lookup (`lookup`, no training, one timing run) | - | 2.04 | 2.50 | 1.13 |
| Training-free layer skipping, 10 of 24 blocks (`layer_skip`, one timing run) | - | 0.98 | 0.94 | 0.80 |
| Jump + exit with prompt lookup (`hybrid`) | 2 | 2.62 ± 0.01 | 2.87 ± 0.01 | 2.17 ± 0.01 |

With sampling (`verifier="rejection"`, three seeds each) the jump + exit drafter reaches 1.63 ± 0.05, 1.36 ± 0.06
and 1.85 ± 0.03 at T=0.7 against 1.51 ± 0.10, 1.25 ± 0.03 and 1.64 ± 0.04 for EAGLE-3, and 1.21 ± 0.01,
1.17 ± 0.01 and 1.66 ± 0.00 at T=1.0 against 1.10 ± 0.01, 1.07 ± 0.04 and 1.46 ± 0.02. On GSM8K the hybrid never
finds a useful lookup, so its best setting there is the drafter alone. Rows with one seed or one timing run have no
spread and are not compared with the rest at the 0.1 level.

**Reproductions with economize.** Published checkpoints decoded by economize itself on an NVIDIA A100, one timing
run each, against plain decoding. [docs/REPRODUCTION.md](https://github.com/Thiggel/economize/blob/main/docs/REPRODUCTION.md)
gives the protocols, the quality reproductions of the lossy techniques (GPTQ, AWQ, Wanda, SparseGPT, SLEB,
ShortGPT, SVD-LLM, TEAL, KIVI, KV eviction) and every deviation.

| Technique | Model and checkpoint | Benchmark | economize | Published |
|---|---|---|---|---|
| `eagle3`, chain of 5 | Llama-3.1-8B-Instruct, `yuhuili/EAGLE3-LLaMA3.1-Instruct-8B` | MT-bench, HumanEval, GSM8K, Alpaca, CNN/DM | 2.82x (3.73 tokens per cycle) | 4.44x with a draft tree |
| `eagle3`, chain of 3 | Llama-3.1-8B-Instruct, `RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3` | HumanEval / math / summarization | 1.96x / 1.89x / 1.75x | no speedup published |
| `medusa`, tree `mc_sim_7b_63` | Vicuna-7B-v1.3, `FasterDecoding/medusa-vicuna-7b-v1.3` | Spec-Bench | 1.51x | 1.71x |
| `draft_model` | Vicuna-7B-v1.3 with `double7/vicuna-68m` | Spec-Bench | 1.58x | 1.52x |
| `lookup` | Vicuna-7B-v1.3 | Spec-Bench | 1.50x | no speedup published |

The EAGLE-3 rows decode chains, while the paper's 4.44x uses draft trees, and its baseline is plain `transformers`
decoding. The published speedups were measured on other machines.

## Jump drafting

A draft step for the next token predicts its hidden state after layer `m` from the current token's state after
layer `u` and the next token's embedding, runs only layers `m+1..r` and reads out logits. Layers `1..m` (and
`r+1..L`) are skipped. One *cycle* of `N = cycle_length` tokens runs the unmodified model once, in parallel, over
every token not yet computed exactly, which also yields one exact token, followed by up to `N-1` draft steps.

```python
config = economize.JumpConfig(
    target_layer=12,            # m
    source_layer=18,            # u
    exit_layer=None,            # r (default: the last layer; below it a readout adapter is added)
    adaptation="frozen",        # or "lora" (LoRA active only on the jump path) or "full" (trainable layer copies)
)
model = economize.apply(base, config)
list(model.technique_parameters())   # the added parameters only; the base model is untouched
```

`DecodeConfig` chooses how drafts are used:

| Option | Meaning |
|---|---|
| `cycle_length` | Tokens per cycle `N` (`1` is ordinary decoding). |
| `lossless=True` | Speculative decoding: drafts are verified by the exact pass, the output equals greedy decoding. |
| `lossless=False` | Every draft is kept and the exact pass only refreshes the cache (lossy, cheaper). |
| `min_confidence` | Stop drafting in a cycle once the drafter's top probability falls below this value. |
| `verifier` | A registered verifier by name: `"greedy"`, `"rejection"` (lossless speculative sampling at T > 0, chains and trees), `"typical"` (Medusa's typical acceptance, lossy) or `"accept_all"`. |
| `verifier_options` | Keyword options of the verifier, e.g. `{"threshold": 0.09, "alpha": 0.3}` for `"typical"`. |
| `cuda_graphs` | Replay the exact passes and the draft steps from CUDA graphs (see below). |
| `sync_interval` | Ordinary decoding reads tokens back and checks stop conditions every this many steps (default 8; the output is the same as reading every step). |

Other conditioning variants are part of `JumpConfig`: several source layers (`extra_source_layers`) combined by a
registered `mixer` (`"concat"`, `"fused_linear"`, `"layer_attention"`), the predictor (`"mlp"`, `"linear"`, its
width and depth, or `"flow"` / `"flow_linear"`: that mean followed by a few Euler steps of a rectified flow, trained by
flow matching, with `flow_steps`, `flow_sigma`, `flow_width`), and ablations (`use_state`, `use_token`,
`token_residual`).

Several prompts of different lengths decode as one batch, each row accepting its own number of drafts per cycle.
With `verifier="rejection"` and a sampling selector the output follows the model's own sampling distribution
(speculative sampling):

```python
results = model.decode_batch([ids_a, ids_b, ids_c], 128)          # one DecodeResult per prompt
sampled = model.decode_batch(
    [ids_a, ids_b], 128,
    decoding=economize.DecodeConfig(cycle_length=4, verifier="rejection"),
    selector=economize.TopPSampler(temperature=0.7, top_p=0.9, seed=0, device="cuda"),
)
```

Each row may also have its own selector, stop condition and token limit (`selectors=`, `stops=`, and a list for
`max_new_tokens`). Rows that share a selector object are drafted and verified together.

At batch size 1 a draft step and a short exact pass are bound by kernel launches. With
`DecodeConfig(cuda_graphs=True)` the decoder records its exact chain passes and the jump drafter's steps as CUDA
graphs on a static KV cache and replays them. The model keeps one decoder per configuration, so the graphs recorded
in the first generation serve the following ones; `model.refresh()` drops them after the weights or the techniques
change. Draft trees, dynamic-depth policies, KV-cache techniques and models with recurrent layers decode eagerly
whatever the flag says.

## Combining techniques

`economize.apply` takes any number of configurations and attaches them in a fixed order, whatever the order of the
arguments: weight transforms (by their `order`: depth pruning, pruning, low rank, quantization), activation hooks,
KV-cache wrappers, dynamic depth, the drafter, then the rest. Techniques that need data calibrate on the model as it
is at that point, so quantization sees the pruned weights.

```python
from economize.techniques.prune import PruneConfig
from economize.techniques.quant import QuantConfig
from economize.techniques.sparsity import TealConfig

model = economize.apply(
    base,
    PruneConfig(method="wanda", sparsity="2:4"),
    QuantConfig(format="int4", method="gptq", group_size=128),
    TealConfig(sparsity=0.4),
    economize.JumpConfig(target_layer=12, source_layer=18),
    calibration=calibration_ids,          # LongTensor [N, S], an iterable of batches, or a callable
)
```

See [docs/weights.md](https://github.com/Thiggel/economize/blob/main/docs/weights.md) for every weight and activation technique and its options.

## Training

`economize.train` trains what the techniques add (drafters, readout adapters, LoRA on a drafter's path, early-exit
heads) and, where a technique needs it, the base model:

```python
from economize.train import fit

model = economize.apply(base, economize.JumpConfig(target_layer=16, source_layer=17, exit_layer=17))
fit(model, "data/qwen/train.bin", steps=2000, output_dir="runs/jump_m16_u17")   # uint32 token file
model = economize.load(base, "runs/jump_m16_u17/final")
```

`fit` picks the objective registered for the attached drafter (`jump_rollout` for `jump`, the objective of
`tskip/train.py`). An existing `transformers.Trainer` pipeline uses `EffTrainer` with the same objectives:

```python
from economize.train import EffTrainer, JumpRollout, economize_training_args

objective = JumpRollout(horizon=16)
args = economize_training_args(objective, "runs/jump_trainer", max_steps=2000, per_device_train_batch_size=8)
EffTrainer(model, objective, args, train_dataset=dataset).train()
```

`pip install "economize[train]"` adds `accelerate`, `datasets` and `trl` (`fit` itself needs only torch).
[docs/training.md](https://github.com/Thiggel/economize/blob/main/docs/training.md) lists the objectives, the data formats, data-parallel runs under `accelerate launch` and how
to register an objective.

## Ecosystem

| Library | Entry point | Details |
|---|---|---|
| `transformers` | `model.generate(...)`, `custom_generate` Hub recipe | logits processors, stopping criteria, `return_dict_in_generate` |
| Hugging Face Hub | `model.save_pretrained`, `push_to_hub`, `economize.from_pretrained` | only the technique weights are uploaded |
| PEFT | `economize.apply(peft_model, ...)` | the adapter finds the base model inside a `PeftModel` |
| `kernels` | `economize.integrations.kernels.kernelize(model)`, `op(name)` | hub kernels for the base model and every added module |
| lm-evaluation-harness | `economize-lm-eval run --model economize --model_args pretrained=...,economize=...` | generation runs economize's decoding loop |
| TRL | `economize.integrations.trl.attach_rollouts` | in-process GRPO / RLOO rollouts with the drafter refreshed after each step |
| vLLM | `economize.backends.vllm` | drafter export (`export_drafter(model, dir, backend="vllm")`), plugin entry point, jump drafting in the V2 model runner (tested with vLLM 0.30) |
| SGLang | `economize.backends.sglang` | interface only (`JumpDraftBackend`); its registered exporter raises `UnsupportedBackendError`, and nothing is patched in SGLang yet |

[docs/ecosystem.md](https://github.com/Thiggel/economize/blob/main/docs/ecosystem.md) describes each integration and its tests.

## Saving and loading

```python
model.save_pretrained("my_drafter")       # economize_config.json + technique weights in safetensors
model = economize.load(base, "my_drafter")
```

Only the technique weights are saved. Checkpoints of the research code (`tskip`, a `jump.pt` file or its run
directory) are detected and converted on load. Further formats are added through the `FORMATS` registry.

## Extending economize

Everything that varies is a registry, and new entries are added by registration, never by editing economize:

| Registry | Module | Entry-point group |
|---|---|---|
| `TECHNIQUES` | `economize.core` | `economize.techniques` |
| `ARCHITECTURES`, `LAYER_KINDS` | `economize.arch` | `economize.architectures`, `economize.layer_kinds` |
| `FORMATS` | `economize.core` | `economize.formats` |
| `VERIFIERS`, `POLICIES` | `economize.decoding` | `economize.verifiers`, `economize.draft_policies` |
| `CACHES`, `MASKS`, `MASK_FORMATS` | `economize.runtime` | `economize.caches`, `economize.masks`, `economize.mask_formats` |
| `PREDICTORS`, `MIXERS` | `economize.techniques.jump` | `economize.jump_predictors`, `economize.jump_mixers` |
| `SKIP_SEARCHES`, `SKIP_METRICS` | `economize.techniques.layer_skip` | `economize.skip_searches`, `economize.skip_metrics` |
| `MEDUSA_TREES` | `economize.techniques.medusa` | `economize.medusa_trees` |
| `MTP_FAMILIES` | `economize.techniques.vendor_mtp` | `economize.mtp_families` |
| `CONFIDENCE_MEASURES`, `FLEXI_ROUTERS` | `economize.techniques.dynamic_depth` | `economize.calm_confidence`, `economize.flexi_routers` |
| `PRUNE_METHODS`, `DEPTH_CRITERIA` | `economize.techniques.prune` | `economize.prune_methods`, `economize.depth_criteria` |
| `QUANT_METHODS`, `QUANT_BACKENDS`, `QUANT_FORMATS`, `AWQ_SCALE_GROUPS`, `TORCHAO_TENSORS` | `economize.techniques.quant` | `economize.quant_methods`, `economize.quant_backends`, `economize.quant_formats`, `economize.awq_scale_groups`, `economize.torchao_tensors` |
| `FACTORIZERS` | `economize.techniques.lowrank` | `economize.factorizers` |
| `TEAL_BACKENDS` | `economize.techniques.sparsity` | `economize.teal_backends` |
| `EVICTION_METHODS` | `economize.techniques.kv_evict` | `economize.kv_eviction` |
| `KERNEL_OPS` | `economize.integrations.kernels` | `economize.kernel_ops` |
| `OBJECTIVES`, `DIVERGENCES`, `SCHEDULES`, `STATE_LOSSES` | `economize.train` | `economize.objectives`, `economize.divergences`, `economize.schedules`, `economize.state_losses` |
| `PROMPT_SOURCES`, `HARDWARE`, `SHAPE_READERS` | `economize.bench` | `economize.prompt_sources`, `economize.hardware`, `economize.shape_readers` |
| `EXPORTERS` | `economize.backends` | `economize.exporters` |
| `VLLM_DRAFTERS` | `economize.backends.vllm` | `economize.vllm_drafters` |
| `CT_SCHEMES` | `economize.integrations.compressed_tensors` | `economize.ct_schemes` |

A new predictor, for example:

```python
from economize.techniques.jump import PREDICTORS, Predictor

@PREDICTORS.register("residual_linear")
class ResidualLinear(Predictor):
    def __init__(self, hidden_size, config):
        super().__init__(hidden_size, config)
        self.proj = torch.nn.Linear(2 * hidden_size, hidden_size)

    def forward(self, source, token_embedding):
        return token_embedding.float() + self.proj(torch.cat([source, token_embedding], -1).float())

model = economize.apply(base, economize.JumpConfig(target_layer=12, source_layer=18, predictor="residual_linear"))
```

An installed package can contribute the same entry without an import, through the entry-point group in its
`pyproject.toml`:

```toml
[project.entry-points."economize.jump_predictors"]
residual_linear = "my_pkg.predictors:ResidualLinear"
```

A new technique subclasses `economize.Technique` with a frozen `TechniqueConfig` dataclass and is registered with
`@economize.register_technique`. A drafting technique declares `Capability.DRAFTER` and returns an object that follows
the `economize.decoding.Drafter` protocol. A new model family registers an `ArchAdapter` under its `model_type`;
most decoder-only families share the built-in `StandardDecoderAdapter`. Family facts that techniques and the cost
model need are set on the adapter: `layer_layout = "llama"` gives AWQ the Llama scale groups, `shape_fields`
corrects the cost model's config reading, and `module_path(module)` names modules for exporters. A drafter is
written for a serving backend by a `DrafterExporter` registered with
`@economize.backends.register_exporter("<backend>", "<technique>")`; a new backend is one such class.
[docs/ARCHITECTURE.md](https://github.com/Thiggel/economize/blob/main/docs/ARCHITECTURE.md) describes every contract.

## Scope of version 0.3

- Techniques: the table at the top. Jump drafting comes frozen, with LoRA or with full layer copies, from one or
  several source layers mixed by concatenation, a fused linear map or layer attention, with MLP, linear,
  transformer-block or flow-matching predictors.
- Decoding: batched ragged prompts with per-row selectors, stops and limits, chains and draft trees, greedy
  verification, lossless speculative sampling on chains and trees, typical acceptance, lossy drafting, fixed or
  confidence-based draft length, CUDA graphs for chain decoding, a log of every model call for cost accounting.
- Training: `economize.train` (`fit`, `EffTrainer`, registered objectives).
- Checkpoints record a fingerprint of the base weights; loading onto different weights warns
  (`economize.TeacherMismatchWarning`, off with `check_teacher=False`).
- Not yet: flex-attention masks for draft trees (trees use `sdpa` or `eager` masks), draft trees on models with
  recurrent layers, and the vLLM / SGLang draft step.

## Tests

```bash
cd TJEPA-Eff && source env.sh
PYTHONPATH=economize/src:economize/tests:. $PY -m pytest economize/tests            # unit, technique, integration, oracle (CPU)
PYTHONPATH=economize/src:economize/tests:. $PY -m pytest economize/tests -m slow    # trained checkpoints on Qwen2.5-0.5B
PYTHONPATH=economize/src:economize/tests:. $PY -m pytest economize/tests -m gpu     # batched decoding on a CUDA device
```

The lint (`uvx ruff check src tests hub`) and type (`uvx mypy --python-executable <python> src/economize`) gates and
the design rules a change must follow are in [docs/CONTRIBUTING.md](https://github.com/Thiggel/economize/blob/main/docs/CONTRIBUTING.md).

The unit tests check that the exact path equals the model's own logits on every supported family, that batched
lossless decoding equals greedy decoding of each prompt alone (chains, trees and jump drafters, with partial
acceptance and hybrid recurrent layers), that rejection sampling outputs the target distribution (chi-square
tests with a negative control), the recipe's ordering and rollback, and the checkpoint round trips. The oracle
tests check that economize produces the same tokens, the same log of model calls and the same acceptance as the
research decoder `tskip.decode.Decoder`, on random drafters of a tiny model and on three trained drafters of
Qwen2.5-0.5B.

## License

Apache License 2.0, see [LICENSE](https://github.com/Thiggel/economize/blob/main/LICENSE).
