# Real-query attention capture, Qwen3.5-4B (HF bf16 weights, NF4 4-bit on GPU), chunked prefill.
import os, sys, json, time, math, numpy as np, torch
from transformers import AutoTokenizer, AutoModelForImageTextToText, BitsAndBytesConfig, AttentionInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
CTX=int(sys.argv[1]); CHUNK=int(sys.argv[2]) if len(sys.argv)>2 else 2048
NQ=64; OUT=f"out_{CTX}"; os.makedirs(OUT,exist_ok=True)
torch.cuda.set_per_process_memory_fraction(min(1.0,6.0*1024**3/torch.cuda.get_device_properties(0).total_memory))
MP=os.environ.get("MP","Qwen/Qwen3.5-4B")
tok=AutoTokenizer.from_pretrained(MP)
from datasets import load_dataset
txt="\n".join(load_dataset("Salesforce/wikitext","wikitext-2-raw-v1",split="test")["text"])
ids=tok(txt,return_tensors="pt").input_ids[0]
reps=math.ceil(CTX/len(ids)); ids=ids.repeat(reps)[:CTX]   # wikitext-2 test ~300K tokens, no repeat for <=256K
print("tokens",len(ids),"corpus_tokens",len(ids)//reps if reps else 0,"reps",reps,flush=True)
# sampled query positions: spread across context, skip first 1024
qpos=np.unique(np.linspace(1024,CTX-1,NQ).astype(int)); qset=set(qpos.tolist())
Qs={}; Ks={}
def cap(module,q,k,v,mask,**kw):
    L=module.layer_idx; T=q.shape[2]; S=k.shape[2]; start=S-T
    sel=[p-start for p in qpos if start<=p<S]
    if sel: Qs.setdefault(L,[]).append(q[0,:,sel,:].float().cpu())
    if S==CTX or start+T==CTX: Ks[L]=k[0].float().cpu()
    return sdpa_attention_forward(module,q,k,v,mask,**kw)
AttentionInterface.register("capture",cap)
bnb=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16)
t0=time.time()
m=AutoModelForImageTextToText.from_pretrained(MP,quantization_config=bnb,attn_implementation="capture",dtype=torch.bfloat16,device_map={"":0})
m.eval(); lm=m.model.language_model
print("loaded",round(time.time()-t0,1),"s mem",torch.cuda.memory_allocated()/2**20,flush=True)
torch.cuda.reset_peak_memory_stats(); pkv=None; t0=time.time()
with torch.no_grad():
    for s in range(0,CTX,CHUNK):
        x=ids[s:s+CHUNK].unsqueeze(0).cuda()
        o=lm(input_ids=x,past_key_values=pkv,use_cache=True,cache_position=torch.arange(s,s+x.shape[1],device="cuda"))
        pkv=o.past_key_values; del o
        if s//CHUNK%8==0: print("pos",s+x.shape[1],"peak MiB",torch.cuda.max_memory_allocated()//2**20,"t",round(time.time()-t0,1),flush=True)
prefill=time.time()-t0; peak=torch.cuda.max_memory_allocated()/2**20; peakr=torch.cuda.max_memory_reserved()/2**20
del pkv; torch.cuda.empty_cache()
layers=sorted(Ks); print("kv layers",layers,flush=True)
scale=256**-0.5; res={}; Qsave=[]
def conc(p):  # p: [..., S] probs over causal keys
    sp,_=torch.sort(p,dim=-1,descending=True); c=sp.cumsum(-1); n=p.shape[-1]
    f90=((c<0.9).sum(-1)+1).float()/n; f99=((c<0.99).sum(-1)+1).float()/n
    return f90,f99,c[...,0],c[...,min(7,n-1)],c[...,min(63,n-1)],((c<0.9).sum(-1)+1).float(),((c<0.99).sum(-1)+1).float()
for L in layers:
    Q=torch.cat(Qs[L],1)  # [16,nq,256]
    Qsave.append(Q.half().numpy()); K=Ks[L].cuda()  # [4,S,256]
    rows=[]
    for i,p in enumerate(qpos):
        for h in range(16):
            g=h//4  # GQA: q heads 4g..4g+3 -> kv head g (repeat_kv order)
            sc=(K[g,:p+1]@Q[h,i].cuda())*scale
            pr=torch.softmax(sc.float(),-1)
            # proxy-comparable: exclude self, renormalize
            pr_ns=pr[:-1]/pr[:-1].sum()
            a=[t.item() for t in conc(pr)]; b=[t.item() for t in conc(pr_ns)]
            rows.append([L,h,g,int(p)]+a+b+[pr[-1].item(),pr[:4].sum().item(),pr[-64:].sum().item()])
    res[L]=rows; print("layer",L,"done",flush=True)
cols=["layer","qhead","kvgroup","pos","f90","f99","top1","top8","top64","n90","n99","f90_noself","f99_noself","top1_noself","top8_noself","top64_noself","n90_noself","n99_noself","self_mass","first4_sink_mass","last64_local_mass"]
A=np.array([r for L in layers for r in res[L]])
def agg(M):
    d={}
    for j,c in enumerate(cols[4:],4): v=M[:,j]; d[c]={"mean":float(v.mean()),"median":float(np.median(v)),"p10":float(np.percentile(v,10)),"p90":float(np.percentile(v,90))}
    return d
out={"model":"Qwen3.5-4B (HF bf16 weights, bitsandbytes NF4 4-bit, compute bf16)","machine":"dedicated test rig, RTX 3080 10GB","ctx":CTX,"chunk":CHUNK,
 "corpus":"wikitext-2-raw-v1 test, first tokens","query_positions":qpos.tolist(),"n_qheads":16,"n_kvheads":4,"gqa_ratio":4,"kv_layers_hf_index":layers,
 "scale":scale,"Q_is":"post q_norm, post RoPE (partial rotary 0.25), pre-scale","prefill_s":prefill,"peak_alloc_MiB":peak,"peak_reserved_MiB":peakr,
 "overall":agg(A),"by_layer":{int(L):agg(A[A[:,0]==L]) for L in layers},
 "by_pos_tercile":{t:agg(A[(A[:,3]>=lo)&(A[:,3]<hi)]) for t,(lo,hi) in {"early":(0,CTX/3),"mid":(CTX/3,2*CTX/3),"late":(2*CTX/3,CTX+1)}.items()},
 "columns":cols}
json.dump(out,open(f"{OUT}/attn_concentration_real.json","w"),indent=1)
np.save(f"{OUT}/Q_sample.npy",np.stack(Qsave)); np.save(f"{OUT}/rows.npy",A)
print("DONE peak",peak,peakr,"prefill",prefill,flush=True)