NVFP4-1.1: Replace threshold rounding with inline PTX cvt.rni/rz/rmi

- Add f32_to_i32_rni (cvt.rni.s32.f32) for round-to-nearest-even
- Add f32_to_i32_rz (cvt.rzi.s32.f32) for round-toward-zero
- Add f32_to_i32_rmi (cvt.rmi.s32.f32) for round-to-minus-infinity
- Replace round_rne_u0_8 and abs_scaled_to_e2m1_idx threshold hacks
  with proper PTX hardware rounding in fp8_e4m3_from_float32
- quantize_e2m1_nibble now uses f32_to_i32_rni + LUT logic for half_step
- Add test_ptx_convert.py for inline PTX conversion verification
- This is the CORRECT approach per NVFP4-1.1_INLINE_PTX_APPROACH.md option 1
This commit is contained in:
2026-05-28 04:40:17 +00:00
parent 2777ebfe8e
commit 1cbb3cf752
2 changed files with 240 additions and 77 deletions

View File

@@ -13,103 +13,106 @@ FP8 E4M3 format (VERIFIED against PyTorch — bias is 7, NOT 8):
- 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
- float-to-int via inline PTX cvt.rni.s32.f32 (proper hardware rounding)
- `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
- `@dsl_user_op` + `llvm.inline_asm` for PTX instructions not exposed by CuTeDSL
"""
import cutlass
import cutlass.cute as cute
from cutlass.cutlass_dsl import dsl_user_op, T
from cutlass._mlir.dialects import llvm
from cutlass.cute.typing import Float32, Int32
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
# ── Inline PTX float-to-int conversion ──────────────────────────────
# CuTeDSL has no built-in f32→i32 conversion. We wrap PTX cvt instructions
# via @dsl_user_op + llvm.inline_asm. This is the PROPER approach — hardware
# rounding, no threshold hacks, no approximation.
@cute.jit
def round_rne_u0_8(x: cutlass.Float32) -> cutlass.Int32:
"""Round-to-nearest-even for x in [0, 8).
@dsl_user_op
def f32_to_i32_rni(x: Float32, *, loc=None, ip=None) -> Int32:
"""Convert Float32 to Int32 with round-to-nearest-even (RNE).
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().
Wraps PTX: cvt.rni.s32.f32 $0, $1;
Equivalent to CUDA __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
return Int32(
llvm.inline_asm(
T.i32(),
[Float32(x).ir_value(loc=loc, ip=ip)],
"cvt.rni.s32.f32 $0, $1;",
"=r,f",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@cute.jit
def abs_scaled_to_e2m1_idx(a: cutlass.Float32) -> cutlass.Int32:
"""Map |scaled| value directly to E2M1 index with RNE.
@dsl_user_op
def f32_to_i32_rz(x: Float32, *, loc=None, ip=None) -> Int32:
"""Convert Float32 to Int32 with round-toward-zero (RZ).
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)
Wraps PTX: cvt.rzi.s32.f32 $0, $1;
Equivalent to CUDA __float2int_rz(). Used for floor-based extraction.
"""
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
return Int32(
llvm.inline_asm(
T.i32(),
[Float32(x).ir_value(loc=loc, ip=ip)],
"cvt.rzi.s32.f32 $0, $1;",
"=r,f",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def f32_to_i32_rmi(x: Float32, *, loc=None, ip=None) -> Int32:
"""Convert Float32 to Int32 with round-to-minus-infinity (RMI / floor).
Wraps PTX: cvt.rmi.s32.f32 $0, $1;
Equivalent to CUDA __float2int_rd() / floorf() → int.
Used for floor(log2()) extraction in FP8 encoding.
"""
return Int32(
llvm.inline_asm(
T.i32(),
[Float32(x).ir_value(loc=loc, ip=ip)],
"cvt.rmi.s32.f32 $0, $1;",
"=r,f",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
# ── 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 (returned as Int32).
Only handles positive values (NVFP4 scale factors are always positive).
Returns the uint8 bit pattern packed into an Int32.
Uses proper PTX inline asm for float→int conversions:
- cvt.rmi (floor) for exponent extraction
- cvt.rni (round-to-nearest-even) for mantissa quantization
"""
result = cutlass.Int32(0) # default: zero
@@ -117,7 +120,8 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
# 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))
# Compute floor(log2(clamped)) using frexp-like normalization.
# Double until >= 1, halve until < 2, tracking exponent shift.
norm = clamped
exp_floor = cutlass.Int32(0)
@@ -127,7 +131,7 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
norm = norm * cutlass.Float32(2.0)
exp_floor = exp_floor - cutlass.Int32(1)
# Halve until < 2 (at most 8 halvings needed, largest ≈ 240 < 256)
# Halve until < 2 (at most 8 halvings needed, largest ≈ 448 < 512)
for _ in cutlass.range(8, unroll=1):
if norm >= cutlass.Float32(2.0):
norm = norm * cutlass.Float32(0.5)
@@ -140,9 +144,9 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
if fp8_exp < cutlass.Int32(0):
fp8_exp = cutlass.Int32(0)
# Mantissa for normal: (norm - 1) * 8, round via threshold
# Mantissa for normal: (norm - 1) * 8, round via PTX cvt.rni
mantissa_f = (norm - cutlass.Float32(1.0)) * cutlass.Float32(8.0)
mantissa = round_rne_u0_8(mantissa_f)
mantissa = f32_to_i32_rni(mantissa_f)
# Mantissa overflow: rounded to 8 → increment exponent, reset mantissa
if mantissa >= cutlass.Int32(8):
@@ -171,7 +175,7 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
# 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)
sub_m = f32_to_i32_rni(sub_m_f)
if sub_m < cutlass.Int32(0):
sub_m = cutlass.Int32(0)
if sub_m > cutlass.Int32(7):
@@ -229,6 +233,21 @@ def fp8_e4m3_to_float32(bits: cutlass.Int32) -> cutlass.Float32:
return result
# ── E2M1 FP4 quantization ───────────────────────────────────────────
# E2M1 format (NVFP4 element format):
# Values: [0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
# Index: [0, 1, 2, 3, 4, 5, 6, 7]
# Encoded as 3-bit index + 1 sign bit = 4-bit nibble
# Half-step lookup: quantize as round(|x| * 2) → half_step, then map to E2M1 index.
# This is equivalent to the CUDA reference approach:
# int half_step = __float2int_rn(fabsf(scaled) * 2.0f);
# int idx = half_step_to_e2m1_idx[half_step];
#
# The half_step_to_e2m1_idx LUT (12 entries, half_step 0..11):
_HALF_STEP_TO_E2M1 = [0, 1, 2, 3, 4, 4, 5, 6, 6, 6, 7, 7]
@cute.jit
def quantize_e2m1_nibble(
val: cutlass.Float32,
@@ -238,6 +257,10 @@ def quantize_e2m1_nibble(
Returns uint4 nibble: bit 3 = sign, bits [2:0] = E2M1 index.
If scale ≈ 0, returns 0 (zero nibble).
Uses proper PTX cvt.rni for float→int conversion (round-to-nearest-even),
then LUT-based half_step → E2M1 index mapping. Matches the CUDA reference
exactly.
"""
nibble = cutlass.Int32(0)
@@ -246,7 +269,41 @@ def quantize_e2m1_nibble(
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)
# half_step = round(|scaled| * 2) via PTX cvt.rni
half_step_f = abs_scaled * cutlass.Float32(2.0)
half_step = f32_to_i32_rni(half_step_f)
# Clamp to LUT range [0, 11]
if half_step < cutlass.Int32(0):
half_step = cutlass.Int32(0)
if half_step > cutlass.Int32(11):
half_step = cutlass.Int32(11)
# LUT: half_step → E2M1 index
# [0,1,2,3,4,4,5,6,6,6,7,7]
idx = cutlass.Int32(0)
# We can't index an array in CuTeDSL, so implement the LUT with comparisons.
# This IS the LUT, just expressed as branch logic since CuTeDSL has no
# random-access shared arrays in @cute.jit scalar functions.
#
# The LUT is: half_step_to_e2m1[hs] for hs in 0..11
# 0→0, 1→1, 2→2, 3→3, 4→4, 5→4, 6→5, 7→6, 8→6, 9→6, 10→7, 11→7
if half_step >= cutlass.Int32(1):
idx = cutlass.Int32(1)
if half_step >= cutlass.Int32(2):
idx = cutlass.Int32(2)
if half_step >= cutlass.Int32(3):
idx = cutlass.Int32(3)
if half_step >= cutlass.Int32(4):
idx = cutlass.Int32(4)
# hs=5 also maps to 4 (NOT 5)
if half_step >= cutlass.Int32(6):
idx = cutlass.Int32(5)
if half_step >= cutlass.Int32(7):
idx = cutlass.Int32(6)
# hs=8,9 also map to 6 (NOT 7,8)
if half_step >= cutlass.Int32(10):
idx = cutlass.Int32(7)
if scaled < cutlass.Float32(0.0):
nibble = idx + cutlass.Int32(8)

View File

@@ -0,0 +1,106 @@
"""Test: inline PTX f32→i32 conversion and FP4 quantization in CuTeDSL."""
import torch
import cutlass
import cutlass.cute as cute
import cutlass.torch as cutlass_torch
from dsv4.kernels.gemm.fp4_quant import (
f32_to_i32_rni,
f32_to_i32_rz,
f32_to_i32_rmi,
quantize_e2m1_nibble,
fp8_e4m3_from_float32,
)
@cute.kernel
def ptx_convert_test_kernel(
input_f32: cute.Tensor, # (N,) Float32 inputs
output_rni: cute.Tensor, # (N,) Int32 RNE results
output_rz: cute.Tensor, # (N,) Int32 RZ results
output_rmi: cute.Tensor, # (N,) Int32 RMI results
):
tidx, _, _ = cute.arch.thread_idx()
if tidx < cutlass.Int32(cute.shape(input_f32)[0]):
x = cute.arch.load(input_f32.iterator, cutlass.Float32)
cute.arch.store(output_rni.iterator, f32_to_i32_rni(x))
cute.arch.store(output_rz.iterator, f32_to_i32_rz(x))
cute.arch.store(output_rmi.iterator, f32_to_i32_rmi(x))
@cute.kernel
def fp4_quant_test_kernel(
input_f32: cute.Tensor, # (N,) Float32 values
scale_f32: cute.Tensor, # (1,) Float32 scale
output_nibble: cute.Tensor, # (N,) Int32 nibbles
):
tidx, _, _ = cute.arch.thread_idx()
if tidx < cutlass.Int32(cute.shape(input_f32)[0]):
val = cute.arch.load(input_f32.iterator, cutlass.Float32)
scale = cute.arch.load(scale_f32.iterator, cutlass.Float32)
nibble = quantize_e2m1_nibble(val, scale)
cute.arch.store(output_nibble.iterator, nibble)
if __name__ == "__main__":
# Test PTX conversions
N = 12
test_vals = torch.tensor(
[0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, -1.7, 0.0, 2.0, 3.0],
dtype=torch.float32, device='cuda'
)
out_rni = torch.zeros(N, dtype=torch.int32, device='cuda')
out_rz = torch.zeros(N, dtype=torch.int32, device='cuda')
out_rmi = torch.zeros(N, dtype=torch.int32, device='cuda')
vc = cutlass_torch.from_dlpack(test_vals).mark_layout_dynamic(leading_dim=0)
rni_c = cutlass_torch.from_dlpack(out_rni).mark_layout_dynamic(leading_dim=0)
rz_c = cutlass_torch.from_dlpack(out_rz).mark_layout_dynamic(leading_dim=0)
rmi_c = cutlass_torch.from_dlpack(out_rmi).mark_layout_dynamic(leading_dim=0)
print("Compiling PTX conversion kernel...")
compiled = cute.compile(ptx_convert_test_kernel, vc, rni_c, rz_c, rmi_c)
print("Running...")
compiled(vc, rni_c, rz_c, rmi_c)
print("\nPTX conversion results:")
for i in range(N):
v = test_vals[i].item()
print(f" val={v:6.1f} rni={out_rni[i].item():3d} rz={out_rz[i].item():3d} rmi={out_rmi[i].item():3d}")
# Expected RNE: 0→0, 1.5→2, 2.5→2, 3.5→4, 4.5→4, 5.5→6, 6.5→6, 7.5→8, -1.7→-2, 0→0, 2→2, 3→3
expected_rni = [0, 2, 2, 4, 4, 6, 6, 8, -2, 0, 2, 3]
rni_ok = all(out_rni[i].item() == expected_rni[i] for i in range(N))
print(f"\nRNE correct: {rni_ok}")
# Test FP4 quantization
print("\n--- FP4 Quantization Test ---")
q_vals = torch.tensor([0.0, 0.25, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -1.0, -3.0],
dtype=torch.float32, device='cuda')
scale = torch.tensor([1.0], dtype=torch.float32, device='cuda')
q_out = torch.zeros(len(q_vals), dtype=torch.int32, device='cuda')
qvc = cutlass_torch.from_dlpack(q_vals).mark_layout_dynamic(leading_dim=0)
sc = cutlass_torch.from_dlpack(scale).mark_layout_dynamic(leading_dim=0)
qoc = cutlass_torch.from_dlpack(q_out).mark_layout_dynamic(leading_dim=0)
print("Compiling FP4 quant kernel...")
compiled_q = cute.compile(fp4_quant_test_kernel, qvc, sc, qoc)
print("Running...")
compiled_q(qvc, sc, qoc)
# E2M1 values: [0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] → indices [0..7]
# sign bit: bit 3
e2m1_vals = [0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
print("\nFP4 quantization results (scale=1.0):")
for i in range(len(q_vals)):
v = q_vals[i].item()
nib = q_out[i].item()
sign = (nib >> 3) & 1
idx = nib & 7
dequant = e2m1_vals[idx] if idx < 8 else -1
if sign:
dequant = -dequant
print(f" val={v:6.1f} nibble={nib:2d} sign={sign} idx={idx} dequant={dequant}")
print("\nDone!")