Release frozen ThunderWorld FP4 1.5x and 2x weights, calibration provenance and loader (part 2)
18e8cca verified | """K16 NVFP4 base/residual quantizers; FP32 outer factors.""" | |
| import triton | |
| import triton.language as tl | |
| def quantize(x, R: tl.constexpr, C: tl.constexpr, REPAIR: tl.constexpr): | |
| # A common outer factor for base and residual permits accumulator fusion. | |
| factor = tl.maximum(tl.max(tl.abs(x), 1) / 1536.0, 1.0e-12) | |
| z = x / factor[:, None] | |
| z3 = tl.reshape(z, (R, C // 16, 16)) | |
| scale = tl.maximum(tl.max(tl.abs(z3), 2) / 6.0, 0.001953125).to(tl.float8e4nv) | |
| normalized = tl.reshape(z3 / scale.to(tl.float32)[:, :, None], (R, C)) | |
| pairs = tl.reshape(normalized, (R, C // 2, 2)) | |
| lo, hi = tl.split(pairs) | |
| packed = tl.inline_asm_elementwise( | |
| "{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1; cvt.u32.u8 $0, tmp; }", | |
| constraints="=r,f,f", args=[lo, hi], dtype=tl.uint32, is_pure=True, pack=1).to(tl.uint8) | |
| if REPAIR: | |
| code = tl.reshape(tl.join(packed & 15, packed >> 4), (R, C)).to(tl.int32) | |
| mag = code & 7 | |
| val = tl.where(mag < 4, mag * 0.5, tl.where(mag < 6, mag - 2., 2. * mag - 8.)) | |
| val = tl.where(code >= 8, -val, val) | |
| residual = z3 - tl.reshape(val, (R, C // 16, 16)) * scale.to(tl.float32)[:, :, None] | |
| rs = tl.maximum(tl.max(tl.abs(residual), 2) / 6.0, 0.001953125).to(tl.float8e4nv) | |
| rn = tl.reshape(residual / rs.to(tl.float32)[:, :, None], (R, C // 2, 2)) | |
| rlo, rhi = tl.split(rn) | |
| rp = tl.inline_asm_elementwise( | |
| "{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1; cvt.u32.u8 $0, tmp; }", | |
| constraints="=r,f,f", args=[rlo, rhi], dtype=tl.uint32, is_pure=True, pack=1).to(tl.uint8) | |
| else: | |
| rp = packed | |
| rs = scale | |
| return packed, scale, rp, rs, factor | |
| def reciprocal(x): | |
| # All callers clamp scales/factors positive and normal before reaching here. | |
| return tl.inline_asm_elementwise("rcp.approx.ftz.f32 $0, $1;", constraints="=f,f", args=[x], dtype=tl.float32, is_pure=True, pack=1) | |
| def quantize_fast(x, R: tl.constexpr, C: tl.constexpr, REPAIR: tl.constexpr): | |
| # A common outer factor for base and residual permits accumulator fusion. | |
| factor = tl.maximum(tl.max(tl.abs(x), 1) * (1.0 / 1536.0), 1.0e-12) | |
| z = x * reciprocal(factor)[:, None] | |
| z3 = tl.reshape(z, (R, C // 16, 16)) | |
| scale = tl.maximum(tl.max(tl.abs(z3), 2) * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv) | |
| normalized = tl.reshape(z3 * reciprocal(scale.to(tl.float32))[:, :, None], (R, C)) | |
| pairs = tl.reshape(normalized, (R, C // 2, 2)) | |
| lo, hi = tl.split(pairs) | |
| packed = tl.inline_asm_elementwise( | |
| "{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1; cvt.u32.u8 $0, tmp; }", | |
| constraints="=r,f,f", args=[lo, hi], dtype=tl.uint32, is_pure=True, pack=1).to(tl.uint8) | |
| if REPAIR: | |
| code = tl.reshape(tl.join(packed & 15, packed >> 4), (R, C)).to(tl.int32) | |
| mag = code & 7 | |
| val = tl.where(mag < 4, mag * 0.5, tl.where(mag < 6, mag - 2., 2. * mag - 8.)) | |
| val = tl.where(code >= 8, -val, val) | |
| residual = z3 - tl.reshape(val, (R, C // 16, 16)) * scale.to(tl.float32)[:, :, None] | |
| rs = tl.maximum(tl.max(tl.abs(residual), 2) * (1.0 / 6.0), 0.001953125).to(tl.float8e4nv) | |
| rn = tl.reshape(residual * reciprocal(rs.to(tl.float32))[:, :, None], (R, C // 2, 2)) | |
| rlo, rhi = tl.split(rn) | |
| rp = tl.inline_asm_elementwise( | |
| "{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1; cvt.u32.u8 $0, tmp; }", | |
| constraints="=r,f,f", args=[rlo, rhi], dtype=tl.uint32, is_pure=True, pack=1).to(tl.uint8) | |
| else: | |
| rp = packed | |
| rs = scale | |
| return packed, scale, rp, rs, factor | |
| def quantize_p(x, R:tl.constexpr, C:tl.constexpr, REPAIR:tl.constexpr, FAST:tl.constexpr): | |
| if FAST:return quantize_fast(x,R,C,REPAIR) | |
| else:return quantize(x,R,C,REPAIR) | |
| def quantize_linear(x, factor, R: tl.constexpr, C: tl.constexpr, REPAIR: tl.constexpr): | |
| # A common outer factor for base and residual permits accumulator fusion. | |
| z = x / factor | |
| z3 = tl.reshape(z, (R, C // 16, 16)) | |
| scale = tl.maximum(tl.max(tl.abs(z3), 2) / 6.0, 0.001953125).to(tl.float8e4nv) | |
| normalized = tl.reshape(z3 / scale.to(tl.float32)[:, :, None], (R, C)) | |
| pairs = tl.reshape(normalized, (R, C // 2, 2)) | |
| lo, hi = tl.split(pairs) | |
| packed = tl.inline_asm_elementwise( | |
| "{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1; cvt.u32.u8 $0, tmp; }", | |
| constraints="=r,f,f", args=[lo, hi], dtype=tl.uint32, is_pure=True, pack=1).to(tl.uint8) | |
| if REPAIR: | |
| code = tl.reshape(tl.join(packed & 15, packed >> 4), (R, C)).to(tl.int32) | |
| mag = code & 7 | |
| val = tl.where(mag < 4, mag * 0.5, tl.where(mag < 6, mag - 2., 2. * mag - 8.)) | |
| val = tl.where(code >= 8, -val, val) | |
| residual = z3 - tl.reshape(val, (R, C // 16, 16)) * scale.to(tl.float32)[:, :, None] | |
| rs = tl.maximum(tl.max(tl.abs(residual), 2) / 6.0, 0.001953125).to(tl.float8e4nv) | |
| rn = tl.reshape(residual / rs.to(tl.float32)[:, :, None], (R, C // 2, 2)) | |
| rlo, rhi = tl.split(rn) | |
| rp = tl.inline_asm_elementwise( | |
| "{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, $2, $1; cvt.u32.u8 $0, tmp; }", | |
| constraints="=r,f,f", args=[rlo, rhi], dtype=tl.uint32, is_pure=True, pack=1).to(tl.uint8) | |
| else: | |
| rp = packed | |
| rs = scale | |
| return packed, scale, rp, rs, factor | |