参考：可运行微型实验

"""Small CPU-only reference experiments. NumPy required; no external model/API.

This is an educational attention layer, not a complete LM or a serving engine.
Run: python MidPreparation/StudyGuide/tools/mini_lab.py
"""
from collections import Counter
from pathlib import Path
import json
import numpy as np

def stable_softmax(x):
    """Every row must contain at least one finite score."""
    m=np.max(x,axis=-1,keepdims=True)
    if not np.isfinite(m).all():
        raise ValueError('All-masked/invalid row: handle explicitly before softmax')
    e=np.exp(x-m)
    return e/e.sum(axis=-1,keepdims=True)

def merge(seq,pair,new_id):
    out=[];i=0
    while i<len(seq):
        if i+1<len(seq) and (seq[i],seq[i+1])==pair:
            out.append(new_id);i+=2
        else: out.append(seq[i]);i+=1
    return out

def train_byte_bpe(text,num_merges):
    """One-stream toy version: no regex, normalizer, or special tokens."""
    seq=list(text.encode('utf-8'));vocab={i:bytes([i]) for i in range(256)};rules=[]
    for _ in range(num_merges):
        counts=Counter(zip(seq,seq[1:]))
        if not counts:break
        # Deterministic tie-break by pair IDs.
        pair=min(counts,key=lambda x:(-counts[x],x));new_id=len(vocab)
        vocab[new_id]=vocab[pair[0]]+vocab[pair[1]]
        rules.append((pair,new_id));seq=merge(seq,pair,new_id)
    return vocab,rules,seq

def encode_bpe(text,rules):
    seq=list(text.encode('utf-8'))
    for pair,new_id in rules:seq=merge(seq,pair,new_id)
    return seq

def rope_bhtd(x,positions):
    """x [B,H,T,d], positions [T], adjacent pairs, float64 reference."""
    d=x.shape[-1]
    assert d>0 and d%2==0
    freq=10000.0**(-np.arange(0,d,2)/d)
    angle=np.asarray(positions)[:,None]*freq[None,:]
    cos=np.cos(angle)[None,None,:,:];sin=np.sin(angle)[None,None,:,:]
    out=np.empty_like(x,dtype=np.float64)
    even=x[...,0::2];odd=x[...,1::2]
    out[...,0::2]=even*cos-odd*sin
    out[...,1::2]=even*sin+odd*cos
    return out

def grouped_attention(q,k,v,q_positions,k_positions):
    """Q [B,Hq,Tq,d], KV [B,Hkv,Tk,d]; only k_pos <= q_pos visible."""
    hq=q.shape[1];hkv=k.shape[1];assert hq%hkv==0
    # Repeats are pedagogical temporaries; real kernels need not materialize them.
    kr=np.repeat(k,hq//hkv,axis=1);vr=np.repeat(v,hq//hkv,axis=1)
    scores=q@kr.swapaxes(-1,-2)/np.sqrt(q.shape[-1])
    allowed=np.asarray(k_positions)[None,:]<=np.asarray(q_positions)[:,None]
    scores=np.where(allowed[None,None,:,:],scores,-np.inf)
    a=stable_softmax(scores)
    return a@vr,a

def layer(x,weights,hq,hkv,position_start=0,cache=None):
    """Single GQA+RoPE layer, no FFN/norm. x [B,T,D]."""
    b,t,dmodel=x.shape;dh=dmodel//hq
    assert dmodel%hq==0
    wq,wk,wv,wo=weights
    pos=np.arange(position_start,position_start+t)
    q=(x@wq).reshape(b,t,hq,dh).transpose(0,2,1,3)
    k=(x@wk).reshape(b,t,hkv,dh).transpose(0,2,1,3)
    v=(x@wv).reshape(b,t,hkv,dh).transpose(0,2,1,3)
    q=rope_bhtd(q,pos);k=rope_bhtd(k,pos)
    if cache is None:allpos=pos
    else:
        oldk,oldv,oldpos=cache
        assert oldpos[-1]+1==position_start
        k=np.concatenate([oldk,k],axis=2);v=np.concatenate([oldv,v],axis=2)
        allpos=np.concatenate([oldpos,pos])
    z,a=grouped_attention(q,k,v,pos,allpos)
    z=z.transpose(0,2,1,3).reshape(b,t,dmodel)@wo
    return z,(k,v,allpos),a

def finite_gradient(fn,x,eps=1e-6):
    grad=np.zeros_like(x,dtype=float)
    for idx in np.ndindex(x.shape):
        plus=x.copy();minus=x.copy();plus[idx]+=eps;minus[idx]-=eps
        grad[idx]=(fn(plus)-fn(minus))/(2*eps)
    return grad

def main():
    results={}
    text='ababab 你好 abab'
    vocab,rules,seq=train_byte_bpe(text,4)
    encoded=encode_bpe(text,rules)
    decoded=b''.join(vocab[i] for i in encoded).decode('utf-8')
    assert decoded==text and encoded==seq
    results['bpe']={'text':text,'merges':rules,'encoded':encoded,'round_trip':decoded==text}

    rng=np.random.default_rng(324);b,t,dmodel,hq,hkv=2,7,16,4,2;dh=dmodel//hq
    weights=tuple(rng.normal(size=shape)/np.sqrt(dmodel) for shape in [(dmodel,dmodel),(dmodel,hkv*dh),(dmodel,hkv*dh),(dmodel,dmodel)])
    x=rng.normal(size=(b,t,dmodel))
    full,_,a=layer(x,weights,hq,hkv)
    # Prefill first 3 positions; decode one position at a time.
    prefix,cache,_=layer(x[:,:3],weights,hq,hkv);parts=[prefix]
    for i in range(3,t):
        z,cache,_=layer(x[:,i:i+1],weights,hq,hkv,i,cache);parts.append(z)
    streamed=np.concatenate(parts,axis=1)
    err=np.max(np.abs(full-streamed));assert err<1e-12
    assert np.allclose(a.sum(-1),1)
    assert np.allclose(a[...,np.triu_indices(t,1)[0],np.triu_indices(t,1)[1]],0)
    results['gqa_cache']={'max_full_vs_incremental_error':float(err),'cache_K_shape':list(cache[0].shape),'Q_heads':hq,'KV_heads':hkv}

    s=np.array([1.,.2,-.7]);expert_values=np.array([2.,-1.,4.])
    def top1_post_softmax(logits):
        ix=np.argmax(logits);return float(stable_softmax(logits[[ix]])@expert_values[[ix]])
    def full_softmax_truncate(logits):
        ix=np.argmax(logits);return float(stable_softmax(logits)[ix]*expert_values[ix])
    g1=finite_gradient(top1_post_softmax,s);g2=finite_gradient(full_softmax_truncate,s)
    assert np.allclose(g1,0) and np.linalg.norm(g2)>0
    results['router']={'top1_then_softmax_gradient':g1.tolist(),'full_softmax_then_truncate_gradient':g2.tolist()}

    # d_out=3, d_in=4, rank=2; L=1/2 ||(W+sBA)x-target||^2
    w=rng.normal(size=(3,4));la=rng.normal(size=(2,4));lb=rng.normal(size=(3,2));xin=rng.normal(size=4);target=rng.normal(size=3);scale=2.
    def loss_a(a):return .5*np.square((w+scale*lb@a)@xin-target).sum()
    def loss_b(bb):return .5*np.square((w+scale*bb@la)@xin-target).sum()
    y=(w+scale*lb@la)@xin;gw=np.outer(y-target,xin)
    ga=scale*lb.T@gw;gb=scale*gw@la.T
    assert np.allclose(ga,finite_gradient(loss_a,la),atol=1e-6)
    assert np.allclose(gb,finite_gradient(loss_b,lb),atol=1e-6)
    results['lora']={'analytic_gradients_match_finite_difference':True}

    teacher=np.array([.8,.2]);logits=np.array([.3,-.1])
    fn=lambda z:np.sum(teacher*(np.log(teacher)-np.log(stable_softmax(z))))
    analytic=stable_softmax(logits)-teacher
    assert np.allclose(finite_gradient(fn,logits),analytic,atol=1e-7)
    results['distillation']={'soft_CE_KL_gradient_verified':True}
    results['status']='All assertions passed; no production performance claims'
    out=Path(__file__).resolve().parents[1]/'sources'/'微型实验结果.json'
    out.write_text(json.dumps(results,ensure_ascii=False,indent=2),encoding='utf-8')
    print(json.dumps(results,ensure_ascii=False,indent=2))

if __name__=='__main__':main()