257 lines
9.4 KiB
Python
257 lines
9.4 KiB
Python
"""
|
|
NVFP4 quantization primitives for CuTeDSL kernels.
|
|
|
|
Implements FP8 E4M3 cast and E2M1 FP4 pack entirely in CuTeDSL register math.
|
|
No shortcuts — proper bit-level quantization matching the Python/CUDA reference.
|
|
|
|
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)
|
|
- Min positive normal: 2^(-6) ≈ 0.015625
|
|
- Min positive subnormal: 2^(-9) ≈ 0.001953
|
|
|
|
CuTeDSL constraints:
|
|
- NO float-to-int conversion (arith.FloatToSIOp not lowerable to PTX)
|
|
- Use threshold rounding: Float32 comparisons to select Int32 constants
|
|
- `cute.arch.fmax`/`cute.arch.fmin` for float min/max (NOT cute.math.fmin/fmax)
|
|
- `cute.floor` for floor (returns Float32)
|
|
- `@cute.jit` decorator required for CuTeDSL functions with dynamic `if` blocks
|
|
- `cutlass.Int32(N)` creates Int32 constants; `cutlass.Float32(N)` creates Float32 constants
|
|
"""
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
|
|
FP8_E4M3_BIAS = 7
|
|
|
|
|
|
# ── Threshold rounding (avoids float-to-int conversion) ─────────────
|
|
# CuTeDSL cannot convert Float32 → Int32. Instead, we use Float32
|
|
# comparisons to select Int32 constants. This is correct because:
|
|
# 1. The ranges are small and bounded (mantissa 0-8, half_step 0-12)
|
|
# 2. Comparisons implement round-to-nearest-even when thresholds are
|
|
# placed at the 0.5 boundaries (e.g., 0.5, 1.5, 2.5, ...)
|
|
# 3. No arith.FloatToSIOp is generated — only arith.CmpFOp + arith.SelectOp
|
|
|
|
|
|
@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]. Uses threshold comparisons to avoid
|
|
float-to-int conversion. The > vs >= choice implements RNE:
|
|
- 0.5 rounds to 0 (0.5 > 0.5 is False → result stays 0)
|
|
- 1.5 rounds to 2 (1.5 >= 1.5 is True → result becomes 2)
|
|
This matches Python's round() and CUDA's __float2int_rn().
|
|
"""
|
|
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| value directly to E2M1 index with RNE.
|
|
|
|
Equivalent to: hs = round(|scaled| * 2), idx = half_step_to_e2m1_idx(hs)
|
|
but avoids float-to-int conversion entirely.
|
|
|
|
E2M1 values: [0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
|
|
|
|
Thresholds derived from half_step RNE boundaries:
|
|
- hs 0→1: |s| > 0.25 (0.5/2, RNE: round(0.5)=0)
|
|
- hs 1→2: |s| >= 0.75 (1.5/2, RNE: round(1.5)=2)
|
|
- hs 2→3: |s| > 1.25 (2.5/2, RNE: round(2.5)=2)
|
|
- hs 3→4: |s| >= 1.75 (3.5/2, RNE: round(3.5)=4)
|
|
- idx 4→5 at hs=6: |s| > 2.75 (5.5/2, RNE: round(5.5)=6)
|
|
- idx 5→6 at hs=8: |s| >= 3.75 (7.5/2, RNE: round(7.5)=8)
|
|
- idx 6→7 at hs=11: |s| > 5.25 (10.5/2, RNE: round(10.5)=10→idx 6; >10.5→11→idx 7)
|
|
"""
|
|
idx = cutlass.Int32(0)
|
|
if a > cutlass.Float32(0.25): # hs >= 1
|
|
idx = cutlass.Int32(1)
|
|
if a >= cutlass.Float32(0.75): # hs >= 2
|
|
idx = cutlass.Int32(2)
|
|
if a > cutlass.Float32(1.25): # hs >= 3
|
|
idx = cutlass.Int32(3)
|
|
if a >= cutlass.Float32(1.75): # hs >= 4
|
|
idx = cutlass.Int32(4)
|
|
# Note: hs 5 → idx 4 (half_step 5 maps to E2M1 idx 4, same as hs 4)
|
|
# So idx stays 4 for |s| in [1.75, 2.75]
|
|
if a >= cutlass.Float32(2.75): # hs >= 6 → idx 5
|
|
idx = cutlass.Int32(5)
|
|
if a >= cutlass.Float32(3.75): # hs >= 8 → idx 6
|
|
idx = cutlass.Int32(6)
|
|
if a > cutlass.Float32(5.25): # hs >= 11 → idx 7
|
|
idx = cutlass.Int32(7)
|
|
return idx
|
|
|
|
|
|
@cute.jit
|
|
def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
|
|
"""Convert a positive Float32 value to FP8 E4M3 bit pattern (returned as Int32).
|
|
|
|
Only handles positive values (NVFP4 scale factors are always positive).
|
|
Returns the uint8 bit pattern packed into an Int32.
|
|
"""
|
|
result = cutlass.Int32(0) # default: zero
|
|
|
|
if val > cutlass.Float32(0.0):
|
|
# Clamp to FP8 E4M3 max non-NaN value (exp=15, mant=6 = 448.0)
|
|
clamped = cute.arch.fmin(val, cutlass.Float32(448.0))
|
|
|
|
# Normalize to [1, 2) range, tracking floor(log2(clamped))
|
|
norm = clamped
|
|
exp_floor = cutlass.Int32(0)
|
|
|
|
# Double until >= 1 (at most 7 doublings needed, smallest normal ≈ 2^-6)
|
|
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)
|
|
|
|
# Halve until < 2 (at most 8 halvings needed, largest ≈ 240 < 256)
|
|
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 exponent = floor(log2(val)) + bias
|
|
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 for normal: (norm - 1) * 8, round via threshold
|
|
mantissa_f = (norm - cutlass.Float32(1.0)) * cutlass.Float32(8.0)
|
|
mantissa = round_rne_u0_8(mantissa_f)
|
|
|
|
# Mantissa overflow: rounded to 8 → increment exponent, reset mantissa
|
|
if mantissa >= cutlass.Int32(8):
|
|
mantissa = cutlass.Int32(0)
|
|
fp8_exp = fp8_exp + cutlass.Int32(1)
|
|
|
|
# Clamp mantissa to [0, 7]
|
|
if mantissa < cutlass.Int32(0):
|
|
mantissa = cutlass.Int32(0)
|
|
if mantissa > cutlass.Int32(7):
|
|
mantissa = cutlass.Int32(7)
|
|
|
|
# Clamp exponent to [0, 15]
|
|
if fp8_exp < cutlass.Int32(0):
|
|
fp8_exp = cutlass.Int32(0)
|
|
if fp8_exp > cutlass.Int32(15):
|
|
fp8_exp = cutlass.Int32(15)
|
|
|
|
# NaN guard: FP8 E4M3 with exp=15 and mant=7 is NaN.
|
|
# Saturate to max non-NaN (exp=15, mant=6 = 448.0).
|
|
if fp8_exp == cutlass.Int32(15):
|
|
if mantissa == cutlass.Int32(7):
|
|
mantissa = cutlass.Int32(6)
|
|
|
|
# Subnormal handling: if fp8_exp < 1, value is 2^(1-7) * m/8 = m * 2^(-9)
|
|
# m = round(clamped * 2^9) = round(clamped * 512)
|
|
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.
|
|
|
|
Normal: val = 2^(e-7) * (1 + m/8)
|
|
Subnormal (e=0): val = m * 2^(-9) = m / 512
|
|
"""
|
|
mantissa = bits & cutlass.Int32(7)
|
|
exponent = (bits >> cutlass.Int32(3)) & cutlass.Int32(15)
|
|
|
|
# Compute 2^(e-7) by iterative doubling/halving from 1.0
|
|
scale = cutlass.Float32(1.0)
|
|
exp_delta = exponent - cutlass.Int32(FP8_E4M3_BIAS)
|
|
|
|
# Double for positive delta (max delta=8, e=15)
|
|
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)
|
|
|
|
# Halve for negative delta (min delta=-7, e=0)
|
|
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 value
|
|
normal_val = (cutlass.Float32(1.0) + cutlass.Float32(mantissa) / cutlass.Float32(8.0)) * scale
|
|
|
|
# Subnormal value (e=0): val = m / 512
|
|
subnormal_val = cutlass.Float32(mantissa) / cutlass.Float32(512.0)
|
|
|
|
# Select
|
|
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
|
|
|
|
|
|
@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.
|
|
If scale ≈ 0, returns 0 (zero nibble).
|
|
"""
|
|
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
|