# Least-squares fit of dedicated VRAM vs context per KV cache type (E21s numbers)
import numpy as np
d={'q4_0':[(4096,2961),(131072,4127),(262144,5281),(393216,6435),(524288,7589),(655360,8743),(720896,9321)],
   'q8_0':[(4096,2993),(131072,5151),(262144,7329),(393216,9507)],
   'f16':[(4096,3041),(131072,7011),(196608,9061)]}
bits={'q4_0':4.5,'q8_0':8.5,'f16':16}
S={}
for k,v in d.items():
    x=np.array([a for a,b in v],float); y=np.array([b for a,b in v],float)
    s,i=np.polyfit(x,y,1); S[k]=s*1048576
    print(k,'bytes/token',round(S[k],1),'intercept MiB',round(i,1))
A=np.array([[1,bits[k]/8] for k in S]); b=np.array([S[k] for k in S])
(n,E),*_=np.linalg.lstsq(A,b,rcond=None)
print('non-KV bytes/token',round(n,1),'KV elements/token',round(E))
icpt=2961-S['q4_0']*4096/1048576
for budget in (9500,9600):
    btok=budget-icpt   # MiB for 1,048,576 tokens equals bytes per token
    print(budget,'MiB budget: allowed bytes/token',round(btok,1),'KV bits/element',round((btok-n)*8/E,2))