Repository navigation
NVFP4EffMA 0.3.0
NVFP4EffMA 0.3.0
NVFP4 GEMM, GEMV and quantization for NVIDIA GPUs without native FP4 (sm_50 … sm_90).
Install
PyTorch: pick the wheel matching your Python (cpXY), PyTorch (torchA.B) and its CUDA (cuNNN):
pip install nvfp4effma-0.3.0+<cuNNNtorchA.B>-<cpXY>-<cpXY>-linux_x86_64.whlOther combinations: pip install ./binds/python --no-build-isolation from the source archive.
C / C++: unpack nvfp4effma-0.3.0-linux-x86_64-cuda12.tar.xz (tar -xJf) and link lib/libnvfp4effma.a (see its README).
Every artifact carries SASS for sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86, sm_87, sm_89, sm_90 (CUDA 12.0).
Assets
nvfp4effma-0.3.0-linux-x86_64-cuda12.tar.xz(8,5M)nvfp4effma-0.3.0-src.tar.gz(136K)nvfp4effma-0.3.0+cu126torch2.7-cp312-cp312-linux_x86_64.whl(57M)
What's new
- Every NVIDIA NVFP4 variant, read and produced as it is, bit for bit with Transformer Engine's arithmetic: ModelOpt checkpoints (1 × 16, decode factor), Transformer Engine's 1 × 16 and 16 × 16 (its default for weights) recipes, amax globals per tensor or per row (row-scaled activations), 4over6 (map-to-4 / map-to-6, MAE / MSE, E4M3 bound 448 or 256), UE5M3 block scales, rowwise and columnwise data, and the GEMM-swizzled scale layout of cuBLASLt — in GEMM, GEMV, quantization and dequantization, A and B in any combination. Not supported: the random Hadamard transform, stochastic rounding and 4over6's fast-math error mode.
- C API:
nvfp4a_tensorgainsflags(NVFP4A_SCALES_SWIZZLED,NVFP4A_SCALES_UE5M3,NVFP4A_GLOBAL_AMAX,NVFP4A_GLOBAL_PER_ROW) andscale_max;nvfp4a_quantize_ex(rowwise / columnwise, 1 × 16 / 16 × 16, 4over6, amax computed),nvfp4a_gemm_dual(forward and backward GEMMs from rowwise / columnwise halves),nvfp4a_swizzle_scales,nvfp4a_transpose(16 × 16 tensors). - PyTorch:
quantize(x, block=, four_over_six=, scale_max=, row_scaled=, columnwise=, swizzled=, scale_dtype=)with Transformer Engine's current scaling by default (the amax is computed and stored);NVFP4Tensorholds every layout,NVFP4Tensor.from_tewraps a Transformer Engine NVFP4 tensor without copying,.with_scale_layout(),.transpose_2d();lineartakes any formats on either side. - No cost for the formats in the kernels: block scale formats are compile-time kernel instances; the GEMM rows of 0.2.0 are unchanged or slightly faster (NVFP4 × NVFP4 9.42 ms, W4A8 9.33 ms at 4096³), the Transformer Engine formats within 1.5 % of plain NVFP4 in GEMM and −6.5 … +3 % in GEMV (swizzled x scales are faster). Creating them costs more — quantize once, not per call (
docs/performance.md). - Quantization: 16 × 16 tiles at 0.74 ms (BF16, 8192²), 4over6 at 1.9 ms with the division by 6·M computed exactly through its reciprocal (exhaustively checked), the amax pass vectorized.
- Benchmarks: a Transformer Engine block in every benchmark (each row against the plain format of the same run), a warm-up before the first row,
NVFP4A_BENCH_FILTERto run selected rows.
Changes in behaviour. Quantization now follows Transformer Engine's multiplications (x · inv, bmax · S_enc / 6) instead of 0.2.0's divisions, which rarely differ by one rounding step: some blocks get a different code than with 0.2.0. Python quantize(x) without a global_scale now stores an amax global (Transformer Engine's convention) instead of the decode factor amax / (6 · 448); values are the same. The C struct nvfp4a_tensor has two more trailing fields (source-compatible: initializers without them mean the ModelOpt format; not ABI-compatible).
The C tarball is now .tar.xz (8.6 MB; the format instances compress well, gzip would be 53 MB).
Tested and tuned by the author on a GTX 1080 (sm_61) only, plus the code path of GPUs without DP4A (SASS built for sm_60, run on the GTX 1080). Other architectures are compiled, not run. No stack or local memory in any kernel, except on Jetson TX1 / TX2 (sm_53 / sm_62, 32 K registers per block), where the heaviest kernels spill.
The wheel here is built for Python 3.12, PyTorch 2.7 and CUDA 12.6.