[Oleg-Goncharov]

Description

Add a fused grouped MXFP8 requantization path for Blackwell GPUs. The new operation converts compact rowwise MXFP8 input into columnwise E4M3 MXFP8 without materializing an intermediate high-precision tensor in global memory.

This reduces memory traffic compared with separate grouped dequantization and quantization calls. The kernel uses a tiled TMA pipeline and packed data movement/conversion operations, and supports both compact and GEMM-swizzled output scaling factors. The quantization configuration selects either an FP32 intermediate or a faster BF16 intermediate.

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 the nvte_group_requantize_mxfp8 C API for fused grouped rowwise-to-columnwise MXFP8 requantization.
• Add an SM100+ TMA-based kernel supporting grouped tensors, FP8 input types, E4M3 output, and compact or GEMM-swizzled output scales.
• Add optimized BF16 fast-math and FP32 intermediate-precision paths.
• Optimize packed data access and conversion in the requantization kernel.
• Add parameterized C++ correctness tests comparing the fused operation against separate dequantization and quantization.

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]

Confidence Score: 4/5

[Medium impact] The change reaches prequantized MXFP8 inputs used by grouped linear and grouped MLP workflows.

Fix the shared descriptor race before merging so varying-width calls can run safely on separate CUDA streams.

Findings

1.  Concurrent calls overwrite shared descriptors ▶

Summary

Replaces the earlier fused path with grouped CUDA requantization and an optional CuTeDSL backend.

• Grouped MXFP8 tensors can be requantized directly into columnwise E4M3.
• Supported grouped shapes can use the CuTeDSL requantization kernel.
• Prequantized grouped tensors now pass their shape metadata into requantization.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A["nvte_group_requantize"] --> B["Check grouped input and output"]
  B --> C{"CuTeDSL supports this call?"}
  C -->|Yes| D["Run CuTeDSL kernel"]
  C -->|No| E{"Member widths vary?"}
  E -->|No| F["Run CUDA with call-local descriptors"]
  E -->|Yes| G["Update shared g_tensor_maps"]
  G --> H["Run CUDA using shared descriptors"]
  D --> I["Columnwise MXFP8 and optional BF16 output"]
  F --> I
  H --> I

Loading

_{Reviews (13) · Last reviewed commit: "Merge branch 'main' into pr_requantize_m..." · Reviewed by Greptile}

[ptrendx]

@Oleg-Goncharov This PR seems to have some unrelated changes (like the nccl ep submodule). Could you fix that?

[Oleg-Goncharov]

/te-ci

[ptrendx]

@Oleg-Goncharov Please sign every commit in this PR. See https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work

[phu0ngng]

Hi @Oleg-Goncharov, what is the current state of this PR? Could we prioritize merging it soon?

[ptrendx]

/te-ci

[nvMelissa]

Moved to TE2.22

[ptrendx]

/te-ci

[greptile-apps]

Concurrent calls overwrite shared descriptors

For groups with varying widths, two nvte_group_requantize calls on different CUDA streams overwrite the same g_tensor_maps entries. Each call updates those entries before launching its kernel, but stream ordering does not protect them from the other call. One kernel can therefore read or write the other call's buffers, producing wrong output or an illegal memory access.

Keep the descriptors in storage owned by each call until its kernel finishes.