[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel by cael-ling · Pull Request #3315 · NVIDIA/TransformerEngine

cael-ling · 2026-08-04T14:36:12Z

Description

Add nvte_group_swiglu_quantize, which fuses the weighted SwiGLU activation with column-wise MXFP8 quantization over grouped (MoE) tensors:

saved FC1 output [T, 2F]  --(weighted SwiGLU)-->  [T, F]  --(colwise MXFP8)-->  FC2 wgrad input

weighted_swiglu(x) = ( silu(x[:, :F]) * x[:, F:] ) * prob[:, None]

Performance (GB200)

T=98304, F=2048/4096, 16 experts, compact scales:

variant	per-iter	vs fused
fused (this kernel)	191.2 us	1.00x
unfused-compiled (3 elementwise ops fused into 1 kernel)	277.6 us	1.45x
unfused-te-op (`ScaledSwiGLU` + `tex.group_quantize`)	747.4 us	3.91x

Reproduce:

python benchmarks/benchmark_group_scaled_swiglu.py --num-groups 16 --scales compact
# shape sweep: 16/64 experts x compact/swizzled scales x eager/graph timing
python benchmarks/benchmark_group_scaled_swiglu.py --hidden 4096

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

Add nvte_group_swiglu_quantize, which fuses the weighted SwiGLU activation with columnwise MXFP8 quantization over grouped (MoE) tensors. Signed-off-by: Cael Ling <caell@nvidia.com>

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

greptile-apps · 2026-08-04T14:40:26Z

Greptile Summary

The PR adds a grouped weighted-SwiGLU kernel fused with column-wise MXFP8 quantization and exposes it through the C and PyTorch APIs.

Adds the Blackwell CUDA kernel, dispatch path, and public activation API.
Adds PyTorch binding validation, grouped-output construction, and pybind registration.
Adds C++ and PyTorch coverage plus a fused-versus-unfused benchmark.
The shared metadata-validator adjustment leaves existing grouped operations executing non-current-device metadata on the wrong CUDA stream/device.

Confidence Score: 4/5

The PR is not yet safe to merge because the shared metadata-validation change allows existing grouped operations to launch against metadata and outputs on mismatched CUDA devices.

Existing grouped bindings such as group_quantize do not install a device guard, so metadata on a non-current GPU now passes validation before offset kernels use the ambient stream and grouped output allocation follows the ambient CUDA device.

Files Needing Attention: transformer_engine/pytorch/csrc/quantizer.cpp and transformer_engine/pytorch/csrc/extensions/cast.cpp

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh	Implements the grouped weighted-SwiGLU and column-wise MXFP8 CUDA kernel, including grouped layout and scale-layout validation.
transformer_engine/common/cast/dispatch/quantize.cuh	Adds dispatch and validation for the fused grouped activation-quantization operation.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds the guarded PyTorch binding for the fused operation, while existing grouped bindings remain unguarded.
transformer_engine/pytorch/csrc/quantizer.cpp	Relaxes grouped metadata device validation without ensuring all existing grouped callers establish the matching device and stream.
tests/cpp/operator/test_cast_mxfp8_grouped_scaled_swiglu.cu	Adds numerical, dtype, grouped-shape, empty-expert, compact-scale, and swizzled-scale coverage.
tests/pytorch/test_grouped_tensor.py	Adds binding-level output-shape, dtype, and non-contiguous operand rejection tests.

Sequence Diagram

sequenceDiagram
    participant Py as PyTorch caller
    participant Bind as group_scaled_swiglu binding
    participant Quant as MXFP8 grouped allocator
    participant Core as NVTE dispatch
    participant GPU as CUDA kernel
    Py->>Bind: input [T,2F], prob [T], metadata
    Bind->>Bind: validate layout and devices
    Bind->>Bind: guard input CUDA device
    Bind->>Quant: allocate grouped [T,F] output
    Quant-->>Bind: columnwise MXFP8 buffers
    Bind->>Core: nvte_group_scaled_swiglu
    Core->>GPU: weighted SwiGLU + columnwise quantize
    GPU-->>Py: grouped MXFP8 output

_{Reviews (5): Last reviewed commit: "[PyTorch] Add a benchmark for group_scal..." | Re-trigger Greptile}

…tize Signed-off-by: Cael Ling <caell@nvidia.com>

Signed-off-by: Cael Ling <caell@nvidia.com>

…ntize Signed-off-by: Cael Ling <caell@nvidia.com>

vthumbe1503

Thanks for the PR @cael-ling. Could you also please publish the perf numbers for this kernel?

vthumbe1503 · 2026-08-07T21:21:29Z

+ *  \param[in,out] output    Grouped output tensor [T, F] (columnwise MXFP8).
+ *  \param[in]     stream    CUDA stream used for the operation.
+ */
+void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob,


Can we rename this? Quantization details are already encoded in GroupedTensor. It ll be consistent with the other API naming convention

Suggested change

void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob,

void nvte_group_scaled_swiglu(const NVTEGroupedTensor input, const NVTETensor prob,

vthumbe1503 · 2026-08-07T21:24:16Z

                          std::optional<at::Tensor> tensor_offsets,
                          std::optional<at::Tensor> noop_flag);

+py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob,


Similar comment.

Suggested change

py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob,

py::object group_scaled_swiglu(const at::Tensor &input_2f, const at::Tensor &prob,

lines up with TE's existing ScaledSwiGLU op, so nvte_group_swiglu_quantize becomes nvte_group_scaled_swiglu, and the kernel header, dispatch, bindings and tests follow. Also makes the kernel actually faster than the unfused path it replaces. Per-token scales are staged through shared memory once per chunk, the output buffer is single-buffered to raise occupancy, and SiLU uses an approximate exp and divide. The activation is quantized to MXFP8 immediately, so the approximation stays far below one FP8 ULP. Signed-off-by: Cael Ling <caell@nvidia.com>

Times the fused kernel against the unfused path it replaces, with the activation half computed three ways: eager PyTorch, torch.compile, and TE's existing ScaledSwiGLU op. Signed-off-by: Cael Ling <caell@nvidia.com>

cael-ling · 2026-08-08T13:40:46Z

Hi @vthumbe1503, thanks for your review, rename work has been performed, benchmark results have been added in the description, please take a look, thanks!

[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel

5ae2ae8

Add nvte_group_swiglu_quantize, which fuses the weighted SwiGLU activation with columnwise MXFP8 quantization over grouped (MoE) tensors. Signed-off-by: Cael Ling <caell@nvidia.com>

cael-ling requested review from Oleg-Goncharov, ksivaman and ptrendx as code owners August 4, 2026 14:36

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

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

2979881

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

greptile-apps Bot reviewed Aug 4, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/csrc/extensions/cast.cpp

Comment thread transformer_engine/pytorch/csrc/extensions/cast.cpp Outdated

[PyTorch] Require contiguous, colocated operands in group_swiglu_quan…

e092890

…tize Signed-off-by: Cael Ling <caell@nvidia.com>

greptile-apps Bot reviewed Aug 5, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/csrc/extensions/cast.cpp

[PyTorch] Require grouped metadata tensors on the current CUDA device

5bfd791

Signed-off-by: Cael Ling <caell@nvidia.com>

greptile-apps Bot reviewed Aug 5, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/csrc/quantizer.cpp Outdated

[PyTorch] Scope the grouped metadata device check to group_swiglu_qua…

8572e3a

…ntize Signed-off-by: Cael Ling <caell@nvidia.com>

vthumbe1503 reviewed Aug 7, 2026

View reviewed changes

cael-ling added 2 commits August 8, 2026 06:02

[PyTorch] Add a benchmark for group_scaled_swiglu

0b94fcb

Times the fused kernel against the unfused path it replaces, with the activation half computed three ways: eager PyTorch, torch.compile, and TE's existing ScaledSwiGLU op. Signed-off-by: Cael Ling <caell@nvidia.com>

Navigation Menu

Uh oh!

[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel - #3315

[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel#3315
cael-ling wants to merge 7 commits into
NVIDIA:mainfrom
cael-ling:feature/mxfp8-group-swiglu-recompute

cael-ling commented Aug 4, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Aug 4, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vthumbe1503 left a comment

Uh oh!

vthumbe1503 Aug 7, 2026

Uh oh!

cael-ling Aug 8, 2026

Uh oh!

vthumbe1503 Aug 7, 2026

Uh oh!

cael-ling Aug 8, 2026

Uh oh!

cael-ling commented Aug 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob,
	void nvte_group_scaled_swiglu(const NVTEGroupedTensor input, const NVTETensor prob,

	py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob,
	py::object group_scaled_swiglu(const at::Tensor &input_2f, const at::Tensor &prob,

Uh oh!

Conversation

cael-ling commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Performance (GB200)

Type of change

Changes

Checklist:

Uh oh!

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

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vthumbe1503 left a comment

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

cael-ling Aug 8, 2026

Choose a reason for hiding this comment

Uh oh!

vthumbe1503 Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

cael-ling Aug 8, 2026

Choose a reason for hiding this comment

Uh oh!

cael-ling commented Aug 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cael-ling commented Aug 4, 2026 •

edited

Loading

greptile-apps Bot commented Aug 4, 2026 •

edited

Loading