Add Triton grouped-GEMM for MoE experts on Ampere/Ada by hwchen2017 · Pull Request #8180 · deepspeedai/DeepSpeed

hwchen2017 · 2026-07-26T18:32:56Z

Motivation

AutoEP expert computation uses torch._grouped_mm , which only has a fused grouped-GEMM kernel on Hopper (sm90) and newer. On Ampere/Ada GPUs (sm80/sm86 — A100/A6000/A40/RTX 30xx) it silently falls back to a per-group Python loop ( _grouped_mm_fallback ), issuing one at::mm per expert plus a device→host sync on offs . This is slow for MoE, which invokes grouped GEMM 3× per layer in the forward and again in the backward.

What this PR does

Adds deepspeed/moe/group_gemm_triton.py : an autograd-aware Triton grouped GEMM that is a drop-in for torch._grouped_mm (2D×3D + offs ), and wires it into GroupedExperts so it is auto-selected on sm < 9.0 (native path is kept on sm90+).

Key points:

• Single fused kernel per grouped GEMM, no per-group launches and no device→host sync (per-tile → group mapping is computed on device from offs ).
• Support forward + backward for both operands, supporting fp16/bf16/fp32.
• trans_b option: expert weights are passed in their native [E, N, K] layout and transposed via strides inside the autograd Function, keeping .transpose off the tape. The weight gradient is produced directly in the native layout, avoiding a contiguous-materialization copy in backward.
• Device-side group metadata via a small Triton kernel ( _group_meta_kernel ).
• Auto-selection is overridable via DS_DISABLE_TRITON_GROUPED_MM=1 .
• Auto-tunning won't be called twice for same experts shape.

Performance
Configuration: dim=2048, hidden=768, avg_tokens/expert=64,
dtype=bf16, distribution: unbalanced

Experts	Tokens	Method	Forward (ms)	Fwd + Bwd (ms)	Backward (ms)
16	760	Triton	0.107	0.589	0.482
16	760	Torch	0.256	1.093	0.837
16	760	For-loop	0.776	6.337	5.561
32	1,848	Triton	0.213	0.622	0.410
32	1,848	Torch	0.755	2.222	1.466
32	1,848	For-loop	1.603	22.732	21.129

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9c5e0335bc

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

delock · 2026-07-28T08:52:59Z

+        #   BLOCK_M (M/token tile): 32 / 64 / 128
+        #   BLOCK_N (N output tile): 128 / 256
+        #   BLOCK_K (K reduction):   32 / 64
+        return [


Is it possible to abstract _gmm_configs into abstract accelerator interfaces to let different accelerator return their preferred configs?

I'd prefer to keep _gmm_configs local to group_gemm_triton.py for now, for a few reasons:

• These configs are the autotune search space for this specific Triton kernel, not a final device-tuned choice — triton.autotune already picks the best one per problem shape at runtime. So they're an implementation detail of the kernel rather than a device capability.
• The candidates for different kinds of device are likely same.
• Keeping it here also keeps all the Triton-specific logic (kernels + their autotune configs) in one place.

delock · 2026-07-28T10:52:06Z

Hi @hwchen2017 thanks for your PR. I have left my comments in the code. One comment is, this PR is intended for Ampere/Ada, but triton can be used for other accelertors as well. Current code contains many CUDA specific checks, these code should be generized. The suggested solution is inlined. @tjruwase @sfc-gh-truwase for accelerator abstraction topic.

sfc-gh-truwase · 2026-07-28T15:43:00Z

@hwchen2017 thanks for the PR. Please address @delock's feedback. Also, can you please include some perf numbers?

hwchen2017 · 2026-07-30T00:28:31Z

Hi @sfc-gh-truwase , I add some perf numbers and benchmark code.

Signed-off-by: Hongwei <hongweichen@microsoft.com>

sfc-gh-truwase · 2026-07-30T17:37:11Z

+        # grouped-GEMM kernel is preferred there when Triton is available.
+        # Set DS_DISABLE_TRITON_GROUPED_MM=1 to force the native path.
+        import os
+        if os.getenv("DS_DISABLE_TRITON_GROUPED_MM", "0") == "1":


Runtime env vars can be brittle in multi-node settings, and so generally avoided. If this is unavoidable then see https://www.deepspeed.ai/getting-started/#multi-node-environment-variables for proper usage.

@hwchen2017 any thoughts on this?

I will replace it with a variable in config.

Feel free to create an accelerator section in the config.

Please also add some documentation. You can create an accelerator section here
https://deepspeed.readthedocs.io/en/latest/index.html

Code has been updated to replace env var with a variable in config.

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

This reverts commit 51aa775. Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

tohtana

Hi @hwchen2017,

Thank you for submitting this amazing PR! This fills an important gap in AutoEP.
I left a few comments. I don't think any of them are blockers, but please consider addressing them.

tohtana · 2026-08-01T05:54:37Z

+                mask=mask_k[:, None] & mask_n[None, :],
+                other=0.0,
+            )
+            acc += tl.dot(a_tile, b_tile)


Lines 32-34 say "Only float16, bfloat16 and float32 inputs are supported, and both operands must share the same dtype. Matmuls accumulate in fp32 and the result
is cast back to the input dtype, matching torch._grouped_mm numerics." but the default precision of tl.dot is TF32.
I think this GEMM is not precision sensitive. So we could use TF32 but we also should clarify it. If we pass ieee to input_precision, FP32 inputs will be computed in FP32. It would also be good to add a config item to control it.

• The FP32 path is an edge case since half precision is the default precision in training now.
• TF32 is the de-facto standard for FP32 matmuls on Ampere+; torch.matmul /cuBLAS and torch._grouped_mm also use TF32 by default. Forcing ieee would actually diverge from stock PyTorch numerics and give up most of the tensor-core throughput. Fully disabling tensor cores for FP32 is rare these days.
• Adding an input_precision config knob is extra surface area for a path almost nobody exercises; we can add it later if a real use case shows up.

Okay, how about noting it with "Only float16, bfloat16 and float32 inputs are supported"?

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

hwchen2017 · 2026-08-03T06:00:50Z

+                            (``offs[-1] == M``). Empty groups are allowed.
+    out   : ``[M, N]``      same floating dtype as inputs.
+
+Only ``float16``, ``bfloat16`` and ``float32`` inputs are supported, and both


@tohtana It's mentioned here.

tohtana

Thank you for the update, @hwchen2017! Looks good to me.

hwchen2017 requested review from delock and sfc-gh-truwase July 26, 2026 18:32

hwchen2017 requested review from loadams, tjruwase and tohtana as code owners July 26, 2026 18:32

chatgpt-codex-connector Bot reviewed Jul 26, 2026

View reviewed changes

Comment thread deepspeed/moe/ep_experts.py Outdated

Comment thread deepspeed/moe/ep_experts.py Outdated

Comment thread deepspeed/moe/group_gemm_triton.py

delock reviewed Jul 28, 2026

View reviewed changes

Comment thread deepspeed/moe/ep_experts.py Outdated

delock reviewed Jul 28, 2026

View reviewed changes

Comment thread deepspeed/moe/ep_experts.py Outdated

delock reviewed Jul 28, 2026

View reviewed changes

Comment thread deepspeed/moe/ep_experts.py

delock reviewed Jul 28, 2026

View reviewed changes

hwchen2017 added 8 commits July 29, 2026 17:31

A good working version for group gemm

123c030

Signed-off-by: Hongwei <hongweichen@microsoft.com>

Refactor code

e7b962d

Signed-off-by: Hongwei <hongweichen@microsoft.com>

change autotune config

a815d4c

Signed-off-by: Hongwei <hongweichen@microsoft.com>

Update comments

0699071

Signed-off-by: Hongwei <hongweichen@microsoft.com>

Add fused triton kernel for some elementwise kernel

ab3f239

Signed-off-by: Hongwei <hongweichen@microsoft.com>

Add benchmark code

f09c1da

Signed-off-by: Hongwei <hongweichen@microsoft.com>

Fix format error, and add doc

584236a

Signed-off-by: Hongwei <hongweichen@microsoft.com>

Fix format error, address reviewer comments

ca2090b

Signed-off-by: Hongwei <hongweichen@microsoft.com>

hwchen2017 force-pushed the hongwei/group_gemm branch from 264b399 to ca2090b Compare July 30, 2026 00:33

sfc-gh-truwase reviewed Jul 30, 2026

View reviewed changes

Comment thread tests/unit/v1/moe/test_group_gemm_triton.py

Fix review comments

ed93495

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

delock approved these changes Jul 31, 2026

View reviewed changes

Replace disable_triton_group_mm env variabel with config

58ddff6

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

Copilot started work on behalf of hwchen2017 July 31, 2026 20:20 View session

Skip CUDA forward tests when transformer op is unavailable

51aa775

Copilot finished work on behalf of hwchen2017 July 31, 2026 20:22

Revert "Skip CUDA forward tests when transformer op is unavailable"

1cca543

This reverts commit 51aa775. Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

hwchen2017 force-pushed the hongwei/group_gemm branch from fd3a792 to 1cca543 Compare July 31, 2026 21:37

tohtana requested changes Aug 1, 2026

View reviewed changes

Address review comments

5ac1ad3

Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>

hwchen2017 commented Aug 3, 2026

View reviewed changes

tohtana approved these changes Aug 3, 2026

View reviewed changes

hwchen2017 added this pull request to the merge queue Aug 3, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Aug 3, 2026

hwchen2017 added this pull request to the merge queue Aug 3, 2026

Merged via the queue into master with commit df84f6d Aug 3, 2026
14 of 15 checks passed

hwchen2017 deleted the hongwei/group_gemm branch August 3, 2026 07:45

Uh oh!

Conversation

hwchen2017 commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

delock commented Jul 28, 2026

Uh oh!

sfc-gh-truwase commented Jul 28, 2026

Uh oh!

hwchen2017 commented Jul 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tohtana left a comment • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tohtana left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

hwchen2017 commented Jul 26, 2026 •

edited

Loading

tohtana left a comment •

edited

Loading