Embedded AI Engineer – KV Cache, Schedule, Unsloth, KTransformer

1. The formula for KV cache and how to calculate itFormula: ceil(max_model_len / block_size) * page_size_bytespage_size_bytes = 2 * block_size * num_kv_heads * head_size * dtype_sizeExplanation:block_size is the size of token chunks, with a default configuration of 16;max_model_len is the maximum input sequence length that the model can processnum_kv_heads refers to the N heads in the multi-head attention mechanismhidden_size is the size of the hidden layer (due to GQA, in config.jsonhead_sizeis the size of the concatenated queries of num_attention_heads.Here, we require the size of a single head, so we need to usehead_size/num_attention_heads)dtype_size refers to the data type, where one byte (B) is 8 bits, and Qwen2.5 uses FP16, which is 2 bytesSpecific calculation, taking Qwen2.5-72B as an example

https://www.modelscope.cn/models/Qwen/Qwen2.5-72B-Instruct/file/view/master/config.json?status=1

ceil(32768/16) = 2048

page_size_bytes = 2 * 16 * 8 * (8192/64) * 2 = 65536, the first 2 is due to K and V

ceil(32768/16) * page_size_bytes = 128MB

For num_hidden_layers=80, with a model of 80 layers, 128MB*80=10G, so if the GPU memory is insufficient, it will report an error: 10.00GIB KV cache is needed

Leave a Comment