KV-cache quantization (Geron Ch 17, morie.fn hmkvc)

Usage

morie_geron_kv_cache_compress(
  K,
  V,
  n_bits = 8,
  per_head = TRUE,
  dtype_bytes = 2
)

Arguments

K, V: Arrays (n_heads, seq_len, d_head).
n_bits: Bit width.
per_head: Per-head scales.
dtype_bytes: Bytes per element before compression.

Value

List with K_dequantized, V_dequantized, bytes_before, bytes_after, compression_ratio, max_error, estimate, n, method.

Examples

M <- matrix(c(1, 2, 3, 4, 5, 6), nrow = 2)
morie_geron_kv_cache_compress(M, M)
#> $K_dequantized
#> , , 1
#> 
#>          [,1]     [,2]
#> [1,] 0.992126 1.984252
#> 
#> , , 2
#> 
#>          [,1]     [,2]
#> [1,] 3.023622 4.015748
#> 
#> , , 3
#> 
#>          [,1] [,2]
#> [1,] 5.007874    6
#> 
#> 
#> $V_dequantized
#> , , 1
#> 
#>          [,1]     [,2]
#> [1,] 0.992126 1.984252
#> 
#> , , 2
#> 
#>          [,1]     [,2]
#> [1,] 3.023622 4.015748
#> 
#> , , 3
#> 
#>          [,1] [,2]
#> [1,] 5.007874    6
#> 
#> 
#> $scales
#> $scales$K
#> [1] 0.04724409
#> 
#> $scales$V
#> [1] 0.04724409
#> 
#> 
#> $bytes_before
#> [1] 24
#> 
#> $bytes_after
#> [1] 20
#> 
#> $compression_ratio
#> [1] 1.2
#> 
#> $max_error
#> [1] 0.02362205
#> 
#> $n_bits
#> [1] 8
#> 
#> $estimate
#> [1] 1.2
#> 
#> $n
#> [1] 12
#> 
#> $method
#> [1] "KV-cache quantization"
#>