Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:Turboderp org Exllamav2 Ext Cache

From Leeroopedia
Knowledge Sources
Domains KV_Cache, Quantization, CUDA
Last Updated 2026-02-15 00:00 GMT

Overview

C++ extension functions for converting KV cache tensors between FP16, FP8, and quantized formats, with support for both contiguous and paged cache layouts.

Description

ext_cache.cpp implements the CUDA-backed format conversion routines that allow ExLlamaV2 to store key-value caches at reduced precision. The file provides five principal functions:

  • fp16_to_fp8 -- Converts FP16 cache tensors to FP8 (unsigned 8-bit) format in-place on GPU. Validates that input is kHalf and output is kUInt8, then dispatches to array_fp16_to_fp8_cuda. Operates on a strided 4D tensor layout with configurable batch_size, offset, and width for partial cache updates.
  • fp8_to_fp16 -- Reverses the FP8 conversion, restoring full FP16 precision. Same shape validation and strided access pattern as fp16_to_fp8.
  • fp16_to_q_kv -- Converts both K and V FP16 cache tensors to quantized format at the specified bit width (wbits). Supports two execution paths: (1) paged mode when page_size > 0, which uses block_table and cache_seqlens to index into a paged cache layout; (2) contiguous mode which operates on a flat strided layout with block-size alignment enforced via Q_CACHE_BLOCKSIZE_Q.
  • q_to_fp16_kv -- Dequantizes both K and V caches back to FP16 format, with the same paged/contiguous dual-path design as fp16_to_q_kv.
  • count_match -- A CPU-side function that compares two 1D tensors element-by-element (using 64-bit chunks for speed) and returns the length of the matching prefix. Used to determine how much of a cached sequence can be reused when a new prompt shares a common prefix with a previously cached prompt.

Each function uses OptionalCUDAGuard to ensure operations execute on the correct GPU device and obtains the current CUDA stream for asynchronous execution.

Usage

Use these functions from the Python-side cache classes (ExLlamaV2Cache_Q4, ExLlamaV2Cache_Q6, ExLlamaV2Cache_Q8) to convert KV tensors during the attention forward pass. Call fp16_to_fp8/fp8_to_fp16 for 8-bit caching, and fp16_to_q_kv/q_to_fp16_kv for 4-bit or 6-bit caching. Call count_match when checking prefix overlap for cache reuse in dynamic batching.

Code Reference

Source Location

Signature

void fp16_to_fp8(
    torch::Tensor in_tensor,
    torch::Tensor out_tensor,
    int batch_size,
    int offset,
    int width
);

void fp8_to_fp16(
    torch::Tensor in_tensor,
    torch::Tensor out_tensor,
    int batch_size,
    int offset,
    int width
);

void fp16_to_q_kv(
    torch::Tensor k_in,
    torch::Tensor k_out,
    torch::Tensor k_scales,
    torch::Tensor v_in,
    torch::Tensor v_out,
    torch::Tensor v_scales,
    int batch_size,
    int offset,
    int width,
    int page_size,
    torch::Tensor cache_seqlens,
    torch::Tensor block_table,
    int wbits
);

void q_to_fp16_kv(
    torch::Tensor k_in,
    torch::Tensor k_out,
    torch::Tensor k_scales,
    torch::Tensor v_in,
    torch::Tensor v_out,
    torch::Tensor v_scales,
    int batch_size,
    int offset,
    int width,
    int page_size,
    torch::Tensor cache_seqlens,
    torch::Tensor block_table,
    int wbits
);

int count_match(
    torch::Tensor a,
    torch::Tensor b,
    int max_a
);

Import

from exllamav2.ext import exllamav2_ext as ext_c

I/O Contract

Inputs

Parameter Type Description
in_tensor / k_in / v_in torch.Tensor (kHalf or kUInt8) Source cache tensor, 4D shape [batch, layers, heads, dim]
out_tensor / k_out / v_out torch.Tensor (kUInt8 or kHalf) Destination cache tensor, same shape as input
k_scales / v_scales torch.Tensor (kHalf) Per-block quantization scale factors (for q_kv functions)
batch_size int Number of sequences in the batch
offset int Starting position (in sequence dimension) for partial conversion
width int Number of positions to convert starting from offset
page_size int Page size for paged cache layout; 0 for contiguous mode
cache_seqlens torch.Tensor (kInt) Per-sequence cache lengths (paged mode)
block_table torch.Tensor (kInt) Block index table mapping logical to physical pages
wbits int Quantization bit width (4 or 6)
a, b torch.Tensor 1D tensors to compare (count_match)
max_a int Maximum prefix length to check (count_match)

Outputs

Function Return Description
fp16_to_fp8 void Writes FP8 values into out_tensor in-place
fp8_to_fp16 void Writes FP16 values into out_tensor in-place
fp16_to_q_kv void Writes quantized values and scales into k_out/v_out/k_scales/v_scales
q_to_fp16_kv void Writes dequantized FP16 values into k_out/v_out
count_match int Number of consecutive matching 64-bit elements from the start

Usage Examples

from exllamav2.ext import exllamav2_ext as ext_c

# Convert FP16 KV cache to FP8 for memory savings
ext_c.fp16_to_fp8(k_cache_fp16, k_cache_fp8, batch_size=4, offset=0, width=seq_len)

# Quantize KV cache to 4-bit with paged layout
ext_c.fp16_to_q_kv(
    k_in, k_out, k_scales,
    v_in, v_out, v_scales,
    batch_size=4, offset=0, width=new_tokens,
    page_size=256, cache_seqlens=seqlens, block_table=block_tbl,
    wbits=4
)

# Check how much of an existing cache can be reused
match_len = ext_c.count_match(old_token_ids, new_token_ids, max_a=old_seq_len)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment