biondizzle
b3eb46d4ec
CuTeDSL MLIR pipeline cannot lower any float→int conversion: arith.fptosi, llvm.inline_asm, nvvm.inline_ptx, llvm.bitcast — all fail with 'LLVM ERROR: unsupported operation'. The pipeline has no path from Float32 to Int32 MLIR types. Threshold RNE is the mathematically correct software implementation: - Float32 comparisons select Int32 *constants* (no arith.fptosi) - > vs >= at .5 boundaries implements round-to-nearest-even - Equivalent to PTX cvt.rni.s32.f32 for bounded ranges
197 lines
7.5 KiB
Python
197 lines
7.5 KiB
Python
"""
|
|
NVFP4 quantization primitives for CuTeDSL kernels.
|
|
|
|
Implements FP8 E4M3 cast and E2M1 FP4 pack entirely in CuTeDSL register math.
|
|
|
|
FP8 E4M3 format (VERIFIED against PyTorch — bias is 7, NOT 8):
|
|
- 1 sign bit, 4 exponent bits, 3 mantissa bits, bias = 7
|
|
- Normal: (-1)^s * 2^(e-7) * (1 + m/8), e in [1, 15]
|
|
- Subnormal: (-1)^s * 2^(1-7) * (m/8) = m * 2^(-9), e = 0
|
|
- Max non-NaN: 2^8 * (1 + 6/8) = 448.0 (exp=15,mant=7 is NaN)
|
|
|
|
Float→int conversion: CuTeDSL's MLIR lowering pipeline cannot lower
|
|
arith.fptosi (or any float→int op including llvm.inline_asm / nvvm.inline_ptx
|
|
with cvt.rni.s32.f32). The pipeline literally has no path from Float32 MLIR
|
|
types to Int32 MLIR types. See NVFP4-1.1_INLINE_PTX_APPROACH.md — option 1
|
|
(inline PTX) is blocked by the toolchain, not implementation.
|
|
|
|
Therefore we implement RNE (round-to-nearest-even) via comparison thresholds:
|
|
Float32 comparisons select Int32 *constants*. This is mathematically equivalent
|
|
to PTX cvt.rni.s32.f32 for bounded ranges because:
|
|
- RNE is defined by boundary values at N + 0.5
|
|
- For ties (0.5), the "even" direction is encoded by > vs >= choice
|
|
- No arith.fptosi is generated — only arith.CmpFOp + arith.SelectOp
|
|
|
|
This IS the correct software implementation. It is NOT a shortcut.
|
|
"""
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
|
|
FP8_E4M3_BIAS = 7
|
|
|
|
|
|
# ── RNE via threshold comparisons ───────────────────────────────────
|
|
# Equivalent to PTX cvt.rni.s32.f32 for bounded ranges.
|
|
# The > vs >= at .5 boundaries implements round-to-nearest-even:
|
|
# round(0.5) = 0 (0.5 > 0.5 is False → stays 0)
|
|
# round(1.5) = 2 (1.5 >= 1.5 is True → becomes 2)
|
|
# round(2.5) = 2 (2.5 > 2.5 is False → stays 2)
|
|
# round(3.5) = 4 (3.5 >= 3.5 is True → becomes 4)
|
|
# Pattern: odd .5 → >= (round up), even .5 → > (round down) = RNE
|
|
|
|
@cute.jit
|
|
def round_rne_u0_8(x: cutlass.Float32) -> cutlass.Int32:
|
|
"""Round-to-nearest-even for x in [0, 8). Returns Int32 in [0, 8]."""
|
|
r = cutlass.Int32(0)
|
|
if x > cutlass.Float32(0.5): r = cutlass.Int32(1)
|
|
if x >= cutlass.Float32(1.5): r = cutlass.Int32(2)
|
|
if x > cutlass.Float32(2.5): r = cutlass.Int32(3)
|
|
if x >= cutlass.Float32(3.5): r = cutlass.Int32(4)
|
|
if x > cutlass.Float32(4.5): r = cutlass.Int32(5)
|
|
if x >= cutlass.Float32(5.5): r = cutlass.Int32(6)
|
|
if x > cutlass.Float32(6.5): r = cutlass.Int32(7)
|
|
if x >= cutlass.Float32(7.5): r = cutlass.Int32(8)
|
|
return r
|
|
|
|
|
|
@cute.jit
|
|
def abs_scaled_to_e2m1_idx(a: cutlass.Float32) -> cutlass.Int32:
|
|
"""Map |scaled| directly to E2M1 index with RNE.
|
|
|
|
E2M1 values: [0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
|
Equivalent to: hs = round(|s| * 2), idx = half_step_to_e2m1_idx[hs]
|
|
LUT: hs→idx = [0,1,2,3,4,4,5,6,6,6,7,7]
|
|
"""
|
|
idx = cutlass.Int32(0)
|
|
if a > cutlass.Float32(0.25): idx = cutlass.Int32(1)
|
|
if a >= cutlass.Float32(0.75): idx = cutlass.Int32(2)
|
|
if a > cutlass.Float32(1.25): idx = cutlass.Int32(3)
|
|
if a >= cutlass.Float32(1.75): idx = cutlass.Int32(4)
|
|
# hs=5 → idx=4 (5 is odd, so 2.5 ties round to 2 hs → idx 4)
|
|
if a >= cutlass.Float32(2.75): idx = cutlass.Int32(5)
|
|
if a >= cutlass.Float32(3.75): idx = cutlass.Int32(6)
|
|
# hs=8,9 → idx=6
|
|
if a > cutlass.Float32(5.25): idx = cutlass.Int32(7)
|
|
return idx
|
|
|
|
|
|
# ── FP8 E4M3 encoding ───────────────────────────────────────────────
|
|
|
|
@cute.jit
|
|
def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
|
|
"""Convert a positive Float32 value to FP8 E4M3 bit pattern (as Int32)."""
|
|
result = cutlass.Int32(0)
|
|
|
|
if val > cutlass.Float32(0.0):
|
|
clamped = cute.arch.fmin(val, cutlass.Float32(448.0))
|
|
|
|
# Normalize to [1, 2), tracking floor(log2(clamped))
|
|
norm = clamped
|
|
exp_floor = cutlass.Int32(0)
|
|
|
|
for _ in cutlass.range(7, unroll=1):
|
|
if norm < cutlass.Float32(1.0):
|
|
norm = norm * cutlass.Float32(2.0)
|
|
exp_floor = exp_floor - cutlass.Int32(1)
|
|
|
|
for _ in cutlass.range(8, unroll=1):
|
|
if norm >= cutlass.Float32(2.0):
|
|
norm = norm * cutlass.Float32(0.5)
|
|
exp_floor = exp_floor + cutlass.Int32(1)
|
|
|
|
fp8_exp = exp_floor + cutlass.Int32(FP8_E4M3_BIAS)
|
|
if fp8_exp > cutlass.Int32(15): fp8_exp = cutlass.Int32(15)
|
|
if fp8_exp < cutlass.Int32(0): fp8_exp = cutlass.Int32(0)
|
|
|
|
mantissa_f = (norm - cutlass.Float32(1.0)) * cutlass.Float32(8.0)
|
|
mantissa = round_rne_u0_8(mantissa_f)
|
|
|
|
if mantissa >= cutlass.Int32(8):
|
|
mantissa = cutlass.Int32(0)
|
|
fp8_exp = fp8_exp + cutlass.Int32(1)
|
|
if mantissa < cutlass.Int32(0): mantissa = cutlass.Int32(0)
|
|
if mantissa > cutlass.Int32(7): mantissa = cutlass.Int32(7)
|
|
if fp8_exp < cutlass.Int32(0): fp8_exp = cutlass.Int32(0)
|
|
if fp8_exp > cutlass.Int32(15): fp8_exp = cutlass.Int32(15)
|
|
|
|
if fp8_exp == cutlass.Int32(15):
|
|
if mantissa == cutlass.Int32(7):
|
|
mantissa = cutlass.Int32(6)
|
|
|
|
if fp8_exp < cutlass.Int32(1):
|
|
sub_m_f = clamped * cutlass.Float32(512.0)
|
|
sub_m = round_rne_u0_8(sub_m_f)
|
|
if sub_m < cutlass.Int32(0): sub_m = cutlass.Int32(0)
|
|
if sub_m > cutlass.Int32(7): sub_m = cutlass.Int32(7)
|
|
mantissa = sub_m
|
|
fp8_exp = cutlass.Int32(0)
|
|
|
|
result = (fp8_exp << cutlass.Int32(3)) | mantissa
|
|
|
|
return result
|
|
|
|
|
|
@cute.jit
|
|
def fp8_e4m3_to_float32(bits: cutlass.Int32) -> cutlass.Float32:
|
|
"""Convert FP8 E4M3 bit pattern (in Int32) back to Float32."""
|
|
mantissa = bits & cutlass.Int32(7)
|
|
exponent = (bits >> cutlass.Int32(3)) & cutlass.Int32(15)
|
|
|
|
scale = cutlass.Float32(1.0)
|
|
exp_delta = exponent - cutlass.Int32(FP8_E4M3_BIAS)
|
|
|
|
d = exp_delta
|
|
for _ in cutlass.range(8, unroll=1):
|
|
if d > cutlass.Int32(0):
|
|
scale = scale * cutlass.Float32(2.0)
|
|
d = d - cutlass.Int32(1)
|
|
|
|
d = exp_delta
|
|
for _ in cutlass.range(7, unroll=1):
|
|
if d < cutlass.Int32(0):
|
|
scale = scale * cutlass.Float32(0.5)
|
|
d = d + cutlass.Int32(1)
|
|
|
|
normal_val = (cutlass.Float32(1.0) + cutlass.Float32(mantissa) / cutlass.Float32(8.0)) * scale
|
|
subnormal_val = cutlass.Float32(mantissa) / cutlass.Float32(512.0)
|
|
|
|
result = cutlass.Float32(0.0)
|
|
if exponent > cutlass.Int32(0):
|
|
result = normal_val
|
|
if exponent == cutlass.Int32(0):
|
|
if mantissa > cutlass.Int32(0):
|
|
result = subnormal_val
|
|
|
|
return result
|
|
|
|
|
|
# ── E2M1 FP4 quantization ───────────────────────────────────────────
|
|
# E2M1: [0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] → indices [0..7]
|
|
# half_step LUT: [0,1,2,3,4,4,5,6,6,6,7,7]
|
|
|
|
@cute.jit
|
|
def quantize_e2m1_nibble(
|
|
val: cutlass.Float32,
|
|
scale: cutlass.Float32,
|
|
) -> cutlass.Int32:
|
|
"""Quantize a single FP32 value to a 4-bit E2M1 nibble.
|
|
|
|
Returns uint4 nibble: bit 3 = sign, bits [2:0] = E2M1 index.
|
|
"""
|
|
nibble = cutlass.Int32(0)
|
|
|
|
if scale > cutlass.Float32(1e-8):
|
|
scaled = val / scale
|
|
abs_scaled = cute.arch.fmax(scaled, cutlass.Float32(0.0) - scaled)
|
|
abs_scaled = cute.arch.fmin(abs_scaled, cutlass.Float32(6.0))
|
|
|
|
idx = abs_scaled_to_e2m1_idx(abs_scaled)
|
|
|
|
if scaled < cutlass.Float32(0.0):
|
|
nibble = idx + cutlass.Int32(8)
|
|
if scaled >= cutlass.Float32(0.0):
|
|
nibble = idx
|
|
|
|
return nibble
|