参考：手算数值验证

"""CPU numerical checks for the study guide; no model downloads or training."""
from pathlib import Path
import json
import numpy as np

def softmax(x):
    e=np.exp(x-np.max(x,axis=-1,keepdims=True))
    return e/e.sum(axis=-1,keepdims=True)

def rope(x,p):
    x=np.asarray(x,dtype=np.float64); d=x.shape[-1]
    theta=10000.0**(-np.arange(0,d,2)/d); a=p*theta
    y=np.empty_like(x)
    y[...,0::2]=x[...,0::2]*np.cos(a)-x[...,1::2]*np.sin(a)
    y[...,1::2]=x[...,0::2]*np.sin(a)+x[...,1::2]*np.cos(a)
    return y

def main():
    result={}
    X=np.array([[1,0,1],[0,1,1]],dtype=float)
    Wq=np.array([[1,0],[0,1],[1,1]])
    Wk=np.array([[1,1],[1,0],[0,1]])
    Wv=np.array([[1,0],[0,1],[1,-1]])
    Q,K,V=X@Wq,X@Wk,X@Wv
    S=Q@K.T/np.sqrt(2);A=softmax(S)
    Sm=S.copy();Sm[np.triu_indices(2,1)]=-np.inf;Am=softmax(Sm)
    X2=X.copy();X2[1]=10
    S2=(X2@Wq)@(X2@Wk).T/np.sqrt(2)
    S2[np.triu_indices(2,1)]=-np.inf
    assert np.allclose((Am@V)[0],(softmax(S2)@(X2@Wv))[0])
    assert np.allclose(A.sum(-1),1) and np.allclose(Am.sum(-1),1)
    result['attention']={k:v.tolist() for k,v in dict(Q=Q,K=K,V=V,S=S,A=A,Z=A@V,A_causal=Am,Z_causal=Am@V).items()}
    q=np.array([1,2,3,4]); k=np.array([4,3,2,1])
    a=rope(q,1)@rope(k,3); b=rope(q,11)@rope(k,13); c=q@rope(k,2)
    assert np.allclose(a,b,atol=1e-12) and np.allclose(a,c,atol=1e-12)
    assert np.allclose(np.linalg.norm(rope(q,1)),np.linalg.norm(q))
    result['rope']={'q_position_1':rope(q,1).tolist(),'k_position_3':rope(k,3).tolist(),'dot_1_3':float(a),'dot_11_13':float(b)}
    scores=np.log([1.,2.,3.]); values=np.array([1.,2.,4.]);m=-np.inf;ell=0.;u=0.
    for inds in [[0,1],[2]]:
        s=scores[inds];new_m=max(m,s.max());factor=np.exp(m-new_m)
        weights=np.exp(s-new_m);ell=factor*ell+weights.sum();u=factor*u+weights@values[inds];m=new_m
    assert np.allclose(u/ell,softmax(scores)@values)
    result['online_softmax']={'output':float(u/ell),'expected':17/6}
    p=np.array([.6,.4]);draft=np.array([.8,.2]);accept=np.minimum(p,draft)
    residual=np.maximum(p-draft,0);reject=1-accept.sum()
    final=accept+reject*residual/residual.sum()
    assert np.allclose(final,p)
    result['speculative']={'acceptance_mass':accept.tolist(),'rejection_probability':float(reject),'final':final.tolist()}
    kv=2*32*8192*8*128*2
    assert kv==2**30
    result['resources']={'gqa_kv_bytes':kv,'gqa_kv_GiB':kv/2**30,'8B_bf16_GiB':8e9*2/2**30,'27B_bf16_GiB':27e9*2/2**30}
    capacity=int((20*2**30-16*10**9)//131072)
    assert capacity==41769
    qdist=np.array([.8,.2]);pdist=np.array([.5,.5])
    kl_forward=float(np.sum(qdist*np.log(qdist/pdist)))
    kl_reverse=float(np.sum(pdist*np.log(pdist/qdist)))
    result['additional_exercises']={'ideal_cache_tokens_q35':capacity,'KL_q_p':kl_forward,'KL_p_q':kl_reverse,'option1_MoE_layer_parameters':3*256*(8*96+64),'option2_MoE_layer_parameters':3*256*(8*128+128)}
    rewards=np.array([1,1,0,0],dtype=float);adv=(rewards-rewards.mean())/rewards.std()
    assert np.allclose(adv,[1,1,-1,-1])
    ratios=np.array([1.3,.7]);advantages=np.array([1.,-1.])
    terms=np.minimum(ratios*advantages,np.clip(ratios,.8,1.2)*advantages)
    assert np.allclose(terms,[1.2,-.8])
    result['grpo']={'population_std_advantages':adv.tolist(),'clip_terms':terms.tolist()}
    result['status']='All assertions passed'
    dest=Path(__file__).resolve().parents[1]/'sources'/'数值验证结果.json'
    dest.write_text(json.dumps(result,ensure_ascii=False,indent=2),encoding='utf-8')
    print(json.dumps(result,ensure_ascii=False,indent=2))

if __name__=='__main__': main()