import turboquant
print(turboquant.__version__)  # Should print "0.1.0"

from vllm import LLM, SamplingParams

# Use quant.cpp 3-bit KV cache
llm = LLM(
    model="meta-llama/Llama-3-8B",
    kv_cache_dtype="turbo3",        # quant.cpp 3-bit
    max_model_len=32768,
)

outputs = llm.generate(
    ["Tell me about quantum computing."],
    SamplingParams(temperature=0.7, max_tokens=256),
)

# integrations/vllm/tq_cache_engine.py

import turboquant
import torch
from vllm.worker.cache_engine import CacheEngine


class quant.cppCacheEngine(CacheEngine):
    """Custom vLLM cache engine using quant.cpp compression."""

    def __init__(self, cache_config, model_config, parallel_config):
        super().__init__(cache_config, model_config, parallel_config)

        self.tq_ctx = turboquant.quant.cppContext(
            backend=turboquant.BACKEND_CUDA
        )
        self.key_type = turboquant.TURBO_3B
        self.value_bits = 4
        self.head_dim = model_config.get_head_size()
        self.num_heads = model_config.get_num_kv_heads(parallel_config)

    def allocate_gpu_cache(self):
        """Allocate quantized KV cache on GPU.

        Each block stores quantized keys and values instead of FP16,
        reducing memory by ~5x.
        """
        gpu_cache = []
        fp16_bytes_per_token = self.head_dim * 2  # FP16 = 2 bytes
        tq_key_bytes = turboquant.type_bpe(self.key_type) * self.head_dim / 8
        tq_val_bytes = self.value_bits * self.head_dim / 8
        tq_bytes_per_token = tq_key_bytes + tq_val_bytes

        compression = fp16_bytes_per_token * 2 / tq_bytes_per_token
        print(f"[quant.cpp] KV cache compression: {compression:.1f}x")
        print(f"[quant.cpp] Key type: {turboquant.type_name(self.key_type)}")
        print(f"[quant.cpp] Value bits: {self.value_bits}")

        for layer_idx in range(self.num_layers):
            key_blocks = torch.zeros(
                self.num_gpu_blocks,
                self.block_size,
                self.num_heads,
                int(tq_key_bytes),
                dtype=torch.uint8,
                device="cuda",
            )
            value_blocks = torch.zeros(
                self.num_gpu_blocks,
                self.block_size,
                self.num_heads,
                int(tq_val_bytes),
                dtype=torch.uint8,
                device="cuda",
            )
            gpu_cache.append((key_blocks, value_blocks))

        return gpu_cache

    def swap_in(self, src_to_dst):
        """Swap quantized blocks from CPU to GPU."""
        for src, dst in src_to_dst.items():
            for layer_idx in range(self.num_layers):
                src_key, src_val = self.cpu_cache[layer_idx]
                dst_key, dst_val = self.gpu_cache[layer_idx]
                dst_key[dst].copy_(src_key[src])
                dst_val[dst].copy_(src_val[src])

    def swap_out(self, src_to_dst):
        """Swap quantized blocks from GPU to CPU."""
        for src, dst in src_to_dst.items():
            for layer_idx in range(self.num_layers):
                src_key, src_val = self.gpu_cache[layer_idx]
                dst_key, dst_val = self.cpu_cache[layer_idx]
                dst_key[dst].copy_(src_key[src])
                dst_val[dst].copy_(src_val[src])

    def close(self):
        """Release quant.cpp context."""
        if hasattr(self, 'tq_ctx') and self.tq_ctx is not None:
            self.tq_ctx.close()
            self.tq_ctx = None

import turboquant
import vllm.worker.cache_engine as ce

_orig_reshape = ce.reshape_and_cache

def tq_reshape_and_cache(key, value, key_cache, value_cache,
                          slot_mapping, kv_cache_dtype, kv_scale):
    if kv_cache_dtype == "turbo3":
        ctx = turboquant.quant.cppContext()
        # Quantize keys in-place
        for i, slot in enumerate(slot_mapping):
            k = key[i].cpu().numpy()
            v = value[i].cpu().numpy()
            qk = ctx.quantize_keys(k, turboquant.TURBO_3B)
            qv = ctx.quantize_values(v, bits=4)
            # Store in cache block at slot position
            # ... (implementation depends on vLLM version)
        ctx.close()
    else:
        _orig_reshape(key, value, key_cache, value_cache,
                       slot_mapping, kv_cache_dtype, kv_scale)

ce.reshape_and_cache = tq_reshape_and_cache

import time
import turboquant
import numpy as np

ctx = turboquant.quant.cppContext()
head_dim = 128

# Benchmark quantization latency
keys = np.random.randn(1000, head_dim).astype(np.float32)
t0 = time.perf_counter()
qdata = ctx.quantize_keys(keys, turboquant.TURBO_3B)
t1 = time.perf_counter()
print(f"Quantize 1000 keys: {(t1-t0)*1000:.2f} ms")

# Benchmark attention latency
query = np.random.randn(head_dim).astype(np.float32)
t0 = time.perf_counter()
scores = ctx.attention(query, qdata, 1000, turboquant.TURBO_3B)
t1 = time.perf_counter()
print(f"Attention (seq_len=1000): {(t1-t0)*1000:.2f} ms")

ctx.close()