Integration API

Common

Shared integration configuration translation.

Backend capability and version detection.

Shared integration exceptions.

IntegrationError

Bases: RuntimeError

Base class for runtime integration failures.

OptionalDependencyError

Bases: IntegrationError

Raised when an optional backend package is required but unavailable.

ProductionModeError

Bases: IntegrationError

Raised when a requested production path cannot satisfy the contract.

Integration report payloads.

HuggingFace

HuggingFace DynamicCache wrapper with compressed storage.

CompressedDynamicCache

CompressedDynamicCache(cache: Any, *, head_dim: int, bits: int | None = 4, k_bits: int | None = None, v_bits: int | None = None, seed: int = 42, device: device | None = None, rotation_mode: RotationMode = RotationMode.RHT, model_config: Any | None = None)

Patch a Transformers cache and store compressed K/V rows.

The wrapper dequantizes for HuggingFace attention output, so this is a diagnostic route. Production decode should use packed page kernels.

Source code in turboquant/integration/hf/dynamic_cache.py

def __init__(
    self,
    cache: Any,
    *,
    head_dim: int,
    bits: int | None = 4,
    k_bits: int | None = None,
    v_bits: int | None = None,
    seed: int = 42,
    device: torch.device | None = None,
    rotation_mode: RotationMode = RotationMode.RHT,
    model_config: Any | None = None,
):
    resolved_k = k_bits if k_bits is not None else bits
    resolved_v = v_bits if v_bits is not None else bits
    if resolved_k is None or resolved_v is None:
        raise ValueError("provide bits or both k_bits and v_bits")
    self.cache = cache
    self.head_dim = int(head_dim)
    self.k_bits = _bits_to_mode(int(resolved_k))
    self.v_bits = _bits_to_mode(int(resolved_v))
    self.bits = self.k_bits
    self.seed = int(seed)
    self.device = device or torch.device("cpu")
    self.enabled = True
    self.diagnostic_label = "hf_dequant"
    self._original_update = cache.update
    self._original_get_seq_length = getattr(cache, "get_seq_length", None)
    self._key_layers: list[CompressedLayer | None] = []
    self._value_layers: list[CompressedLayer | None] = []
    self._dequant_keys: list[torch.Tensor | None] = []
    self._dequant_values: list[torch.Tensor | None] = []
    self._dtype = torch.float16
    self._codecs = HeadDimCodecs(
        k_bits=self.k_bits,
        v_bits=self.v_bits,
        seed=self.seed,
        device=self.device,
        rotation_mode=rotation_mode,
    )
    self._bypass_layers = _full_attention_layers(model_config)

    if hasattr(cache.update, "__self__") and isinstance(cache.update.__self__, CompressedDynamicCache):
        warnings.warn("cache is already wrapped by TurboQuant", UserWarning, stacklevel=2)

    cache.update = self._update
    if self._original_get_seq_length is not None:
        cache.get_seq_length = self._get_seq_length

Qwen cache integration helpers for HuggingFace.

HFDiagnosticCacheAdapter dataclass

HFDiagnosticCacheAdapter(config: KVQuantConfig, device: device, dtype: dtype = torch.float16)

Diagnostic adapter that stores compressed cache state and returns dense tensors to HF.

DynamicCachePatch

DynamicCachePatch(cache: Any, adapter: HFDiagnosticCacheAdapter)

Context manager that patches a Transformers cache update method for diagnostics.

Source code in turboquant/integration/hf/qwen.py

def __init__(self, cache: Any, adapter: HFDiagnosticCacheAdapter):
    self.cache = cache
    self.adapter = adapter
    self._original_update = None

HuggingFace Qwen capture helpers.

HuggingFace integration configuration.

vLLM

TurboQuant cache recipes used by vLLM integration code.

Recipe vector packing for vLLM TurboQuant cache entries.

Metadata JSON for vLLM TurboQuant channel groups.

Activation calibration for vLLM TurboQuant metadata.

Uint8 paged cache storage for vLLM-style slot mapping.

Paged compressed cache layout for vLLM-style slot mapping.

Paged compressed cache operations.

PagedCompressedKVCache

PagedCompressedKVCache(*, config: KVQuantConfig, spec: PagedCacheSpec, device: device, dtype: dtype = torch.float16)

Owns compressed KV blocks addressed by vLLM-style flat slots.

Source code in turboquant/integration/vllm/ops.py

def __init__(
    self,
    *,
    config: KVQuantConfig,
    spec: PagedCacheSpec,
    device: torch.device,
    dtype: torch.dtype = torch.float16,
):
    self.config = config
    self.spec = spec
    self.device = device
    self.dtype = dtype
    base_policy = policy_from_config(config)
    self.policy = type(base_policy)(
        preset=base_policy.preset,
        k_format=base_policy.k_format,
        v_format=base_policy.v_format,
        recent_window=0,
        boundary_tokens=0,
        preserve_first_layers=base_policy.preserve_first_layers,
        preserve_last_layers=base_policy.preserve_last_layers,
        sparse_v=base_policy.sparse_v,
        allow_v3=base_policy.allow_v3,
        allow_k_low_bits=base_policy.allow_k_low_bits,
        fallback_mode=base_policy.fallback_mode,
        fallback_count=base_policy.fallback_count,
        reason=base_policy.reason,
    )
    self._writer = CompressedKVCache(
        head_dim=spec.head_dim,
        num_kv_heads=spec.num_kv_heads,
        layer_idx=config.layer_idx,
        policy=self.policy,
        device=device,
        dtype=dtype,
        chunk_size=1,
        num_layers=config.num_layers,
        seed=config.seed,
    )
    self._slots: dict[int, CompressedKVBlock] = {}
    self._writes = 0
    self._workspace_bytes = 0

vLLM cache dtype registry for TurboQuant modes.

Runtime dispatch helpers for vLLM TurboQuant integration.

Model-level verification helpers for TurboQuant cache modes.

TurboQuant vLLM integration - thin adapter layer.

Responsibilities: - Detect layer/backend type (flash vs MLA/GDN) - Install minimal monkey-patches that delegate to capture/store/score - Expose clean modes: off | capture_only | hybrid | full_tq - Keep patching surface tiny; all real logic lives in capture/store

Modes: - off: no TQ activity, passthrough - capture_only: capture KV into compressed store, always use flash output - hybrid: use compressed history + exact recent for decode - full_tq: (future) TQ handles everything including prefill

LayerConfig dataclass

LayerConfig(head_dim: int, num_kv_heads: int, num_query_heads: int, key_bits: int = 3, value_bits: int = 2, value_group_size: int = 32, ring_capacity: int = 128, layer_idx: int = 0, backend_kind: str = 'flash', device: device = (lambda: torch.device('cuda'))())

Per-layer TQ configuration.

LayerState dataclass

LayerState(config: LayerConfig, store: object, engine: object, _log_count: int = 0)

Per-layer runtime state. Owns the capture engine and store.

supports_hybrid property

supports_hybrid: bool

Whether this layer supports hybrid attention.

reset

reset()

Reset state.

Source code in turboquant/integration/vllm/hooks.py

def reset(self):
    """Reset state."""
    self.engine.reset()
    self._log_count = 0

set_mode

set_mode(mode: str)

Set global TurboQuant mode.

Source code in turboquant/integration/vllm/hooks.py

def set_mode(mode: str):
    """Set global TurboQuant mode."""
    global _GLOBAL_MODE
    assert mode in _VALID_MODES, f"Invalid mode: {mode}. Valid: {_VALID_MODES}"
    _GLOBAL_MODE = mode
    logger.info(f"[TurboQuant] Mode set to: {mode}")

get_mode

get_mode() -> str

Get current global mode.

Source code in turboquant/integration/vllm/hooks.py

def get_mode() -> str:
    """Get current global mode."""
    return _GLOBAL_MODE

install_hooks

install_hooks(model_runner, key_bits: int = 3, value_bits: int = 2, value_group_size: int = 32, ring_capacity: int = 128, initial_layers_count: int = 4, initial_layers_key_bits: int = None, mode: str = MODE_CAPTURE_ONLY, no_alloc: bool = False) -> Dict[str, LayerState]

Install TurboQuant hooks on all attention layers in a vLLM model runner.

Returns: dict mapping layer_name -> LayerState

Source code in turboquant/integration/vllm/hooks.py

def install_hooks(
    model_runner,
    key_bits: int = 3,
    value_bits: int = 2,
    value_group_size: int = 32,
    ring_capacity: int = 128,
    initial_layers_count: int = 4,
    initial_layers_key_bits: int = None,
    mode: str = MODE_CAPTURE_ONLY,
    no_alloc: bool = False,
) -> Dict[str, LayerState]:
    """Install TurboQuant hooks on all attention layers in a vLLM model runner.

    Returns: dict mapping layer_name -> LayerState
    """
    global _GLOBAL_MODE
    _GLOBAL_MODE = mode

    if initial_layers_key_bits is None:
        initial_layers_key_bits = min(key_bits + 1, 4)

    static_ctx = model_runner.compilation_config.static_forward_context
    device = model_runner.device

    layer_states: Dict[str, LayerState] = {}
    layer_idx = 0

    for layer_name, attn_module in static_ctx.items():
        if not hasattr(attn_module, "impl"):
            continue

        impl = attn_module.impl
        num_kv_heads = getattr(impl, "num_kv_heads", None)
        if num_kv_heads is None:
            continue

        if hasattr(impl, "head_size"):
            head_dim = int(impl.head_size)
        elif hasattr(impl, "kv_lora_rank"):
            head_dim = int(impl.kv_lora_rank)
        else:
            continue

        bits = initial_layers_key_bits if layer_idx < initial_layers_count else key_bits
        backend_kind = "mla" if _is_mla_impl(impl) else "flash"
        num_query_heads = _infer_num_query_heads(attn_module, impl)

        cfg = LayerConfig(
            head_dim=head_dim,
            num_kv_heads=int(num_kv_heads),
            num_query_heads=num_query_heads,
            key_bits=bits,
            value_bits=value_bits,
            value_group_size=min(value_group_size, head_dim),
            ring_capacity=ring_capacity,
            layer_idx=layer_idx,
            backend_kind=backend_kind,
            device=device,
        )

        state = _create_layer_state(cfg)
        layer_states[layer_name] = state

        if backend_kind == "flash":
            has_separate_kv_update = hasattr(impl, "do_kv_cache_update")
            needs_forward_capture = not has_separate_kv_update

            if has_separate_kv_update:
                patched_update = _make_patched_kv_update(
                    impl.do_kv_cache_update.__func__, state, no_alloc=no_alloc
                )
                impl.do_kv_cache_update = types.MethodType(
                    lambda self, *a, _p=patched_update, **kw: _p(self, *a, **kw), impl
                )

            patched_forward = _make_patched_forward(
                impl.forward.__func__, state, no_alloc=no_alloc,
                capture_in_forward=needs_forward_capture,
            )
            impl.forward = types.MethodType(
                lambda self, *a, _p=patched_forward, **kw: _p(self, *a, **kw), impl
            )

            if needs_forward_capture and layer_idx == 0:
                logger.info(
                    "[TurboQuant] No do_kv_cache_update found (vLLM 0.16 FlashInfer); "
                    "capturing K/V in forward()"
                )
        else:
            if hasattr(impl, "do_kv_cache_update"):
                patched_update = _make_patched_mla_update(impl.do_kv_cache_update.__func__, state)
                impl.do_kv_cache_update = types.MethodType(
                    lambda self, *a, _p=patched_update, **kw: _p(self, *a, **kw), impl
                )
            if hasattr(impl, "forward_mqa"):
                patched_fwd = _make_patched_mla_forward(impl.forward_mqa.__func__, state)
                impl.forward_mqa = types.MethodType(
                    lambda self, *a, _p=patched_fwd, **kw: _p(self, *a, **kw), impl
                )

        impl._tq_layer_state = state
        layer_idx += 1

    model_runner._tq_layer_states = layer_states
    model_runner._tq_no_alloc = no_alloc
    logger.info(
        f"[TurboQuant] Hooks on {len(layer_states)} layers "
        f"(mode={mode}, no_alloc={no_alloc})"
    )
    return layer_states

free_kv_cache

free_kv_cache(model_runner) -> int

Free paged KV cache for TQ-hooked layers. Returns bytes freed.

Only frees layers that have TQ state. Non-TQ layers (MLA/GDN) keep their cache.

Source code in turboquant/integration/vllm/hooks.py

def free_kv_cache(model_runner) -> int:
    """Free paged KV cache for TQ-hooked layers. Returns bytes freed.

    Only frees layers that have TQ state. Non-TQ layers (MLA/GDN) keep their cache.
    """
    layer_states = getattr(model_runner, "_tq_layer_states", None)
    if not layer_states:
        logger.warning("[TurboQuant] No layer states found, nothing to free")
        return 0

    static_ctx = model_runner.compilation_config.static_forward_context
    device = model_runner.device
    freed = 0
    tiny = torch.zeros(1, dtype=torch.int8, device=device)

    ptrs_to_free = set()
    for layer_name, state in layer_states.items():
        if not state.supports_hybrid:
            continue
        if layer_name not in static_ctx:
            continue
        attn_module = static_ctx[layer_name]
        kv_list = getattr(attn_module, "kv_cache", None)
        if kv_list and len(kv_list) > 0:
            ptrs_to_free.add(kv_list[0].data_ptr())

    for layer_name, state in layer_states.items():
        if not state.supports_hybrid:
            continue
        if layer_name not in static_ctx:
            continue
        attn_module = static_ctx[layer_name]
        kv_list = getattr(attn_module, "kv_cache", None)
        if kv_list and len(kv_list) > 0:
            old = kv_list[0]
            freed += old.nelement() * old.element_size()
            kv_list[0] = tiny

    for i in range(len(model_runner.kv_caches)):
        entry = model_runner.kv_caches[i]
        if isinstance(entry, list):
            for j in range(len(entry)):
                if hasattr(entry[j], "data_ptr") and entry[j].data_ptr() in ptrs_to_free:
                    entry[j] = tiny
        elif hasattr(entry, "data_ptr") and entry.data_ptr() in ptrs_to_free:
            model_runner.kv_caches[i] = tiny

    torch.cuda.empty_cache()
    logger.info(f"[TurboQuant] Freed {freed / 1e6:.0f} MB KV cache ({len(layer_states)} layers)")
    return freed

get_stats

get_stats(model_runner) -> dict

Return summary statistics for all TQ layer states.

Source code in turboquant/integration/vllm/hooks.py

def get_stats(model_runner) -> dict:
    """Return summary statistics for all TQ layer states."""
    layer_states = getattr(model_runner, "_tq_layer_states", None)
    if not layer_states:
        return {}

    stats = {}
    total_compressed = 0
    total_buffered = 0
    total_memory = 0

    for name, state in layer_states.items():
        compressed = state.store.num_tokens
        buffered = state.engine.ring.size
        mem = state.store.memory_bytes()
        total_compressed += compressed
        total_buffered += buffered
        total_memory += mem

    stats["num_layers"] = len(layer_states)
    stats["total_compressed_tokens"] = total_compressed // max(len(layer_states), 1)
    stats["total_buffered_tokens"] = total_buffered // max(len(layer_states), 1)
    stats["total_memory_bytes"] = total_memory
    stats["mode"] = _GLOBAL_MODE
    return stats