Gearbx: Entropy-Routed Dynamic Quantization

An entropy-routed dynamic quantization engine for LLM inference that adjusts weight precision on a per-token basis during generation.

Website: gearbx.jpdz.app

How It Works

Not all tokens require the same computational fidelity. Producing "Hello, how can I help you?" demands almost no semantic reasoning, while the next token in a partial differential equation imposes high cognitive load on the model.

Gearbx monitors the model's own output entropy at each generation step and routes subsequent forward passes through pre-cached layer weights at the appropriate precision tier:

Autoregressive LLM inference is memory-bandwidth bound: each generated token loads the full weight matrix. Packed storage means fewer bytes transferred per token, translating directly to speed gains proportional to compression ratio.

Architecture

Prompt → [PromptClassifier] → initial_gear
                │
    ┌───────────▼──────────────────────────┐
    │       GENERATION LOOP                │
    │                                      │
    │  input_ids → model.forward() → logits│
    │                    │                 │
    │          [EntropyMonitor]            │
    │      Shannon entropy → gear decision │
    │           (rolling avg + hysteresis) │
    │                    │                 │
    │         [PrecisionManager]           │
    │     lazy quantize + module swap      │
    │     (originals offloaded to CPU)     │
    │                    │                 │
    │          sample next_token           │
    └──────────────────────────────────────┘
                │
    generated_ids → decode → output_text

Core Components

• EntropyMonitor (monitor.py: Computes Shannon entropy (bits, log base 2) from logit distributions. Rolling-window average smooths gear transitions. Supports vocab-size auto-scaling (reference: 32768 tokens), hysteresis to prevent oscillation near boundaries, and auto-calibration from observed entropy distribution (p40/p60 thresholds). Deferred pipeline mode eliminates per-token GPU→CPU sync.

• PrecisionManager (precision.py: Lazy quantization with direct module-reference swaps. Discovers attention layers at init via detect_attn_prefixes() but defers quantization to first shift() call. On shift: quantizes from originals, swaps quantized module into model tree, offloads originals to CPU. Only one quantized variant exists at any time, memory goes down, not up.

• QuantizedLinear Kernels (kernels.py: Real packed integer storage: int8 (per-channel symmetric), int6, int4 (2 values per byte), int2 (4 values per byte, ternary {-1,0,1}). Dual-path forward: fused Triton kernels on CUDA (in-register dequant, no intermediate tensor), dequant-cache path on MPS/CPU (first forward unpacks to fp16, subsequent forwards reuse cached tensor).

• Native Acceleration (native_kernels.py + csrc/: Three tiers below Triton: MPS native Metal shaders via PyTorch's MetalShaderLibrary (zero-copy, threadgroup 256), legacy Metal command buffers, and CPU NEON vectorized matmul (ARM -O3 -march=armv8.2-a+fp16).

• Fused Triton Kernels (triton_kernels.py: CUDA-only fused quantized matmul. Multiplies activations directly with packed int8/int4/int2 weights without materializing fp16. In-register dequantization - only packed data traverses global memory.

• GearbxModel (model.py: Orchestrates everything. Loads model via transformers, auto-detects attention architecture, passes vocab_size to EntropyMonitor for threshold scaling. Manual autoregressive loop with KV caching and NaN guards for MPS stability.

• StatisticalPromptClassifier (classifier.py: Cold-start heuristic. Scores prompt complexity from math symbols, code keywords, length, avg word length → initial gear before the entropy window fills.

Backends

Transformers backend: Full per-token gear shifting with real weight swaps. Supports MPS, CUDA, and CPU.

MLX backend: Native Apple Silicon via mlx-lm. Unified memory means no CPU offloading. Supports telemetry-only mode for pre-quantized models. Median-based calibration.

Ollama backend: Per-prompt routing via prompt classifier. Talks to Ollama, llama.cpp, vLLM, LM Studio, or any OpenAI-compatible local server over HTTP. Real per-token entropy telemetry via top_logprobs. Supports SSE streaming with StreamChunk per-token telemetry.

Device Support

Quick Start

Transformers Backend

from gearbx import GearbxModel

# Device auto-detected: MPS > CUDA > CPU
gbm = GearbxModel(
    'mistralai/Mistral-7B-Instruct-v0.3',
    num_attn_layers_to_manage=16,
    high_thresh=3.5,
    low_thresh=1.8,
)

r = gbm.generate('What is the capital of Japan?', max_new_tokens=30)
print(r.text)
print(r.gear_stats)  # e.g. {'low': 0.80, 'mid': 0.20}

MLX Backend (Apple Silicon)

from gearbx import MLXGearbxModel

gbm = MLXGearbxModel('mlx-community/Mistral-7B-Instruct-v0.3-4bit')
r = gbm.generate('Explain entropy in information theory.', max_new_tokens=200)
print(r.text)
print(r.gear_stats)
print(f'{r.tokens_per_sec:.1f} tok/s')

Ollama Backend

from gearbx import OllamaGearbx

gbm = OllamaGearbx(
    gear_models={
        'low':  'mistral:7b',
        'mid':  'mistral:7b',
        'high': 'mistral:7b',
    },
    low_thresh=0.4,
    high_thresh=1.2,
)

# Streaming with per-token telemetry
for chunk in gbm.generate_stream('Prove sqrt(2) is irrational.', max_new_tokens=200):
    print(chunk.token, end='', flush=True)

GGUF via Ollama Cache

# Load GGUF from Ollama's local blob cache (no HF download)
python3 run_gguf.py llama3.2:1b

Installation

Apple Silicon (MPS)

python3 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"

python3 -c "import torch; print('MPS:', torch.backends.mps.is_available())"

Apple Silicon (MLX)

python3 -m venv .venv && source .venv/bin/activate
pip install -e ".[mlx,dev]"

NVIDIA (CUDA)

python3 -m venv .venv && source .venv/bin/activate
pip install -e ".[cuda,dev]"

python3 -c "import torch; print('CUDA:', torch.cuda.is_available())"

TUI (Terminal UI)

# From source
cd tui && make build

# Via npm (downloads prebuilt binary)
npm install -g gearbx
gearbx

Requirements

• Python >= 3.10
• PyTorch >= 2.1.0
• transformers >= 4.40.0
• accelerate >= 0.28.0
• bitsandbytes >= 0.43.0 (CUDA only, optional)
• mlx >= 0.12.0, mlx-lm >= 0.12.0 (MLX backend, optional)

Hardware Requirements

Running Tests

# All unit tests (no GPU needed)
pytest tests/ -v

# Individual test files
pytest tests/test_monitor.py -v
pytest tests/test_precision.py -v
pytest tests/test_integration.py -v

Benchmarks

All benchmarks support --unit mode (synthetic data, CPU) and full mode (real model, MPS/CUDA).

# Unit mode, no GPU needed
python3 benchmarks/bench_entropy.py --unit
python3 benchmarks/bench_latency.py --unit
python3 benchmarks/bench_quality.py --unit

# Full mode, auto-detects MPS or CUDA
python3 benchmarks/bench_entropy.py --model mistralai/Mistral-7B-Instruct-v0.3
python3 benchmarks/bench_latency.py --model mistralai/Mistral-7B-Instruct-v0.3
python3 benchmarks/bench_quality.py --model mistralai/Mistral-7B-Instruct-v0.3 --n 200

# Production benchmark
python3 benchmarks/bench_production.py

Benchmark Targets

Threshold Tuning

Thresholds auto-scale by log2(vocab_size) / log2(32768) when vocab_size is provided (GearbxModel does this automatically).

Auto-calibration: calibrate_from_logits() computes entropy from a prefill pass and sets thresholds at p40/p60 of the distribution. Post-calibration caps prevent inflated thresholds: high ≤ log2(vocab) * 0.18, low ≤ log2(vocab) * 0.06.

Suggested Thresholds by Model

For Ollama backend, top-k truncated entropy has a narrower range, use lower thresholds (e.g., low=0.4, high=1.2).

Known Limitations

1. Single-Stream Only: Designed for single-sequence local inference, not batched serving. generate_batch() runs prompts sequentially with per-sequence gear routing.

2. KV Cache Precision Mismatch: Mid-sequence gear shifts create mixed-precision KV cache entries. Use min_gear_duration=4 to reduce noise, or use_cache=False during strict validation.

3. MPS fp16 Base: On Apple Silicon (transformers backend), the base model runs at fp16 since bitsandbytes doesn't support MPS. Memory usage is higher than CUDA 4-bit but entropy routing and quality benefits still apply.

4. Double Quantization: Loading GGUF→fp16→int4 compounds quantization error. PrecisionManager warns when detected. Consider higher bit-width for already-quantized source models.

5. Ollama Coarse Routing: Ollama backend routes per-prompt (not per-token) since the server is a black box. Entropy telemetry is still per-token.

6. Catastrophic Entropy: If entropy exceeds 90% of theoretical max for 2+ consecutive tokens, generation auto-falls-back to mid gear. Catches model failures from bad quantization.

License

Proprietary - Jpdz Labs. All rights reserved.