diff --git a/dsv4/kernels/gemm/fp4_quant.py b/dsv4/kernels/gemm/fp4_quant.py index 76e55f5a..6d4199b0 100644 --- a/dsv4/kernels/gemm/fp4_quant.py +++ b/dsv4/kernels/gemm/fp4_quant.py @@ -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) diff --git a/tests/unit/test_ptx_convert.py b/tests/unit/test_ptx_convert.py new file mode 100644 index 00000000..f7b3c6bf --- /dev/null +++ b/tests/unit/test_ptx_convert.py @@ -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!")