KV-cache memory footprint (Geron Ch 16, morie.fn grkvc)

bytes = ceil(seq * L * H * d * 2 * batch * bits / 8); note the parenthesised floor-division plus remainder bump, exactly as Python.

Usage

morie_geron_kv_cache_compression(
  seq_len,
  num_layers,
  num_heads,
  d_head,
  bits = 16,
  batch_size = 1,
  baseline_bits = 16
)

Arguments

seq_len, num_layers, num_heads, d_head

Positive integers.

bits

Storage width per value.

batch_size

Positive batch size.

baseline_bits

Reference width for the compression ratio.

Value

List with cache_bytes, megabytes, gigabytes, baseline_bytes, compression_ratio, bytes_per_token, n_values.

Examples

morie_geron_kv_cache_compression(2048, 32, 32, 128, bits = 8,
    batch_size = 2)
#> $cache_bytes
#> [1] 1073741824
#> 
#> $megabytes
#> [1] 1024
#> 
#> $gigabytes
#> [1] 1
#> 
#> $baseline_bytes
#> [1] 2147483648
#> 
#> $compression_ratio
#> [1] 2
#> 
#> $bytes_per_token
#> [1] 524288
#> 
#> $n_values
#> [1] 1073741824
#> 
#> $bits
#> [1] 8
#> 
#> $estimate
#> [1] 1073741824
#> 
#> $n
#> [1] 2048
#>