File size: 5,471 Bytes
18e8cca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """K16 NVFP4 base/residual quantizers; FP32 outer factors."""
import triton
import triton.language as tl
@triton.jit
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
@triton.jit
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)
@triton.jit
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
@triton.jit
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)
@triton.jit
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
|