Оценка памяти KV cache

В LLM inference KV cache хранит key и value в каждом transformer-слое для уже обработанных токенов.

Напишите функцию, которая возвращает размер KV cache в байтах.

Для каждого токена, каждого слоя и каждой KV-head нужно хранить два вектора: key и value.

Сигнатура

def kv_cache_bytes(
    batch_size: int,
    context_length: int,
    num_layers: int,
    num_kv_heads: int,
    head_dim: int,
    bytes_per_element: int = 2,
) -> int: