Add MXFP8 support with cuBLASMp

[almogsegal]

Description

Please include a brief summary of the changes, relevant motivation and context.

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:

• Add MXFP8 support in comm_gemm.cpp (cuBLASMp path).
• Add cuBLASMp + MXFP8 tests in test_comm_gemm_overlap.py.

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

[almogsegal]

MXFP8 comparison of cuBLASMp vs UB on DGX-B200:

GPUs	Op	UB time	cuBLASMp time	Speedup	Faster
2	AG	0.0969 ms	0.0802 ms	1.21x	cuBLASMp
2	RS	0.0867 ms	0.0730 ms	1.19x	cuBLASMp
4	AG	0.1215 ms	0.1255 ms	1.03x	UB
4	RS	0.1357 ms	0.1049 ms	1.29x	cuBLASMp
8	AG	0.2380 ms	0.2146 ms	1.11x	cuBLASMp
8	RS	0.2191 ms	0.2141 ms	1.02x	cuBLASMp

[greptile-apps]

Greptile Summary

This PR adds MXFP8 (block scaling) support to the cuBLASMp-backed comm+GEMM path in comm_gemm.cpp and removes the corresponding skip conditions in the test suite.

• Extends cublasmp_gemm with MXFP8 canonicalization logic: non-transposed A uses columnwise data, transposed A uses rowwise data (mirroring the cublaslt_gemm.cu pattern), and correctly keeps the transpose flag unchanged (unlike FP8, where columnwise data is the transposed view).
• Adds CUBLASMP_MATMUL_MATRIX_SCALE_VEC32_UE8M0 scale-mode selection for MXFP8 behind a CUBLASMP_VERSION >= 801 compile-time guard, and gates with_gemm_swizzled_scales validation on the same version check.
• Refactors test parameterization to use a combined COMM_GEMM_QUANTIZATION_PARAMS list and enables the previously-skipped cublasmp-mxfp8 test combinations for all-gather and reduce-scatter overlap tests.

Confidence Score: 5/5

The change is safe to merge; the MXFP8 cuBLASMp path is well-guarded by both compile-time version checks and runtime tensor-validation assertions.

The MXFP8 canonicalization logic correctly mirrors the cublaslt path, the transpose flag handling is consistent with the MXFP8 'logical shape preserved in columnwise data' invariant, and the scale-mode selection is version-gated. No correctness defects were found in the changed code paths.

No files require special attention beyond the minor test-skip gap for cuBLASMp < 0.8.1 noted in the review comment.

Important Files Changed

Filename	Overview
transformer_engine/common/comm_gemm/comm_gemm.cpp	Extends cublasmp_gemm to handle MXFP8 scaling mode: adds input canonicalization (selects rowwise vs columnwise data without flipping the transpose flag), validates with_gemm_swizzled_scales, and sets the VEC32_UE8M0 scale mode for cuBLASMp >= 0.8.1. Logic is well-structured and correctly mirrors the cublaslt path.
tests/pytorch/distributed/test_comm_gemm_overlap.py	Removes the cuBLASMp+MXFP8 skip conditions and refactors test parameterization into a combined COMM_GEMM_QUANTIZATION_PARAMS list. Tests correctly check mxfp8_available and the cublasmp build flag, but there is no Python-level check to skip if cuBLASMp < 0.8.1.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["cublasmp_gemm(a, b, d, transa, transb, ...)"] --> B{"tensor_scaling || mxfp8?"}
    B -- "No" --> ERR1["NVTE_ERROR: Unsupported scaling mode"]
    B -- "Yes" --> C{"mxfp8?"}
    C -- "Yes" --> D{"CUBLASMP_VERSION >= 801?"}
    D -- "No" --> ERR2["NVTE_ERROR: Requires cuBLASMp 0.8.1+"]
    D -- "Yes" --> E["Check with_gemm_swizzled_scales for A and B"]
    C -- "No (tensor FP8)" --> F["Skip swizzled-scales check"]
    E --> G["canonicalize_input(a, transa, is_a=true)"]
    F --> G
    G --> H{"MXFP8 scaling?"}
    H -- "is_a=true, !transa" --> I["use_columnwise(current_trans) — same trans flag"]
    H -- "is_a=true, transa" --> J["use rowwise data — same trans flag"]
    H -- "is_a=false, transb" --> K["use_columnwise(current_trans) — same trans flag"]
    H -- "is_a=false, !transb" --> L["use rowwise data — same trans flag"]
    H -- "FP8" --> M["hopper TN-swap / Blackwell missing-rowwise\n→ possibly use_columnwise(!trans)"]
    I & J & K & L & M --> N["canonicalize_input(b, transb, is_a=false)"]
    N --> O["init_matrices_fn (AgGemm / GemmRs)"]
    O --> P{"is_fp8 or is_mxfp8?"}
    P -- "Yes" --> Q["get_input_scale_mode:\nmxfp8 → VEC32_UE8M0 / fp8 → SCALAR_FP32"]
    P -- "No" --> R["No scale attributes"]
    Q & R --> S["cublasMpMatmul_bufferSize → cublasMpMatmul"]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["cublasmp_gemm(a, b, d, transa, transb, ...)"] --> B{"tensor_scaling || mxfp8?"}
    B -- "No" --> ERR1["NVTE_ERROR: Unsupported scaling mode"]
    B -- "Yes" --> C{"mxfp8?"}
    C -- "Yes" --> D{"CUBLASMP_VERSION >= 801?"}
    D -- "No" --> ERR2["NVTE_ERROR: Requires cuBLASMp 0.8.1+"]
    D -- "Yes" --> E["Check with_gemm_swizzled_scales for A and B"]
    C -- "No (tensor FP8)" --> F["Skip swizzled-scales check"]
    E --> G["canonicalize_input(a, transa, is_a=true)"]
    F --> G
    G --> H{"MXFP8 scaling?"}
    H -- "is_a=true, !transa" --> I["use_columnwise(current_trans) — same trans flag"]
    H -- "is_a=true, transa" --> J["use rowwise data — same trans flag"]
    H -- "is_a=false, transb" --> K["use_columnwise(current_trans) — same trans flag"]
    H -- "is_a=false, !transb" --> L["use rowwise data — same trans flag"]
    H -- "FP8" --> M["hopper TN-swap / Blackwell missing-rowwise\n→ possibly use_columnwise(!trans)"]
    I & J & K & L & M --> N["canonicalize_input(b, transb, is_a=false)"]
    N --> O["init_matrices_fn (AgGemm / GemmRs)"]
    O --> P{"is_fp8 or is_mxfp8?"}
    P -- "Yes" --> Q["get_input_scale_mode:\nmxfp8 → VEC32_UE8M0 / fp8 → SCALAR_FP32"]
    P -- "No" --> R["No scale attributes"]
    Q & R --> S["cublasMpMatmul_bufferSize → cublasMpMatmul"]

Loading

_{Reviews (3): Last reviewed commit: "Enable cuBLASMp MXFP8 overlap tests" | Re-trigger Greptile}

[denera]

LGTM overall, but our build flow supports cuBLASMp 0.8.0 while the new MXFP8 path requires 0.8.1. This is not accounted for in the tests, so they're going to blanket fail on MXFP8 when cuBLASMp version is 0.8.0.

I'd recommend just bumping the minimum cuBLASMp version to 0.8.1 in the build flow (see here).

[denera]

/te-ci L1