"""Raw KV element counts only, not allocation or measured VRAM."""
import json
def raw_kv_bytes(layers,batch,length,kv_heads,head_dim,bytes_per_value):
    vals=(layers,batch,length,kv_heads,head_dim,bytes_per_value)
    if any(not isinstance(v,int) or isinstance(v,bool) or v<=0 for v in vals):raise ValueError("positive integers required")
    return 2*layers*batch*length*kv_heads*head_dim*bytes_per_value
if __name__=="__main__":
    a=raw_kv_bytes(2,1,8,2,4,2);b=raw_kv_bytes(2,1,9,2,4,2);c=raw_kv_bytes(2,2,8,2,4,2)
    assert (a,b,c)==(512,576,1024)
    lengths=[2,3,4]
    recompute=sum(lengths);reuse=lengths[0]+sum(lengths[i]-lengths[i-1] for i in range(1,len(lengths)))
    assert (recompute,reuse)==(9,4)
    print(json.dumps({"raw_bytes":[a,b,c],"one_token_delta":b-a,"token_positions_processed_to_build_KV_without_cache":recompute,"with_cache":reuse},indent=2))