| """W1b exactness/speed screens on an isolated current worker, no listeners. |
| |
| Owned process and GPU window only. All replacements are local to this process. |
| Materialised weights stay fp16. Timings include transforms, not GEMM alone. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from statistics import median |
| import time |
| from typing import TYPE_CHECKING |
|
|
| from .profile_support import PROCESS_LIMIT_MIB, require_gpu_window |
| from .experiment_support import paired_graphs |
| from .qualify_donor import differences |
|
|
| if TYPE_CHECKING: |
| import torch |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("phase", choices=("material", "wo", "reply", "mgemm", "misc", "experts")) |
| p.add_argument("root", type=Path) |
| p.add_argument("baseline", type=Path) |
| p.add_argument("out", type=Path) |
| p.add_argument("--subset", choices=("small", "shared", "dense", "head"), default="small") |
| p.add_argument("--donor-root", type=Path) |
| p.add_argument("--build", type=Path) |
| args = p.parse_args() |
| require_gpu_window() |
| import numpy as np |
| import torch |
| from .. import fused as FU |
| from ..checkpoint import Checkpoint |
| from ..exl3 import kernels as K |
| from ..model import DSparkDrafter, Exl3Linear |
| from ..server import GpuBackend |
| from ..window import WindowRange |
| torch.set_num_threads(2) |
| dev = torch.device("cuda", 0) |
| ck = Checkpoint(str(args.root)) |
| m = DSparkDrafter(ck, dev, slots=8) |
| saved = torch.load(args.baseline/"synthetic.pt", weights_only=True) |
| taps = saved["taps"].to(dev) |
| for slot, last in enumerate((300,500,700,900)): |
| pos = torch.arange(last-127,last+1,device=dev) |
| m.insert(torch.full_like(pos,slot),pos,m.main_x(taps[pos])) |
| args.out.parent.mkdir(parents=True,exist_ok=True) |
| output = args.out.open("x") |
|
|
| def record(row: dict[str, object]) -> None: |
| |
| import os |
| import subprocess |
| process = subprocess.run(["nvidia-smi", "--id="+os.environ["CUDA_VISIBLE_DEVICES"], |
| "--query-compute-apps=pid,used_gpu_memory", "--format=csv,noheader,nounits"], |
| capture_output=True,text=True,check=True).stdout |
| own = next(float(line.split(',')[1]) for line in process.splitlines() if int(line.split(',')[0])==os.getpid()) |
| if own >= PROCESS_LIMIT_MIB: |
| raise RuntimeError(f"W1b footprint {own} MiB exceeds ceiling") |
| row["process_mib"] = own |
| row["peak_allocated_mib"] = torch.cuda.max_memory_allocated()/2**20 |
| output.write(json.dumps(row)+"\n") |
| output.flush() |
| print(json.dumps({k:v for k,v in row.items() if k!='timing'}),flush=True) |
|
|
| class Material: |
| def __init__(self, lin: Exl3Linear): |
| self.lin = lin |
| self.in_features,self.out_features=lin.in_features,lin.out_features |
| self.w=torch.empty(self.in_features,self.out_features,dtype=torch.float16,device=dev) |
| K.ext().reconstruct(self.w,lin.trellis,lin.k,False,True) |
| def __call__(self,x: torch.Tensor,out: torch.Tensor|None=None) -> torch.Tensor: |
| x=x.to(torch.float16).contiguous() |
| xh=torch.empty_like(x) |
| K.ext().had_r_128(x,xh,self.lin.suh,None,1.) |
| y=torch.empty(x.shape[0],self.out_features,dtype=torch.float32,device=dev) |
| K.ext().hgemm(xh,self.w,y) |
| K.ext().had_r_128(y,y,None,self.lin.svh,1.) |
| if out is not None: |
| out.copy_(y) |
| return out |
| return y |
|
|
| gen=torch.Generator(device=dev).manual_seed(41) |
| if args.phase == "material": |
| attrs={"small":("wq_a","wkv"), "shared":("wq_a","wkv","sh_w1","sh_w2","sh_w3"), |
| "dense":("wq_a","wkv","wq_b","wo_b","sh_w1","sh_w2","sh_w3")} |
| changes=[] |
| if args.subset == "head": |
| changes=[(m,"head",m.head,Material(m.head))] |
| else: |
| changes=[(st,a,getattr(st,a),Material(getattr(st,a))) for st in m.stages for a in attrs[args.subset]] |
| extra=sum(new.w.numel()*2 for _,_,_,new in changes) |
| record({"phase":"material_budget","subset":args.subset,"decoded_bytes":extra, |
| "retains_original_packed":True,"margin_mib":512}) |
| for _,a,lin,new in changes: |
| for rows in (1,2,3,4,5,10,15,20): |
| require_gpu_window() |
| x=(torch.randn(rows,lin.in_features,generator=gen,device=dev)*.3).half() |
| got,ref=new(x),lin(x) |
| record({"phase":"material_op","attribute":a,"rows":rows, |
| "vs_current":differences(got,ref),"timing":paired_graphs(lambda:lin(x),lambda:new(x))}) |
| def baseline(): |
| for owner,a,lin,new in changes: setattr(owner,a,lin) |
| return m.draft(slots,anchor,last,5,n_out=n) |
| def candidate(): |
| for owner,a,lin,new in changes: setattr(owner,a,new) |
| return m.draft(slots,anchor,last,5,n_out=n) |
| for b,n in ((1,5),(2,5),(3,4),(4,4)): |
| require_gpu_window() |
| slots=torch.arange(b,device=dev) |
| last=torch.tensor((300,500,700,900)[:b],device=dev) |
| anchor=torch.tensor([0]*b,device=dev) |
| ri,rc=(v.clone() for v in baseline()) |
| ci,cc=candidate() |
| record({"phase":"material_full","B":b,"n_out":n,"ids":differences(ci,ri), |
| "confidence":differences(cc,rc),"timing":paired_graphs(baseline,candidate)}) |
| elif args.phase == "wo": |
| st=m.stages[0] |
| for b in (1,2,3,4): |
| rows=b*5 |
| x=torch.randn(rows,m.s.o_groups*4096,generator=gen,device=dev).bfloat16() |
| ref=FU.wo_a_fp8(x,st.wo_a8,st.wo_a_sc,m.s.o_groups) |
| for bn,bk,warps,stages in ((32,128,2,2),(32,128,4,2),(32,128,4,3),(16,128,2,3),(32,64,2,3),(32,256,2,2)): |
| require_gpu_window() |
| def base(): return FU.wo_a_fp8(x,st.wo_a8,st.wo_a_sc,m.s.o_groups) |
| def candidate(): return FU.wo_a_fp8(x,st.wo_a8,st.wo_a_sc,m.s.o_groups,BN=bn,BK=bk,num_warps=warps,num_stages=stages) |
| got=candidate() |
| record({"phase":"wo","B":b,"rows":rows,"config":[bn,bk,warps,stages], |
| "vs_current":differences(got,ref),"timing":paired_graphs(base,candidate)}) |
| elif args.phase == "reply": |
| require_gpu_window() |
| m.capture(batch_sizes=(1,2,3,4),ks=(5,),insert_rows=(8,16,32,64,128),n_outs=(1,2,3,4,5),pack_reply=True) |
| ordinary=GpuBackend(m,512,30720,graph_reply=False) |
| captured=GpuBackend(m,512,30720,graph_reply=True) |
| checked = 0 |
| for b in (1,2,3,4): |
| for n in (1,2,3,4,5): |
| for margin in (False,True): |
| for cache in ("hit","miss"): |
| require_gpu_window() |
| cand=torch.where((m.markov.slot_of_token>=0) if cache=="hit" else (m.markov.slot_of_token<0))[0][:b].cpu().tolist() |
| reqs=[(i,cand[i],lp+1,WindowRange(128,lp-127,lp),n) for i,lp in enumerate((300,500,700,900)[:b])] |
| ref=ordinary.draft(reqs,5,margin) |
| got=captured.draft(reqs,5,margin) |
| if not all(a is None and c is None or a is not None and c is not None and np.array_equal(a.view(np.uint8),c.view(np.uint8)) for a,c in zip(ref,got)): |
| raise RuntimeError(f"reply mismatch B{b} n{n} margin{margin} {cache}") |
| checked += 1 |
| |
| n=5 if b<3 else 4 |
| reqs=[(i,0,lp+1,WindowRange(128,lp-127,lp),n) for i,lp in enumerate((300,500,700,900)[:b])] |
| for cadence in (0,.042): |
| for trial in range(3): |
| pairs=[] |
| for i in range(45): |
| require_gpu_window() |
| values={} |
| for tag,backend in (("base",ordinary),("candidate",captured)) if i%2==0 else (("candidate",captured),("base",ordinary)): |
| if cadence: time.sleep(cadence) |
| t0=time.perf_counter_ns() |
| backend.draft(reqs,5,False) |
| values[tag]=(time.perf_counter_ns()-t0)/1000 |
| if i>=5: pairs.append({**values,"saving_us":values['base']-values['candidate']}) |
| record({"phase":"reply_timing","B":b,"trial":trial,"cadence_ms":cadence*1000,"preceding_reply_cases_bit_exact":checked, |
| "timing":{"pairs":pairs,"median":{k:median(r[k] for r in pairs) for k in pairs[0]}}}) |
| elif args.phase == "misc": |
| import triton |
| st=m.stages[0] |
| for b in (1,2,3,4): |
| require_gpu_window() |
| rows=b*5 |
| slots=torch.arange(b,device=dev) |
| last=torch.tensor((300,500,700,900)[:b],device=dev) |
| lo=last-127 |
| q=torch.randn(b,5,m.s.n_heads,512,generator=gen,device=dev).bfloat16() |
| kb=torch.randn(b,5,512,generator=gen,device=dev).bfloat16() |
| def base(): return FU.block_attention_ring(q,m.ring_kv[:,0],m.ring_pos,slots,lo,last,kb,st.attn_sink,m.scale) |
| ref=base() |
| for bm,do,bd,warps in ((16,256,64,4),(32,256,64,4),(16,128,64,4),(32,128,64,4),(16,256,32,8)): |
| def candidate(): return FU.block_attention_ring(q,m.ring_kv[:,0],m.ring_pos,slots,lo,last,kb,st.attn_sink,m.scale,BM=bm,DO=do,BD=bd,num_warps=warps) |
| record({"phase":"attention","B":b,"config":[bm,do,bd,warps],"vs_current":differences(candidate(),ref), |
| "timing":paired_graphs(base,candidate)}) |
| x=torch.randn(rows,5120,generator=gen,device=dev).bfloat16() |
| def base(): return FU.rms_norm(x,st.attn_norm,m.eps) |
| ref=base() |
| for warps in (1,2,8): |
| def candidate(): |
| out=torch.empty_like(x) |
| FU._rms_norm_kernel[(rows,)](x,st.attn_norm,out,5120,m.eps,BD=1024,num_warps=warps,enable_fp_fusion=False) |
| return out |
| record({"phase":"norm","B":b,"warps":warps,"vs_current":differences(candidate(),ref),"timing":paired_graphs(base,candidate)}) |
| stream=torch.randn(rows,4,5120,generator=gen,device=dev).bfloat16() |
| fn,scale,bias=st.hc_attn |
| original=FU._HC_SPLITS |
| def base(): |
| FU._HC_SPLITS=original |
| return FU.hc_mixes(stream,fn,scale,bias,m.eps,m.s.hc_eps,m.s.hc_iters) |
| ref=tuple(v.clone() for v in base()) |
| for ns in (10,20,80): |
| def candidate(): |
| FU._HC_SPLITS=ns |
| return FU.hc_mixes(stream,fn,scale,bias,m.eps,m.s.hc_eps,m.s.hc_iters) |
| got=candidate() |
| record({"phase":"mhc","B":b,"splits":ns,"vs_current":[differences(a,c) for a,c in zip(got,ref)], |
| "timing":paired_graphs(base,candidate)}) |
| FU._HC_SPLITS=original |
| elif args.phase == "experts": |
| if args.donor_root is None or args.build is None: |
| p.error("experts requires --donor-root and --build") |
| from .donor_adapter import CudaExl3Dense |
| donor=CudaExl3Dense(args.donor_root,args.build) |
| st=m.stages[0] |
| matrices={} |
| for e in range(60): |
| matrices[e]={} |
| for name in ("w1","w3","w2"): |
| tr,su,sv,_=ck.exl3(f"mtp.0.ffn.experts.{e}.{name}") |
| matrices[e][name]=(tr.to(dev),su.to(dev)[None],sv.to(dev)) |
| for rows in (1,2,3,4,5,10,15,20): |
| for pattern in ("shared","distinct"): |
| require_gpu_window() |
| x=(torch.randn(rows,5120,generator=gen,device=dev)*.3).half() |
| ids=(torch.arange(3).expand(rows,3) if pattern=="shared" else torch.arange(rows*3).view(rows,3)).to(dev) |
| weights=torch.full((rows,3),m.s.route_scale/3,device=dev,dtype=torch.float16) |
| sel=st.perm[ids].contiguous() |
| counts=torch.bincount(ids.flatten(),minlength=128).cpu().tolist() |
| populated=[(e,c) for e,c in enumerate(counts) if c] |
| def base(): return K.run_decode_coop(x,sel,weights,st.groups,m.coop,m.s.swiglu_limit,False,True)[:rows] |
| def candidate(): |
| |
| |
| order=torch.argsort(ids.flatten(),stable=True) |
| tokens=order//3 |
| contrib=torch.zeros(rows*3,5120,device=dev,dtype=torch.float32) |
| start=0 |
| for e,count in populated: |
| xx=x[tokens[start:start+count]] |
| def lin(name,xx): |
| tr,su,sv=matrices[e][name] |
| return donor.linear(xx,tr,su,sv,[sv.numel()]) |
| gate=lin("w1",xx).float().clamp(max=m.s.swiglu_limit) |
| up=lin("w3",xx).float().clamp(-m.s.swiglu_limit,m.s.swiglu_limit) |
| down=lin("w2",(torch.nn.functional.silu(gate)*up).half()).float() |
| pick=order[start:start+count] |
| contrib[pick]=down*weights.flatten()[pick,None] |
| start+=count |
| out=torch.zeros(rows,5120,device=dev,dtype=torch.float32) |
| for pick in range(3): out.add_(contrib.view(rows,3,5120)[:,pick]) |
| return out |
| ref=base().clone() |
| got=candidate() |
| record({"phase":"experts","rows":rows,"pattern":pattern,"counts":counts, |
| "vs_current":differences(got,ref),"static_synthetic_group_counts":True, |
| "timing":paired_graphs(base,candidate)}) |
| else: |
| |
| st=m.stages[0] |
| for rows in (1,2,3,4,5,10,15,20): |
| require_gpu_window() |
| x=(torch.randn(rows,5120,generator=gen,device=dev)*.3).half() |
| linears=(st.sh_w1,st.sh_w3) |
| ptr=lambda ts:torch.tensor([t.data_ptr() for t in ts],dtype=torch.int64,device=dev) |
| tr,su,sv=(ptr([getattr(l,a) for l in linears]) for a in ('trellis','suh','svh')) |
| y=torch.empty(2,rows,2304,dtype=torch.float32,device=dev) |
| xh=torch.empty(2,rows,5120,dtype=torch.float16,device=dev) |
| def base(): return tuple(l(x) for l in linears) |
| def candidate(): |
| K.ext().exl3_mgemm(x[None],tr,y,su,xh,sv,None,None,5,0,False,True,-1,-1,0) |
| return y |
| ref=base() |
| got=candidate() |
| record({"phase":"mgemm_pair","rows":rows,"vs_current":[differences(got[i],ref[i]) for i in (0,1)], |
| "timing":paired_graphs(base,candidate)}) |
| record({"phase":"complete","kind":args.phase}) |
| output.close() |
|
|
| if __name__ == "__main__": main() |
|
|
