[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

YangFei1990and 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' ofhttps://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

YangFei1990and 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:mainAug 21, 2026
30 of 35 checks passed

mmarcinkiewicz mentioned this pull request Aug 22, 2026
Fix include after ptx move #3416
Merged
13 tasks

timmoon10 mentioned this pull request Aug 25, 2026
[PyTorch] Fix Test's Reference Implementation of Clamped Swiglu #3424
Merged
13 tasks

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-appsBot commented Aug 12, 2026•
edited
Loading

Uh oh!

Oleg-Goncharov left a comment

Uh oh!

phu0ngngAug 20, 2026•
edited
Loading

Uh oh!

YangFei1990Aug 20, 2026

Uh oh!

ptrendxAug 20, 2026

Uh oh!

vthumbe1503Aug 21, 2026

Uh oh!

YangFei1990Aug 21, 2026

Uh oh!

vthumbe1503Aug 21, 2026

Uh oh!

vthumbe1503Aug 21, 2026

Uh oh!

YangFei1990Aug 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-appsBot 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!

phu0ngngAug 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!

YangFei1990Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

ptrendxAug 20, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

YangFei1990Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

YangFei1990Aug 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-appsBot commented Aug 12, 2026•
edited
Loading

phu0ngngAug 20, 2026•
edited
Loading