File size: 1,007 Bytes
371818e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | 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()
|