turboquant-llm 0.1.0

Reference implementation of TurboQuant-style vector/KV quantization (rotation + scalar quant + QJL residual).

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.

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, blog Google Research) 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

pip install turboquant-llm

Opcional (integrações futuras com modelos):

pip install turboquant-llm[torch]

Uso rápido

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. Envie o conteúdo deste diretório para o branch principal (git push).

1. Com API token do PyPI:

python -m pip install build twine
python -m build
python -m twine upload dist/*

Teste antes no TestPyPI.

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)
• Google Research blog