biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

NVFP4-1.1: replace cute.math.fmin with cute.arch.fmin (correct API)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 03:55:54 +00:00
parent 20d5ddfa3d
commit e565ebce91
2 changed files with 18 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
NVFP4-1.1_APPROACH_PART_2.md
View File

@@ -192,4 +192,14 @@ Actually, I just realized I can look at the B200's CUTLASS Python installation t
Looking at the README:
> `cute.arch.fmax` is impure for the vectorizer. Use it inside plain `range()`, never inside `vectorize=True`.
So it's `cute.arch.fmax`, not `cute.math.fmax`! And for `fmin`, it's `cute.math.fmin`. Let me check the codebase for the exact APIs:So `cute.arch.fmax` exists, and `cute.math.fmin` exists. But `cute.math.fmax` does NOT exist. Let me fix `fp4_quant.py` to use `cute.arch.fmax` instead of `cute.math.fmax`:Good, all `cute.math.fmax` replaced with `cute.arch.fmax`. Now fix the test kernel too:
So it's `cute.arch.fmax`, not `cute.math.fmax`! And for `fmin`, it's `cute.math.fmin`. Let me check the codebase for the exact APIs:So `cute.arch.fmax` exists, and `cute.math.fmin` exists. But `cute.math.fmax` does NOT exist. Let me fix `fp4_quant.py` to use `cute.arch.fmax` instead of `cute.math.fmax`:Good, all `cute.math.fmax` replaced with `cute.arch.fmax`. Now fix the test kernel too:The error is clear: `if val > cutlass.Float32(0.0)` — CuTeDSL can't use dynamic boolean values in `if` statements. The `if` condition must be a compile-time constant or use `cutlass.const_expr()`. For runtime conditionals, CuTeDSL uses a different pattern.
Looking at the existing kernel code, runtime conditionals in CuTeDSL work because the `if` blocks are compiled as MLIR `scf.if` operations. But the condition must be a CuTeDSL boolean (not a Python bool from `>` comparison). The `>` operator returns a CuTeDSL Boolean, which can't be converted to a Python bool for the `if` statement.
The fix: don't use `if val > 0.0` in `fp8_e4m3_from_float32`. Instead, restructure the code to avoid runtime conditionals, or use CuTeDSL's conditional select pattern.
Looking at the SwiGLU epilogue: it uses `if is_gate:` where `is_gate` is computed as `(subtile_idx % cutlass.Int32(2)) == cutlass.Int32(0)`. This works because the `if` with a CuTeDSL comparison creates an `scf.if` in MLIR. But the error says the boolean can't be converted.
Wait — let me re-read the error. It says "Unable to convert dynamic Boolean value to bool at compile time." This happens when Python tries to evaluate the `if` condition as a Python bool, but the CuTeDSL `>` returns a dynamic MLIR value. In the `@cute.kernel` context, this should be handled by the DSL's `if` rewriter. But the call goes through `fp8_e4m3_from_float32` which is a plain Python function, not a `@cute.jit` function.
The fix: **decorate `fp8_e4m3_from_float32` with `@cute.jit`** so that CuTeDSL handles the `if` blocks correctly.Also add `@cute.jit` to `half_step_to_e2m1_idx`:Indentation error — the `@cute.jit` decorator placement is wrong. Let me fix:There's an extra space before `def`. Let me fix:There's still a leading space. Let me fix all of them:Good, all fixed. Let me commit and push:

14
dsv4/kernels/gemm/fp4_quant.py
View File

@@ -77,7 +77,7 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
if val > cutlass.Float32(0.0):
# Clamp to FP8 E4M3 max non-NaN value (exp=15, mant=6 = 448.0)
clamped = cute.math.fmin(val, cutlass.Float32(448.0))
clamped = cute.arch.fmin(val, cutlass.Float32(448.0))
# Normalize to [1, 2) range, tracking floor(log2(clamped))
norm = clamped
@@ -99,7 +99,7 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
# FP8 exponent = floor(log2(val)) + bias
fp8_exp = exp_floor + cutlass.Int32(FP8_E4M3_BIAS)
fp8_exp = cute.math.fmin(fp8_exp, cutlass.Int32(15))
fp8_exp = cute.arch.fmin(fp8_exp, cutlass.Int32(15))
fp8_exp = cute.arch.fmax(fp8_exp, cutlass.Int32(0))
# Mantissa for normal: (norm - 1) * 8, round
@@ -112,9 +112,9 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
mantissa = cutlass.Int32(0)
fp8_exp = fp8_exp + cutlass.Int32(1)
mantissa = cute.math.fmin(mantissa, cutlass.Int32(7))
mantissa = cute.arch.fmin(mantissa, cutlass.Int32(7))
mantissa = cute.arch.fmax(mantissa, cutlass.Int32(0))
fp8_exp = cute.math.fmin(fp8_exp, cutlass.Int32(15))
fp8_exp = cute.arch.fmin(fp8_exp, cutlass.Int32(15))
fp8_exp = cute.arch.fmax(fp8_exp, cutlass.Int32(0))
# NaN guard: FP8 E4M3 with exp=15 and mant=7 is NaN.
@@ -128,7 +128,7 @@ def fp8_e4m3_from_float32(val: cutlass.Float32) -> cutlass.Int32:
if fp8_exp < cutlass.Int32(1):
sub_m_f = clamped * cutlass.Float32(512.0)
sub_m = cutlass.Int32(sub_m_f) # round-to-nearest-even
sub_m = cute.math.fmin(sub_m, cutlass.Int32(7))
sub_m = cute.arch.fmin(sub_m, cutlass.Int32(7))
sub_m = cute.arch.fmax(sub_m, cutlass.Int32(1))
mantissa = sub_m
fp8_exp = cutlass.Int32(0)
@@ -198,11 +198,11 @@ def quantize_e2m1_nibble(
if scale > cutlass.Float32(1e-8):
scaled = val / scale
abs_scaled = cute.arch.fmax(scaled, cutlass.Float32(0.0) - scaled)
abs_scaled = cute.math.fmin(abs_scaled, cutlass.Float32(6.0))
abs_scaled = cute.arch.fmin(abs_scaled, cutlass.Float32(6.0))
# half_step = round(|scaled| * 2) — round-to-nearest-even (matches __float2int_rn)
hs = cutlass.Int32(abs_scaled * cutlass.Float32(2.0))
hs = cute.math.fmin(hs, cutlass.Int32(12))
hs = cute.arch.fmin(hs, cutlass.Int32(12))
hs = cute.arch.fmax(hs, cutlass.Int32(0))
idx = half_step_to_e2m1_idx(hs)