[Common/PyTorch] Fused grouped MXFP8 requantization by YangFei1990 · Pull Request #3359 · NVIDIA/TransformerEngine

YangFei1990 · 2026-08-12T22:16:53Z

Description

Replace the group_dequantize -> group_quantize(columnwise) -> grouped_swizzle(rowwise scales) chain in group_requantize_inplace with a single kernel (NVTE_FUSED_GROUP_REQUANTIZE=0 restores the unfused path).

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

Replace the group_dequantize -> group_quantize(columnwise) -> grouped_swizzle(rowwise scales) chain in group_requantize_inplace with a single kernel (NVTE_FUSED_GROUP_REQUANTIZE=0 restores the unfused path). Co-authored-by: Oleg Goncharov <ogoncharov@nvidia.com> Signed-off-by: YangFei1990 <feiw@nvidia.com>

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

greptile-apps · 2026-08-12T22:22:19Z

Greptile Summary

The PR replaces grouped MXFP8 dequantize, columnwise requantize, and scale swizzle operations with a fused Blackwell CUDA kernel while retaining an environment-controlled fallback.

Adds a public common-core requantization API and PyTorch dispatch path.
Supports optional BF16 dequantized output, capacity tails, zero-sized groups, and E4M3/E5M2 rowwise inputs.
Adds parameterized C++ coverage comparing fused output with the existing unfused operations.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains in the eligible follow-up-review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/fused_group_requantize.cu	Implements the Blackwell fused grouped MXFP8 dequantization, requantization, scale emission, and optional BF16-output kernel.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Selects the fused operation for supported PyTorch grouped tensors and preserves the existing unfused fallback.
transformer_engine/common/include/transformer_engine/cast.h	Declares and documents the new common C API.
tests/cpp/operator/test_fused_group_requantize_mxfp8.cu	Adds reference comparisons across group layouts, input FP8 types, fast-math modes, optional dequantization, and capacity tails.

Sequence Diagram

sequenceDiagram
  participant PT as PyTorch GroupedTensor
  participant Bind as group_requantize_inplace
  participant Core as nvte_group_requantize
  participant GPU as Fused CUDA kernel
  PT->>Bind: Rowwise MXFP8 data, scales, offsets
  alt Fused path supported and enabled
    Bind->>Core: Input/output descriptors and offsets
    Core->>GPU: Launch grouped requantization
    GPU-->>Core: Columnwise E4M3 data and swizzled scales
    Core-->>Bind: Optional BF16 dequantized output
    Bind-->>PT: Replace scale and columnwise buffers
  else Fallback path
    Bind->>Bind: Group dequantize
    Bind->>Bind: Group columnwise quantize
    Bind->>Bind: Grouped scale swizzle
  end

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

Signed-off-by: YangFei1990 <feiw@nvidia.com>

…90/TransformerEngine into fused_group_requantize

Oleg-Goncharov

Hi @YangFei1990, I’ve done some further optimization work on this implementation, and the version in this GitLab branch is ~10% faster than the current implementation in this PR:
https://gitlab-master.nvidia.com/ogoncharov/transformerengine/-/tree/pr_requantize_mxfp8

The main improvements came from using data swizzling in the TMA descriptor and the redux.sync instruction, which reduces the number of reads/writes to the shared buffer holding intermediate results.

Could you please update the PR to incorporate the corresponding changes before merging?

phu0ngng · 2026-08-20T01:45:16Z

+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
+void nvte_fused_group_requantize_mxfp8(const NVTETensor input, NVTETensor output,
+                                       const NVTETensor tensor_offsets, NVTETensor dequantized,


In which case do callers need the dequantized?

Could we rename this function to nvte_group_requantize since it already contains "fusion" and is also more general, so that we can extend it in the future? Ofc we should make sure to check and only support MXFP8 for now.

dequantized is required to compute dbias when it is needed. I will update the func name.

ptrendx · 2026-08-20T16:19:52Z

+ *                                 path.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
+void nvte_fused_group_requantize_mxfp8(const NVTETensor input, NVTETensor output,


Could we leave out the "mxfp8" from the name? We can note in the documentation that currently it only supports mxfp8, but we should still do a general API name.

vthumbe1503 · 2026-08-21T00:57:49Z

+  // unfused chain's dedicated empty-input handling accepts and the kernel's pointer
+  // validation (correctly) rejects.
+  const bool use_fused_kernel =
+      transformer_engine::getenv<bool>("NVTE_FUSED_GROUP_REQUANTIZE", true) && need_columnwise &&


Why env variable? If we are confident of its perf, then we should always enable it.

I would like to keep it just for debugging purpose. The default path is fused op. Please let me know if you really want to get rid of it.

vthumbe1503 · 2026-08-21T01:04:03Z

+  const bool use_fused_kernel =
+      transformer_engine::getenv<bool>("NVTE_FUSED_GROUP_REQUANTIZE", true) && need_columnwise &&
+      has_usable_offsets && total_tokens > 0 && otype == DType::kBFloat16 &&
+      quantizer.attr("dtype").cast<DType>() == DType::kFloat8E4M3 &&


Doing attr access for the same attribute again and again should be avoided, since it can be pretty heavy considering cpu overheads.

Could we initialize them once for the fused and unfused cases in the start of the function and reuse them?

vthumbe1503 · 2026-08-21T01:05:44Z

             "Requantizing a grouped input requires dims that are multiples of 128, but got (",
             total_tokens, ", ", hidden_dim, ").");

+  // Fused path (default; NVTE_FUSED_GROUP_REQUANTIZE=0 recovers the unfused chain): one


Do we already have pytorch side unit test to test the fused kernel?

Yes all tests in test_mxfp8_group_quantize_graph_safe.py will run this path

Oleg-Goncharov

The kernel LGTM for now. A follow-up PR with performance optimizations will come later

Signed-off-by: YangFei1990 <feiw@nvidia.com>

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

YangFei1990 · 2026-08-21T19:11:35Z

/te-ci L1 pytorch

Based on the latest comment that the optimizations will be part of the subsequent PR.

YangFei1990 requested review from Oleg-Goncharov, ksivaman and ptrendx as code owners August 12, 2026 22:16

YangFei1990 and others added 2 commits August 12, 2026 15:17

Merge branch 'main' into fused_group_requantize

a084a52

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

609ca6f

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

YangFei1990 added 2 commits August 12, 2026 16:17

fix zero tensor dim

07214aa

Signed-off-by: YangFei1990 <feiw@nvidia.com>

Merge branch 'fused_group_requantize' of https://github.com/YangFei19…

a3f2f97

…90/TransformerEngine into fused_group_requantize

phu0ngng requested review from phu0ngng and vthumbe1503 August 12, 2026 23:37

Oleg-Goncharov previously requested changes Aug 17, 2026

View reviewed changes

phu0ngng reviewed Aug 20, 2026

View reviewed changes

phu0ngng requested a review from timmoon10 August 20, 2026 01:50

ptrendx reviewed Aug 20, 2026

View reviewed changes

vthumbe1503 reviewed Aug 21, 2026

View reviewed changes

buptzyb mentioned this pull request Aug 21, 2026

[ROADMAP][2026 Q3] Megatron Core MoE Roadmap NVIDIA/Megatron-LM#6757

Open

Oleg-Goncharov reviewed Aug 21, 2026

View reviewed changes

Oleg-Goncharov self-requested a review August 21, 2026 15:10

YangFei1990 and others added 3 commits August 21, 2026 12:03

rename func; read less attrs

b88bd27

Signed-off-by: YangFei1990 <feiw@nvidia.com>

Merge branch 'main' into fused_group_requantize

19ca1ad

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

3beabd2

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

ptrendx approved these changes Aug 21, 2026

View reviewed changes

ptrendx merged commit 8d9325f into NVIDIA:main Aug 21, 2026
30 of 35 checks passed

Navigation Menu

Uh oh!

[Common/PyTorch] Fused grouped MXFP8 requantization - #3359

[Common/PyTorch] Fused grouped MXFP8 requantization#3359
ptrendx merged 8 commits into
NVIDIA:mainfrom
YangFei1990:fused_group_requantize

YangFei1990 commented Aug 12, 2026

Uh oh!

greptile-apps Bot commented Aug 12, 2026 •

edited

Loading

Uh oh!

Oleg-Goncharov left a comment

Uh oh!

phu0ngng Aug 20, 2026 •

edited

Loading

Uh oh!

YangFei1990 Aug 20, 2026

Uh oh!

ptrendx Aug 20, 2026

Uh oh!

vthumbe1503 Aug 21, 2026

Uh oh!

YangFei1990 Aug 21, 2026

Uh oh!

vthumbe1503 Aug 21, 2026

Uh oh!

vthumbe1503 Aug 21, 2026

Uh oh!

YangFei1990 Aug 21, 2026

Uh oh!

Oleg-Goncharov left a comment

Uh oh!

YangFei1990 commented Aug 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

YangFei1990 commented Aug 12, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

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

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Oleg-Goncharov left a comment

Choose a reason for hiding this comment

Uh oh!

phu0ngng Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YangFei1990 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

ptrendx Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

YangFei1990 Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

YangFei1990 Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

Oleg-Goncharov left a comment

Choose a reason for hiding this comment

Uh oh!

YangFei1990 commented Aug 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

greptile-apps Bot commented Aug 12, 2026 •

edited

Loading

phu0ngng Aug 20, 2026 •

edited

Loading