cuBLASMp NVFP4 support by almogsegal · Pull Request #3634 · NVIDIA/TransformerEngine

almogsegal · 2026-10-06T13:37:53Z

Description

This PR adds NVFP4 support using the cuBLASMp backend

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

greptile-apps · 2026-10-06T13:44:32Z

Confidence Score: 3/5

[High impact] Adds NVFP4 quantization support to distributed training.

The PR still needs the two previously reported test-version compatibility issues addressed before merging.

Findings

Runtime version misses build capability. ▶
Version probe can crash collection. ▶
Aligned AllReduce loses overlap. ▶

Summary

Adds NVFP4 inputs to the cuBLASMp backend and extends distributed overlap tests.

cuBLASMp GEMMs now accept NVFP4 inputs.
NVFP4 keeps both data layouts for backward GEMMs.

_{Reviews (7) · Last reviewed commit: "Merge branch 'main' into cublasmp-nvfp4" · Reviewed by Greptile}

greptile-apps · 2026-10-06T13:44:36Z

+  if (nvfp4 && init_matrices_fn != AgGemmInitMatrices &&
+      (algo_attr == CUBLASMP_MATMUL_ALGO_TYPE_SPLIT_P2P ||
+       algo_attr == CUBLASMP_MATMUL_ALGO_TYPE_SPLIT_MULTICAST) &&
+      (n % ctx->nranks != 0 || (n / ctx->nranks) % 128 != 0)) {


Aligned AllReduce loses overlap. The AllReduce output descriptor gives each rank n columns, but this check tests n / nranks for 128-column alignment. For example, with n=128 and eight ranks, it replaces an aligned split algorithm with the default algorithm. Results remain correct, but the operation loses communication overlap and emits a misleading warning.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

greptile-apps · 2026-10-06T20:01:53Z

+
+
+# NVFP4 with cuBLASMp applies the global scales through a device-side alpha (cuBLASMp 0.11.0+)
+cublasmp_nvfp4_available = _loaded_cublasmp_version() >= 1100


Runtime version misses build capability. If TE was built with cuBLASMp headers older than 0.11.0 but loads a 0.11.0-or-newer library, this check allows the NVFP4 tests to run. The compiled GEMM path still rejects NVFP4, so the tests fail instead of skipping.

greptile-apps · 2026-10-06T20:01:53Z

+        return 0
+    try:
+        version = ctypes.c_int(0)
+        ctypes.CDLL("libcublasmp.so.0").cublasMpGetVersion(ctypes.byref(version))


Version probe can crash collection. If an older supported cuBLASMp library uses the version getter that takes a handle and a version pointer, this one-argument call uses the wrong signature. Because it runs when the test module is imported, it can crash collection even for tests that do not use NVFP4.

timmoon10 · 2026-10-12T17:40:41Z

Pipeline 73310509

NVFP4 GEMMs apply the tensors' global scales through alpha, which TE computes on the GPU from the amaxes (nvte_nvfp4_compute_per_tensor_scale). Without a host sync this needs device pointer mode for alpha and beta, which cuBLASMp added in 0.11.0 (CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_POINTER_MODE), so the cuBLASMp backend rejected NVFP4 until now. The cuBLASMp GEMM now mirrors the cuBLASLt NVFP4 path: it uses the TN layout (the transposed column-wise data and its amax replace the row-wise data where needed), passes the FP4 data with VEC16_UE4M3 block scales, and computes alpha on the stream into a small per-context device buffer that also holds the betas 0 and 1. Row-scaled NVFP4, UE5M3 block scales and quantized outputs are rejected with an error. NVFP4 needs cuBLASMp 0.11.0 or newer both at build time and at run time; the run-time check uses cublasMpGetVersion(), so an older loaded library reports an error instead of misreading the device pointers. When n / tp_size is not a multiple of 128, NVFP4 GEMM + ReduceScatter and GEMM + AllReduce use cuBLASMp's default algorithm and warn once, since the explicitly selected split algorithms need a separate B scale tensor for each output-column chunk. With cuBLASMp, the wgrad GEMM after the dgrad AllGather + GEMM all-gathers grad_output again. An NVFP4 grad output is now quantized once, with both usages, before the dgrad GEMM, which uses the row-wise data, and the wgrad all-gather takes the column-wise data, as the non-overlapped all-gather does. Quantizing it a second time would also draw different stochastic-rounding numbers than the non-overlapped path. Signed-off-by: Almog Segal <asegal@nvidia.com>

Add NVFP4 to the cuBLASMp AllGather + GEMM and GEMM + ReduceScatter cases and to the layer overlap tests. Userbuffers does not support NVFP4, so those layer cases are skipped. The layer runner creates BF16 inputs for NVFP4, since the NVFP4 random Hadamard transform supports only BF16 input. The GEMM overlap check compares NVFP4 with the GEMM of its dequantized operands, since NVFP4 quantization error would dominate a comparison with the unquantized operands; this also checks the global scales applied through alpha. The all-gathered input reduces its amax across ranks, as TE's sequence-parallel layers do, while the reduce-scatter inputs keep per-rank global scales. Signed-off-by: Almog Segal <asegal@nvidia.com>

NVFP4 with cuBLASMp needs cuBLASMp 0.11.0 or newer, and the GEMM reports an error with an older library. The GEMM and layer overlap tests now read the version of the cuBLASMp library that TE loaded (cublasMpGetVersion()) and skip their NVFP4 cuBLASMp cases when it is older, as they skip on GPUs without NVFP4 support. Signed-off-by: Almog Segal <asegal@nvidia.com>

…ands The GEMM overlap test compared FP8 and MXFP8 outputs with a GEMM of the unquantized operands, so quantization error set its tolerances. For every quantized format, compute the reference from the dequantized test operands instead, as was already done for NVFP4. Quantization then cancels out, the test checks only the distributed GEMM, and it uses the unquantized tolerances unless the GEMM output itself is quantized. Signed-off-by: Almog Segal <asegal@nvidia.com>

Apply consistent quant-dequant pattern for reference tensors. Remove dead FP8 logic. Signed-off-by: Tim Moon <tmoon@nvidia.com>

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 · 2026-10-14T22:45:57Z

Pipeline 73930633

for more information, see https://pre-commit.ci

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 · 2026-10-19T16:53:14Z

Pipeline 75062152

almogsegal requested review from ksivaman and ptrendx as code owners October 6, 2026 13:37

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Oct 6, 2026

almogsegal force-pushed the cublasmp-nvfp4 branch from 1f3d589 to 97a4cd4 Compare October 6, 2026 13:38

greptile-apps Bot reviewed Oct 6, 2026

View reviewed changes

timmoon10 requested a review from vcherepanov-nv October 12, 2026 17:25

timmoon10 reviewed Oct 13, 2026

View reviewed changes

Comment thread tests/pytorch/distributed/run_gemm_with_overlap.py Outdated

almogsegal added 4 commits October 14, 2026 10:08

almogsegal force-pushed the cublasmp-nvfp4 branch from 11ee632 to 9949b9f Compare October 14, 2026 08:59

timmoon10 reviewed Oct 14, 2026

View reviewed changes

Comment thread tests/pytorch/distributed/run_gemm_with_overlap.py Outdated

Comment thread tests/pytorch/distributed/run_gemm_with_overlap.py Outdated

Comment thread tests/pytorch/distributed/run_gemm_with_overlap.py Outdated

timmoon10 reviewed Oct 14, 2026

View reviewed changes

Comment thread tests/pytorch/distributed/run_layer_with_overlap.py Outdated

timmoon10 reviewed Oct 14, 2026

View reviewed changes

Comment thread tests/pytorch/distributed/run_gemm_with_overlap.py

timmoon10 and others added 2 commits October 14, 2026 15:43

Refactor UB/cuBLASMp GEMM tests

9a845b8

Apply consistent quant-dequant pattern for reference tensors. Remove dead FP8 logic. Signed-off-by: Tim Moon <tmoon@nvidia.com>

Use BF16 input data in UB/cuBLASMp linear tests

f040759

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

pre-commit-ci Bot and others added 2 commits October 14, 2026 22:49

[pre-commit.ci] auto fixes from pre-commit.com hooks

3b61a35

for more information, see https://pre-commit.ci

Merge branch 'main' into cublasmp-nvfp4

453fdfb

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

Navigation Menu

Repository navigation

cuBLASMp NVFP4 support - #3634

cuBLASMp NVFP4 support#3634
almogsegal wants to merge 8 commits into
NVIDIA:mainfrom
almogsegal:cublasmp-nvfp4

almogsegal commented Oct 6, 2026

Uh oh!

greptile-apps Bot commented Oct 6, 2026 •

edited

Loading

Uh oh!

Uh oh!

greptile-apps Bot Oct 6, 2026

Uh oh!

greptile-apps Bot Oct 6, 2026

Uh oh!

greptile-apps Bot Oct 6, 2026

Uh oh!

timmoon10 commented Oct 12, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Oct 14, 2026

Uh oh!

timmoon10 commented Oct 19, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		# NVFP4 with cuBLASMp applies the global scales through a device-side alpha (cuBLASMp 0.11.0+)
		cublasmp_nvfp4_available = _loaded_cublasmp_version() >= 1100

Conversation

almogsegal commented Oct 6, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 3/5

Findings

Summary

Uh oh!

Uh oh!

greptile-apps Bot Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

timmoon10 commented Oct 12, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

timmoon10 commented Oct 14, 2026

Uh oh!

timmoon10 commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

greptile-apps Bot commented Oct 6, 2026 •

edited

Loading

timmoon10 commented Oct 19, 2026 •

edited

Loading