tqai 0.1.0__tar.gz

tqai-0.1.0/.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,55 @@
+name: Bug Report
+description: Report a bug in tqai
+labels: [bug]
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Description
+      description: A clear description of the bug
+    validations:
+      required: true
+
+  - type: textarea
+    id: reproduce
+    attributes:
+      label: Steps to Reproduce
+      description: Minimal code or commands to reproduce the issue
+      render: python
+
+  - type: textarea
+    id: expected
+    attributes:
+      label: Expected Behavior
+      description: What you expected to happen
+
+  - type: textarea
+    id: actual
+    attributes:
+      label: Actual Behavior
+      description: What actually happened (include error messages)
+
+  - type: dropdown
+    id: backend
+    attributes:
+      label: Backend
+      options:
+        - PyTorch (CPU)
+        - PyTorch (CUDA)
+        - MLX (Apple Silicon)
+    validations:
+      required: true
+
+  - type: input
+    id: model
+    attributes:
+      label: Model
+      description: Model ID you were using
+      placeholder: e.g., mlx-community/Llama-3.1-8B-Instruct-4bit
+
+  - type: textarea
+    id: environment
+    attributes:
+      label: Environment
+      description: Paste the output of `tqai info`
+      render: text

tqai-0.1.0/.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,23 @@
+name: Feature Request
+description: Suggest a feature for tqai
+labels: [enhancement]
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Description
+      description: What feature would you like to see?
+    validations:
+      required: true
+
+  - type: textarea
+    id: use-case
+    attributes:
+      label: Use Case
+      description: How would you use this feature?
+
+  - type: textarea
+    id: alternatives
+    attributes:
+      label: Alternatives Considered
+      description: Any workarounds or alternatives you've tried

tqai-0.1.0/.github/homebrew/tqai.rb

@@ -0,0 +1,33 @@
+class Tqai < Formula
+  include Language::Python::Virtualenv
+
+  desc "TurboQuant KV cache compression for local LLM inference"
+  homepage "https://github.com/AlphaWaveSystems/tqai"
+  url "https://files.pythonhosted.org/packages/source/t/tqai/tqai-0.1.0.tar.gz"
+  sha256 "PLACEHOLDER_SHA256"
+  license "MIT"
+
+  depends_on "python@3.12"
+
+  resource "numpy" do
+    url "https://files.pythonhosted.org/packages/source/n/numpy/numpy-2.2.5.tar.gz"
+    sha256 "PLACEHOLDER_SHA256"
+  end
+
+  def install
+    virtualenv_install_with_resources
+  end
+
+  test do
+    assert_match "tqai v", shell_output("#{bin}/tqai info")
+  end
+end
+
+# NOTE: After publishing to PyPI, update the sha256 values:
+#   1. Download the sdist: pip download tqai==0.1.0 --no-deps --no-binary :all:
+#   2. Get sha256: shasum -a 256 tqai-0.1.0.tar.gz
+#   3. Same for numpy
+#
+# To install this formula into the existing tap:
+#   cp .github/homebrew/tqai.rb /path/to/homebrew-tap/Formula/tqai.rb
+#   cd /path/to/homebrew-tap && git add Formula/tqai.rb && git commit -s -m "Add tqai formula"

tqai-0.1.0/.github/pull_request_template.md

@@ -0,0 +1,21 @@
+## Summary
+
+<!-- What does this PR do? -->
+
+## Type of Change
+
+- [ ] Bug fix
+- [ ] New feature
+- [ ] Performance improvement
+- [ ] Documentation
+- [ ] Refactoring
+
+## Testing
+
+- [ ] `pytest tests/` passes
+- [ ] `ruff check .` passes
+- [ ] Added tests for new functionality (if applicable)
+
+## DCO
+
+- [ ] All commits are signed off (`git commit -s`)

tqai-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,45 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: pip install ruff
+      - run: ruff check .
+
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          pip install -e ".[dev]"
+          if [[ "$RUNNER_OS" == "macOS" ]]; then
+            pip install mlx mlx-lm || true
+          fi
+
+      - name: Run tests
+        run: |
+          pytest tests/ \
+            --ignore=tests/test_e2e_models.py \
+            --ignore=tests/test_e2e_large_models.py \
+            -v --tb=short

tqai-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,78 @@
+name: Release
+
+on:
+  push:
+    tags: ["v*"]
+
+permissions:
+  contents: write
+  id-token: write
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Build package
+        run: |
+          pip install hatchling build
+          python -m build
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish-pypi:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+      - name: Create GitHub Release
+        env:
+          GH_TOKEN: ${{ github.token }}
+          TAG_NAME: ${{ github.ref_name }}
+        run: |
+          gh release create "$TAG_NAME" dist/* \
+            --title "$TAG_NAME" \
+            --notes-file CHANGELOG.md
+
+  notify-homebrew:
+    needs: publish-pypi
+    runs-on: ubuntu-latest
+    steps:
+      - name: Trigger Homebrew formula update
+        env:
+          GH_TOKEN: ${{ secrets.HOMEBREW_TAP_TOKEN }}
+          VERSION: ${{ github.ref_name }}
+        run: |
+          gh api repos/AlphaWaveSystems/homebrew-tap/dispatches \
+            -f event_type=tqai-release \
+            -f "client_payload[version]=$VERSION" || true

tqai-0.1.0/.gitignore

@@ -0,0 +1,40 @@
+# Python
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+dist/
+build/
+*.egg
+
+# Testing
+.pytest_cache/
+htmlcov/
+.coverage
+coverage.xml
+
+# Tools
+.ruff_cache/
+.mypy_cache/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Converted model outputs (user-generated)
+*-tqai/
+
+# Keep shipped codebook data
+!src/tqai/codebook/data/*.npz

tqai-0.1.0/CHANGELOG.md

@@ -0,0 +1,22 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.1.0] - 2026-04-04
+
+### Added
+
+- Core PolarQuantizer implementing TurboQuant Stage 1 (random rotation + Lloyd-Max scalar quantization)
+- Multi-backend support: PyTorch (CPU/CUDA) and MLX (Apple Silicon)
+- HuggingFace Transformers integration via `DynamicCache` subclass
+- mlx-lm integration via `KVCache` replacement
+- Precomputed Lloyd-Max codebooks for head dimensions 64, 96, 128, 256 at 2/3/4 bits
+- Asymmetric K/V bit allocation (e.g., K4/V2 for 80% memory savings)
+- Sink token support (keep first N tokens uncompressed)
+- CLI tool with commands: `info`, `benchmark`, `run`, `compare`, `convert`
+- Offline model conversion (`tqai convert`) for faster startup
+- Comprehensive test suite (179 tests) covering accuracy, edge cases, and cross-backend consistency
+- Verified quality-neutral on 8B+ parameter models (Llama 3.1 8B, Qwen 14B)

tqai-0.1.0/CODE_OF_CONDUCT.md

@@ -0,0 +1,39 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, caste, color, religion, or sexual
+identity and orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior:
+
+* The use of sexualized language or imagery, and sexual attention or advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a professional setting
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the project maintainers at **github@alphawavesystems.com**.
+
+All complaints will be reviewed and investigated promptly and fairly.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org/), version 2.1.

tqai-0.1.0/CONTRIBUTING.md

@@ -0,0 +1,89 @@
+# Contributing to tqai
+
+Thank you for your interest in contributing to tqai! This guide explains how to get involved.
+
+## Reporting Bugs
+
+Open an issue on [GitHub Issues](https://github.com/AlphaWaveSystems/tqai/issues) with:
+
+- A clear, descriptive title
+- Steps to reproduce the problem
+- Expected vs actual behavior
+- tqai version (`tqai info`), OS, Python version
+- Model ID and backend (torch/mlx) you were using
+
+## Pull Request Flow
+
+1. **Fork** the repository on GitHub
+2. **Branch** from `main` — use a descriptive name (e.g., `fix/dtype-mismatch`, `feat/bit-packing`)
+3. **Develop** your changes following the code style guidelines below
+4. **Test** your changes (see Testing section)
+5. **Submit a PR** against `main` with a clear description of what and why
+
+Keep PRs focused on a single concern. If you have multiple unrelated changes, submit separate PRs.
+
+## Code Style
+
+- **Linting**: Run `ruff check .` before committing
+- **Formatting**: Follow PEP 8, 100-char line length
+- **Imports**: Sorted by ruff (isort compatible)
+- **Types**: Use type hints for public API functions
+
+## DCO Sign-Off
+
+All commits must include a Developer Certificate of Origin sign-off line:
+
+```
+Signed-off-by: Your Name <your@email.com>
+```
+
+Add this automatically with:
+
+```bash
+git commit -s -m "Your commit message"
+```
+
+This certifies that you wrote or have the right to submit the code under the project's license. See [developercertificate.org](https://developercertificate.org/) for the full text.
+
+## Testing
+
+Before submitting a PR, run the test suite:
+
+```bash
+# Install dev dependencies
+pip install -e ".[dev]"
+
+# Run unit + accuracy tests
+pytest tests/ --ignore=tests/test_e2e_models.py --ignore=tests/test_e2e_large_models.py
+
+# Run linter
+ruff check .
+```
+
+For changes to the quantizer or cache, also run the accuracy tests:
+
+```bash
+pytest tests/test_accuracy.py -v
+```
+
+## Adding Support for New Head Dimensions
+
+If you need codebooks for a head dimension not in `{64, 96, 128, 256}`:
+
+```bash
+pip install tqai[codegen]
+python scripts/generate_codebooks.py
+```
+
+## Architecture Overview
+
+- `src/tqai/quantizer.py` — Core PolarQuantizer algorithm
+- `src/tqai/backend/` — PyTorch + MLX abstraction (Protocol-based)
+- `src/tqai/codebook/` — Lloyd-Max codebook generation and loading
+- `src/tqai/cache/` — HuggingFace and mlx-lm cache wrappers
+- `src/tqai/cli.py` — CLI tool
+- `src/tqai/convert.py` — Offline model conversion
+
+## License
+
+By contributing to tqai, you agree that your contributions will be licensed under the MIT License.

tqai-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 AlphaWaveSystems
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tqai-0.1.0/PKG-INFO

@@ -0,0 +1,210 @@
+Metadata-Version: 2.4
+Name: tqai
+Version: 0.1.0
+Summary: TurboQuant KV cache compression for local LLM inference
+Project-URL: Homepage, https://github.com/AlphaWaveSystems/tqai
+Project-URL: Repository, https://github.com/AlphaWaveSystems/tqai
+Project-URL: Issues, https://github.com/AlphaWaveSystems/tqai/issues
+Project-URL: Changelog, https://github.com/AlphaWaveSystems/tqai/blob/main/CHANGELOG.md
+Author: pbertsch
+License-Expression: MIT
+License-File: LICENSE
+Keywords: apple-silicon,compression,kv-cache,llm,mlx,quantization,turboquant
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+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: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.24
+Provides-Extra: codegen
+Requires-Dist: scipy>=1.10; extra == 'codegen'
+Provides-Extra: cuda
+Requires-Dist: torch>=2.1; extra == 'cuda'
+Requires-Dist: triton>=2.0; extra == 'cuda'
+Provides-Extra: dev
+Requires-Dist: pytest-cov>=4.0; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Requires-Dist: scipy>=1.10; extra == 'dev'
+Requires-Dist: torch>=2.1; extra == 'dev'
+Requires-Dist: transformers>=4.38; extra == 'dev'
+Provides-Extra: mlx
+Requires-Dist: mlx-lm>=0.12; extra == 'mlx'
+Requires-Dist: mlx>=0.14; extra == 'mlx'
+Provides-Extra: torch
+Requires-Dist: torch>=2.1; extra == 'torch'
+Description-Content-Type: text/markdown
+
+# tqai
+
+[![PyPI version](https://img.shields.io/pypi/v/tqai)](https://pypi.org/project/tqai/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
+[![CI](https://github.com/AlphaWaveSystems/tqai/actions/workflows/ci.yml/badge.svg)](https://github.com/AlphaWaveSystems/tqai/actions/workflows/ci.yml)
+[![Python 3.10+](https://img.shields.io/pypi/pyversions/tqai)](https://pypi.org/project/tqai/)
+
+TurboQuant KV cache compression for local LLM inference.
+
+Compresses the KV cache to ~3 bits per channel with **80%+ memory savings** and near-zero quality loss on 8B+ models. Supports both PyTorch (CPU/CUDA) and MLX (Apple Silicon).
+
+Based on [TurboQuant](https://arxiv.org/abs/2504.19874) (Google Research, ICLR 2026).
+
+## Installation
+
+```bash
+# Homebrew (macOS)
+brew install alphawavesystems/tap/tqai
+
+# PyPI
+pip install tqai
+
+# With PyTorch backend
+pip install tqai[torch]
+
+# With MLX backend (Apple Silicon)
+pip install tqai[mlx]
+```
+
+## Quick Start
+
+### HuggingFace Transformers (PyTorch)
+
+```python
+import tqai
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-3B-Instruct", dtype="auto")
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B-Instruct")
+
+# One line to enable KV cache compression
+cache = tqai.patch(model, bits_k=4, bits_v=2)
+
+inputs = tokenizer("Explain quantum entanglement:", return_tensors="pt")
+output = model.generate(**inputs, past_key_values=cache, max_new_tokens=200)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+### MLX (Apple Silicon)
+
+```python
+import tqai
+import mlx_lm
+
+model, tokenizer = mlx_lm.load("mlx-community/Llama-3.1-8B-Instruct-4bit")
+
+# One line to enable KV cache compression
+tqai.patch(model, bits_k=4, bits_v=2, backend="mlx")
+
+response = mlx_lm.generate(model, tokenizer, prompt="Explain quantum entanglement:", max_tokens=200)
+print(response)
+
+# Restore original behaviour when done
+tqai.unpatch(model)
+```
+
+## Compression Configs
+
+| Config | Avg Bits | Memory Saved | Recommended For |
+|--------|----------|--------------|-----------------|
+| `bits_k=4, bits_v=2` | 3.0 | **80%** | Production (best quality/compression balance) |
+| `bits_k=3, bits_v=2` | 2.5 | **84%** | Extended context windows |
+| `bits_k=4, bits_v=3` | 3.5 | **78%** | Quality-sensitive applications |
+
+## How It Works
+
+tqai implements Stage 1 of TurboQuant (PolarQuant):
+
+1. **Random orthogonal rotation** — Rotates KV vectors by a fixed Haar-distributed matrix to spread information across all coordinates
+2. **Lloyd-Max scalar quantization** — Quantizes each coordinate independently using precomputed optimal codebooks
+3. **Norm preservation** — Stores vector norms separately in FP16
+
+No training, calibration, or model-specific tuning required. The same codebooks work for any model.
+
+## Quality Results
+
+Tested on Apple Silicon with various model sizes:
+
+| Model | Baseline | + tqai K4/V2 | + tqai K3/V2 |
+|-------|----------|--------------|--------------|
+| Qwen 0.5B | Good | Degraded | Poor |
+| Qwen 3B bf16 | Excellent | Good | Degraded |
+| Llama 8B Q4 | Excellent | **Excellent** | **Excellent** |
+| Qwen 14B Q4 | Excellent | **Excellent** | **Excellent** |
+
+Quality is near-identical to baseline on 8B+ parameter models.
+
+## CLI
+
+tqai includes a command-line tool for quick testing without writing code:
+
+```bash
+# Show environment and library info
+tqai info
+
+# Run quantization accuracy benchmark
+tqai benchmark
+tqai benchmark --bits-k 3 --bits-v 2 --head-dim 128
+
+# Generate text with TurboQuant compression
+tqai run "Explain gravity" --model mlx-community/Llama-3.1-8B-Instruct-4bit
+tqai run "Explain gravity" --model Qwen/Qwen2.5-3B-Instruct --backend torch
+
+# Compare baseline vs compressed output side by side
+tqai compare "Explain gravity" --model mlx-community/Llama-3.1-8B-Instruct-4bit
+
+# Pre-convert a model for faster startup
+tqai convert --model mlx-community/Llama-3.1-8B-Instruct-4bit --output ./llama-tqai/
+tqai run "Explain gravity" --model mlx-community/Llama-3.1-8B-Instruct-4bit --tqai-config ./llama-tqai/
+
+# Run without compression (baseline)
+tqai run "Explain gravity" --model mlx-community/Llama-3.1-8B-Instruct-4bit --no-tqai
+```
+
+## Advanced Options
+
+```python
+cache = tqai.patch(
+    model,
+    bits_k=4,           # Bits per key coordinate (2, 3, or 4)
+    bits_v=2,           # Bits per value coordinate (2, 3, or 4)
+    sink_tokens=4,      # Keep first N tokens uncompressed (attention sinks)
+    backend="torch",    # Force backend: "torch" or "mlx"
+    device="cuda",      # PyTorch device (ignored for MLX)
+)
+```
+
+## Running Tests
+
+```bash
+# Install dev dependencies
+pip install tqai[dev]
+
+# Unit + accuracy tests (175 tests, <1s)
+pytest tests/ --ignore=tests/test_e2e_models.py --ignore=tests/test_e2e_large_models.py
+
+# End-to-end with real models (requires model downloads)
+pytest tests/test_e2e_models.py -v -s
+```
+
+## Project Structure
+
+```
+src/tqai/
+├── __init__.py          # patch(), unpatch(), TurboQuantConfig
+├── config.py            # Configuration dataclass
+├── quantizer.py         # PolarQuantizer (core algorithm)
+├── backend/             # PyTorch + MLX abstraction
+├── codebook/            # Lloyd-Max codebooks (precomputed)
+└── cache/               # HuggingFace + mlx-lm integrations
+```
+
+## License
+
+MIT