| """Portable Triton unpacking for E2M1 + E4M3 block16 scaled FP4 weights.""" |
| import torch |
| import triton |
| import triton.language as tl |
|
|
|
|
| @triton.jit |
| def _unpack(W, S, G, O, SIZE:tl.constexpr, BLOCK:tl.constexpr): |
| i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK) |
| packed=tl.load(W+i//2, i<SIZE, other=0).to(tl.int32) |
| code=tl.where(i%2==0,packed&15,packed>>4) |
| magnitude=code&7 |
| value=tl.where(magnitude<4,magnitude.to(tl.float32)*.5, |
| tl.where(magnitude<6,magnitude.to(tl.float32)-2,(magnitude.to(tl.float32)-4)*2)) |
| value=tl.where((code&8)!=0,-value,value) |
| scale_byte=tl.load(S+i//16,i<SIZE,other=0).to(tl.int32) |
| mantissa=scale_byte&7 |
| exponent=(scale_byte>>3)&15 |
| power=((exponent+120)<<23).to(tl.float32,bitcast=True) |
| scale=tl.where(exponent==0,mantissa.to(tl.float32)*.001953125,(1+mantissa.to(tl.float32)*.125)*power) |
| result=value*(scale*tl.load(G)) |
| tl.store(O+i,result,i<SIZE) |
|
|
|
|
| def dequantize_weight(weight, scale, global_scale): |
| assert weight.is_cuda and weight.is_contiguous() and scale.is_contiguous() |
| n,k2=weight.shape |
| result=torch.empty((n,k2*2),device=weight.device,dtype=torch.bfloat16) |
| _unpack[(triton.cdiv(result.numel(),4096),)](weight,scale.view(torch.uint8),global_scale,result,result.numel(),4096, |
| num_warps=4,enable_fp_fusion=False) |
| return result |
|
|
