import json,sys
from pathlib import Path
import torch,triton
ROOT=Path(__file__).parent;sys.path.insert(0,str(ROOT.parent/'release'))
from fp8_runtime import _quantize_rows
@torch.inference_mode()
def main():
 rows=[]
 for m,k in [(4096,4096),(4096,12288),(16384,4096),(16384,12288)]:
  x=torch.randn(m,k,device='cuda',dtype=torch.bfloat16);s=torch.rand(k,device='cuda')+.05
  q=torch.empty_like(x,dtype=torch.float8_e4m3fn);a=torch.empty(m,1,device='cuda')
  refq=torch.empty_like(q);refa=torch.empty_like(a)
  _quantize_rows[(m,)](x,s,refq,refa,k,triton.next_power_of_2(k),num_warps=8)
  for warps in [4,8,16,32]:
   def fn():_quantize_rows[(m,)](x,s,q,a,k,triton.next_power_of_2(k),num_warps=warps)
   fn();ms=triton.testing.do_bench(fn)
   row={'m':m,'k':k,'warps':warps,'ms':ms,'codes_equal':torch.equal(q.view(torch.uint8),refq.view(torch.uint8)),'scales_equal':torch.equal(a,refa)}
   rows.append(row);print(row,flush=True)
 (ROOT/'quantizer-trial.json').write_text(json.dumps(rows,indent=2))
main()