Skip to content

NVFP4EffMA 0.3.0

Choose a tag to compare

@srose69 srose69 released this 20 Oct 14:22
· 9 commits to main since this release

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.whl

Other 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_tensor gains flags (NVFP4A_SCALES_SWIZZLED, NVFP4A_SCALES_UE5M3, NVFP4A_GLOBAL_AMAX, NVFP4A_GLOBAL_PER_ROW) and scale_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); NVFP4Tensor holds every layout, NVFP4Tensor.from_te wraps a Transformer Engine NVFP4 tensor without copying, .with_scale_layout(), .transpose_2d(); linear takes 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_FILTER to 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.