turboquant-llm 0.1.0__tar.gz

turboquant_llm-0.1.0/.gitignore

@@ -0,0 +1,11 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+dist/
+build/
+*.egg-info/
+.venv/
+venv/
+.env

turboquant_llm-0.1.0/PKG-INFO

@@ -0,0 +1,120 @@
+Metadata-Version: 2.4
+Name: turboquant-llm
+Version: 0.1.0
+Summary: Reference implementation of TurboQuant-style vector/KV quantization (rotation + scalar quant + QJL residual).
+Project-URL: Homepage, https://github.com/PestanaRobson/turboquant_llm
+Project-URL: Documentation, https://github.com/PestanaRobson/turboquant_llm#readme
+Project-URL: Repository, https://github.com/PestanaRobson/turboquant_llm
+Project-URL: Issues, https://github.com/PestanaRobson/turboquant_llm/issues
+License-Expression: Apache-2.0
+Keywords: kv-cache,llm,quantization,turboquant,vector-quantization
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software 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: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.22
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == 'torch'
+Description-Content-Type: text/markdown
+
+# turboquant_llm
+
+Implementação de **referência** em NumPy do fluxo estilo **TurboQuant**: rotação ortogonal → quantização escalar → estágio **QJL** (sinais 1-bit da projeção JL do resíduo). Serve para experimentos, testes e desenho de API; **não** substitui kernels CUDA/Triton em produção.
+
+No PyPI o nome do pacote é **`turboquant-llm`** (hífen); em código use **`import turboquant_llm`**. Código-fonte: [github.com/PestanaRobson/turboquant_llm](https://github.com/PestanaRobson/turboquant_llm).
+
+> Os nomes `turboquant` e `turboquant-kv` já existem no PyPI como outros projetos.
+
+## Por que TurboQuant importa (o ganho real)
+
+Em LLMs com **contexto longo**, o **cache KV** (chaves e valores guardados por token para não recalcular atenção) domina **memória** e **largura de banda** entre HBM e compute — é um gargalo clássico de servir modelo grande com muitos tokens.
+
+O método **TurboQuant** ([paper](https://arxiv.org/abs/2504.19874), [blog Google Research](https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/)) foi desenhado para **comprimir vetores de alta dimensão** usados nesse caminho (em especial **K/V** e produtos internos em atenção) com quantização **online** e **sem calibração em dataset** (“data-oblivious”): não exige coletar estatísticas do seu domínio para treinar codebooks.
+
+**Resultados reportados pelo Google** (ambiente e modelo nos artigos; números indicativos, não garantidos para o seu hardware ou checkpoint):
+
+| Eixo | O que se ganha (ordem de grandeza) |
+|------|-------------------------------------|
+| **Memória do KV** | Redução da ordem de **~6×** no footprint do cache KV em relação a representações densas não comprimidas nos experimentos citados. |
+| **Velocidade da atenção** | Até **~8×** de aceleração no cálculo de **logits de atenção** vs. chaves em precisão total (ex.: FP32), em GPU **H100** e implementação otimizada (baseline JAX no blog). |
+| **Qualidade** | No paper: **~3,5 bits por canal** com neutralidade de qualidade forte; **~2,5** com degradação marginal. No blog: quantização do KV na ordem de **~3 bits** mantendo desempenho em benchmarks longos (LongBench, needle-in-haystack, etc.). |
+| **Overhead de quantização** | Evita o custo de armazenar **constantes de quantização completas por bloco** como em muitos esquemas clássicos — parte central do argumento de eficiência do TurboQuant. |
+
+**O que *esta* biblioteca faz:** reproduz a **lógica** (rotação + quantização + estágio tipo QJL) para pesquisa e integração futura. **Não** reproduz sozinha os números de H100 nem a memória real do servidor — isso exige **kernels GPU** e encaixe no **motor de inferência** (vLLM, TensorRT-LLM, etc.). Os ganhos acima são o **alvo** do algoritmo publicado; o caminho até lá é integração nativa no stack de KV.
+
+## Instalação
+
+```bash
+pip install turboquant-llm
+```
+
+Opcional (integrações futuras com modelos):
+
+```bash
+pip install turboquant-llm[torch]
+```
+
+## Uso rápido
+
+```python
+import numpy as np
+from turboquant_llm import TurboQuantCompressor, TurboQuantConfig
+
+dim = 128
+x = np.random.default_rng(0).standard_normal((10, dim))
+
+comp = TurboQuantCompressor(
+    dim,
+    TurboQuantConfig(bits_main=4, qjl_projection_dim=64, seed=0),
+)
+batch = comp.compress(x)
+x_hat = comp.reconstruct_first_stage(batch)  # ignora o estágio QJL no resíduo
+```
+
+## Publicar no PyPI
+
+Repositório: [PestanaRobson/turboquant_llm](https://github.com/PestanaRobson/turboquant_llm). Envie o conteúdo deste diretório para o branch principal (`git push`).
+
+1. Com [API token](https://pypi.org/manage/account/token/) do PyPI:
+
+```bash
+python -m pip install build twine
+python -m build
+python -m twine upload dist/*
+```
+
+Teste antes no [TestPyPI](https://test.pypi.org/).
+
+## Gemma 4 e Qwen 3.x: melhor caminho com este tipo de quantização
+
+TurboQuant (no paper e no blog do Google Research) ataca sobretudo **memória e custo do cache KV** e produtos internos em atenção — não é só “exportar pesos em INT4”.
+
+1. **Hoje (ecossistema Hugging Face / vLLM)**  
+   Integração “de verdade” costuma exigir **suporte no motor de inferência** (vLLM, SGLang, TensorRT-LLM, etc.) ou um **módulo de atenção customizado** que, a cada passo, comprima/descomprima **K/V por cabeça** com a mesma rotação e parâmetros fixos por modelo/camada.
+
+2. **Gemma (ex. Gemma 3 / família Gemma no Hub)** e **Qwen3**  
+   - Confirme `hidden_size`, **GQA** (`num_key_value_heads`), `head_dim` no `config.json` do checkpoint.  
+   - A lib aqui opera em vetores de tamanho `dim` (= `head_dim` por cabeça, ou blocos que vocês definirem).  
+   - Caminho pragmático: usar este pacote para **benchmark offline** (distorsão, memória simulada); em seguida portar o núcleo para **PyTorch custom op** ou **Triton** no caminho de escrita/leitura do KV.
+
+3. **Ordem sugerida**  
+   - Prototipar com esta API em **uma camada** ou tensor KV sintético.  
+   - Medir qualidade (long context / needle) com `bits_main` ~3–4 como no material de referência.  
+   - Só então acoplar ao servidor (vLLM plugin / patch) para latência real.
+
+## Licença
+
+Apache-2.0
+
+## Referências
+
+- [TurboQuant (arXiv)](https://arxiv.org/abs/2504.19874)  
+- [Google Research blog](https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/)

turboquant_llm-0.1.0/README.md

@@ -0,0 +1,93 @@
+# turboquant_llm
+
+Implementação de **referência** em NumPy do fluxo estilo **TurboQuant**: rotação ortogonal → quantização escalar → estágio **QJL** (sinais 1-bit da projeção JL do resíduo). Serve para experimentos, testes e desenho de API; **não** substitui kernels CUDA/Triton em produção.
+
+No PyPI o nome do pacote é **`turboquant-llm`** (hífen); em código use **`import turboquant_llm`**. Código-fonte: [github.com/PestanaRobson/turboquant_llm](https://github.com/PestanaRobson/turboquant_llm).
+
+> Os nomes `turboquant` e `turboquant-kv` já existem no PyPI como outros projetos.
+
+## Por que TurboQuant importa (o ganho real)
+
+Em LLMs com **contexto longo**, o **cache KV** (chaves e valores guardados por token para não recalcular atenção) domina **memória** e **largura de banda** entre HBM e compute — é um gargalo clássico de servir modelo grande com muitos tokens.
+
+O método **TurboQuant** ([paper](https://arxiv.org/abs/2504.19874), [blog Google Research](https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/)) foi desenhado para **comprimir vetores de alta dimensão** usados nesse caminho (em especial **K/V** e produtos internos em atenção) com quantização **online** e **sem calibração em dataset** (“data-oblivious”): não exige coletar estatísticas do seu domínio para treinar codebooks.
+
+**Resultados reportados pelo Google** (ambiente e modelo nos artigos; números indicativos, não garantidos para o seu hardware ou checkpoint):
+
+| Eixo | O que se ganha (ordem de grandeza) |
+|------|-------------------------------------|
+| **Memória do KV** | Redução da ordem de **~6×** no footprint do cache KV em relação a representações densas não comprimidas nos experimentos citados. |
+| **Velocidade da atenção** | Até **~8×** de aceleração no cálculo de **logits de atenção** vs. chaves em precisão total (ex.: FP32), em GPU **H100** e implementação otimizada (baseline JAX no blog). |
+| **Qualidade** | No paper: **~3,5 bits por canal** com neutralidade de qualidade forte; **~2,5** com degradação marginal. No blog: quantização do KV na ordem de **~3 bits** mantendo desempenho em benchmarks longos (LongBench, needle-in-haystack, etc.). |
+| **Overhead de quantização** | Evita o custo de armazenar **constantes de quantização completas por bloco** como em muitos esquemas clássicos — parte central do argumento de eficiência do TurboQuant. |
+
+**O que *esta* biblioteca faz:** reproduz a **lógica** (rotação + quantização + estágio tipo QJL) para pesquisa e integração futura. **Não** reproduz sozinha os números de H100 nem a memória real do servidor — isso exige **kernels GPU** e encaixe no **motor de inferência** (vLLM, TensorRT-LLM, etc.). Os ganhos acima são o **alvo** do algoritmo publicado; o caminho até lá é integração nativa no stack de KV.
+
+## Instalação
+
+```bash
+pip install turboquant-llm
+```
+
+Opcional (integrações futuras com modelos):
+
+```bash
+pip install turboquant-llm[torch]
+```
+
+## Uso rápido
+
+```python
+import numpy as np
+from turboquant_llm import TurboQuantCompressor, TurboQuantConfig
+
+dim = 128
+x = np.random.default_rng(0).standard_normal((10, dim))
+
+comp = TurboQuantCompressor(
+    dim,
+    TurboQuantConfig(bits_main=4, qjl_projection_dim=64, seed=0),
+)
+batch = comp.compress(x)
+x_hat = comp.reconstruct_first_stage(batch)  # ignora o estágio QJL no resíduo
+```
+
+## Publicar no PyPI
+
+Repositório: [PestanaRobson/turboquant_llm](https://github.com/PestanaRobson/turboquant_llm). Envie o conteúdo deste diretório para o branch principal (`git push`).
+
+1. Com [API token](https://pypi.org/manage/account/token/) do PyPI:
+
+```bash
+python -m pip install build twine
+python -m build
+python -m twine upload dist/*
+```
+
+Teste antes no [TestPyPI](https://test.pypi.org/).
+
+## Gemma 4 e Qwen 3.x: melhor caminho com este tipo de quantização
+
+TurboQuant (no paper e no blog do Google Research) ataca sobretudo **memória e custo do cache KV** e produtos internos em atenção — não é só “exportar pesos em INT4”.
+
+1. **Hoje (ecossistema Hugging Face / vLLM)**  
+   Integração “de verdade” costuma exigir **suporte no motor de inferência** (vLLM, SGLang, TensorRT-LLM, etc.) ou um **módulo de atenção customizado** que, a cada passo, comprima/descomprima **K/V por cabeça** com a mesma rotação e parâmetros fixos por modelo/camada.
+
+2. **Gemma (ex. Gemma 3 / família Gemma no Hub)** e **Qwen3**  
+   - Confirme `hidden_size`, **GQA** (`num_key_value_heads`), `head_dim` no `config.json` do checkpoint.  
+   - A lib aqui opera em vetores de tamanho `dim` (= `head_dim` por cabeça, ou blocos que vocês definirem).  
+   - Caminho pragmático: usar este pacote para **benchmark offline** (distorsão, memória simulada); em seguida portar o núcleo para **PyTorch custom op** ou **Triton** no caminho de escrita/leitura do KV.
+
+3. **Ordem sugerida**  
+   - Prototipar com esta API em **uma camada** ou tensor KV sintético.  
+   - Medir qualidade (long context / needle) com `bits_main` ~3–4 como no material de referência.  
+   - Só então acoplar ao servidor (vLLM plugin / patch) para latência real.
+
+## Licença
+
+Apache-2.0
+
+## Referências
+
+- [TurboQuant (arXiv)](https://arxiv.org/abs/2504.19874)  
+- [Google Research blog](https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/)

turboquant_llm-0.1.0/pyproject.toml

@@ -0,0 +1,52 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "turboquant-llm"
+version = "0.1.0"
+description = "Reference implementation of TurboQuant-style vector/KV quantization (rotation + scalar quant + QJL residual)."
+readme = "README.md"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+keywords = ["quantization", "llm", "kv-cache", "turboquant", "vector-quantization"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: Apache Software License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  "numpy>=1.22",
+]
+
+[project.optional-dependencies]
+torch = ["torch>=2.0"]
+dev = [
+  "pytest>=7.0",
+  "ruff>=0.4",
+]
+
+[project.urls]
+Homepage = "https://github.com/PestanaRobson/turboquant_llm"
+Documentation = "https://github.com/PestanaRobson/turboquant_llm#readme"
+Repository = "https://github.com/PestanaRobson/turboquant_llm"
+Issues = "https://github.com/PestanaRobson/turboquant_llm/issues"
+
+[tool.hatch.build.targets.sdist]
+include = ["/src", "/tests", "/README.md"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/turboquant_llm"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"

turboquant_llm-0.1.0/src/turboquant_llm/PACKAGE_ORIGIN.txt

@@ -0,0 +1,3 @@
+Nota interna (não é documentação voltada ao usuário; pode ser removida do wheel em builds futuros se preferir).
+
+Histórico: o trabalho foi iniciado no contexto Pestana Labs. O nome de distribuição e a marca voltados ao público são turboquant-llm / import turboquant_llm.

turboquant_llm-0.1.0/src/turboquant_llm/__init__.py

@@ -0,0 +1,11 @@
+"""TurboQuant-style vector quantization (reference implementation)."""
+
+from turboquant_llm.compressor import CompressedBatch, TurboQuantCompressor
+from turboquant_llm.config import TurboQuantConfig
+
+__all__ = [
+    "CompressedBatch",
+    "TurboQuantCompressor",
+    "TurboQuantConfig",
+]
+__version__ = "0.1.0"

turboquant_llm-0.1.0/src/turboquant_llm/compressor.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from turboquant_llm.config import TurboQuantConfig
+from turboquant_llm.rotation import random_orthogonal_matrix
+from turboquant_llm.scalar_quant import dequantize_global, quantize_global
+
+
+@dataclass
+class CompressedBatch:
+    """Compressed representation for a batch of row vectors ``x`` with shape ``(n, dim)``."""
+
+    # Stage 1: rotated + scalar-quantized
+    codes_main: np.ndarray
+    scale_main: float
+    amin_main: float
+    amax_main: float
+    # Stage 2: QJL signs on residual (columns are projection dimensions)
+    qjl_signs: np.ndarray | None
+    # Fixed per compressor instance (not stored per batch in a real system — kept for round-trip tests)
+    rotation: np.ndarray
+    jl_matrix: np.ndarray
+
+
+class TurboQuantCompressor:
+    """Reference TurboQuant-style pipeline: rotate → scalar quantize → 1-bit JL signs on residual.
+
+    This is intended for experiments and API design, not a drop-in production KV kernel.
+    Shapes: ``x`` is ``(..., dim)``; the last dimension must match ``dim``.
+    """
+
+    def __init__(self, dim: int, config: TurboQuantConfig | None = None) -> None:
+        if dim < 2:
+            raise ValueError("dim must be >= 2")
+        self.dim = dim
+        self.config = config or TurboQuantConfig()
+        rng = np.random.default_rng(self.config.seed)
+        self._R = random_orthogonal_matrix(dim, rng)
+        k = self.config.qjl_projection_dim
+        if k > 0:
+            # Typical JL scaling: entries ~ N(0, 1/dim)
+            self._J = rng.standard_normal((k, dim)) / np.sqrt(dim)
+        else:
+            self._J = np.zeros((0, dim), dtype=np.float64)
+
+    @property
+    def rotation(self) -> np.ndarray:
+        return self._R
+
+    @property
+    def jl_matrix(self) -> np.ndarray:
+        return self._J
+
+    def compress(self, x: np.ndarray) -> CompressedBatch:
+        x = np.asarray(x, dtype=np.float64)
+        if x.shape[-1] != self.dim:
+            raise ValueError(f"last dimension must be {self.dim}, got {x.shape[-1]}")
+        flat = x.reshape(-1, self.dim)
+        # Row vectors: x_rot = x @ R (maps through orthonormal columns of R)
+        xr = flat @ self._R
+
+        codes, sc, lo, hi = quantize_global(xr, bits=self.config.bits_main)
+        recon = dequantize_global(codes, sc, lo, hi)
+        residual = xr - recon
+
+        if self._J.shape[0] == 0:
+            signs = None
+        else:
+            # (n, k) — store 1-bit signs of JL projection of residual rows
+            proj = residual @ self._J.T
+            signs = (proj >= 0).astype(np.uint8)
+
+        return CompressedBatch(
+            codes_main=codes.reshape(x.shape),
+            scale_main=sc,
+            amin_main=lo,
+            amax_main=hi,
+            qjl_signs=signs.reshape(x.shape[:-1] + (self._J.shape[0],)) if signs is not None else None,
+            rotation=self._R,
+            jl_matrix=self._J,
+        )
+
+    def reconstruct_first_stage(self, batch: CompressedBatch) -> np.ndarray:
+        """Reconstruct after rotation + inverse quant (ignores QJL residual)."""
+        recon_flat = dequantize_global(
+            batch.codes_main.reshape(-1, self.dim),
+            batch.scale_main,
+            batch.amin_main,
+            batch.amax_main,
+        )
+        xr_hat = recon_flat
+        x_hat = xr_hat @ self._R.T
+        return x_hat.reshape(batch.codes_main.shape)

turboquant_llm-0.1.0/src/turboquant_llm/config.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class TurboQuantConfig:
+    """Hyperparameters for the reference compressor (aligned with the public TurboQuant narrative).
+
+    - ``bits_main``: scalar quantization after rotation (e.g. 3–4 for strong compression).
+    - ``qjl_bits``: extra sign bits from the JL projection of the *residual* (QJL-style stage).
+    """
+
+    bits_main: int = 4
+    qjl_projection_dim: int = 64
+    seed: int = 0
+
+    def __post_init__(self) -> None:
+        if self.bits_main < 1:
+            raise ValueError("bits_main must be >= 1")
+        if self.qjl_projection_dim < 0:
+            raise ValueError("qjl_projection_dim must be >= 0")

turboquant_llm-0.1.0/src/turboquant_llm/rotation.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def random_orthogonal_matrix(dim: int, rng: np.random.Generator) -> np.ndarray:
+    """Haar-random orthogonal matrix via QR (columns orthonormal). Shape (dim, dim)."""
+    a = rng.standard_normal((dim, dim))
+    q, _ = np.linalg.qr(a, mode="complete")
+    # QR is unique up to sign; stabilize column signs for reproducibility across platforms
+    s = np.sign(np.diag(q))
+    s[s == 0] = 1.0
+    q = q * s
+    return q.astype(np.float64, copy=False)

turboquant_llm-0.1.0/src/turboquant_llm/scalar_quant.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def quantize_global(
+    x: np.ndarray,
+    *,
+    bits: int,
+) -> tuple[np.ndarray, float, float, float]:
+    """Min/max uniform quantization over all elements of ``x``.
+
+    Integer codes are in ``[0, 2**bits - 1]``. Reconstruction:
+    ``x_hat = amin + codes * scale``.
+    """
+    if bits < 1:
+        raise ValueError("bits must be >= 1")
+    n_levels = 2**bits
+    eps = 1e-8
+    amin = float(np.min(x))
+    amax = float(np.max(x))
+    if amax - amin < eps:
+        scale = 1.0
+        codes = np.zeros(x.shape, dtype=np.int64)
+        return codes, scale, amin, amax
+    scale = (amax - amin) / (n_levels - 1)
+    codes = np.rint((x - amin) / scale).clip(0, n_levels - 1).astype(np.int64)
+    return codes, scale, amin, amax
+
+
+def dequantize_global(codes: np.ndarray, scale: float, amin: float, amax: float) -> np.ndarray:
+    _ = amax  # kept for API symmetry / future per-channel max
+    return amin + codes.astype(np.float64) * scale

turboquant_llm-0.1.0/tests/test_compressor.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import numpy as np
+
+from turboquant_llm import TurboQuantCompressor, TurboQuantConfig
+
+
+def test_compress_roundtrip_first_stage_small_dim() -> None:
+    dim = 16
+    rng = np.random.default_rng(42)
+    x = rng.standard_normal((4, dim))
+    cfg = TurboQuantConfig(bits_main=8, qjl_projection_dim=0, seed=0)
+    comp = TurboQuantCompressor(dim, cfg)
+    batch = comp.compress(x)
+    x_hat = comp.reconstruct_first_stage(batch)
+    err = np.linalg.norm(x - x_hat) / (np.linalg.norm(x) + 1e-8)
+    assert err < 0.15
+
+
+def test_qjl_signs_shape() -> None:
+    dim = 32
+    k = 12
+    x = np.random.default_rng(0).standard_normal((2, 3, dim))
+    cfg = TurboQuantConfig(bits_main=4, qjl_projection_dim=k, seed=1)
+    comp = TurboQuantCompressor(dim, cfg)
+    b = comp.compress(x)
+    assert b.qjl_signs is not None
+    assert b.qjl_signs.shape == (2, 3, k)