Add Triton grouped-GEMM for MoE experts on Ampere/Ada

[hwchen2017]

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]

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

[hwchen2017]

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]

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]

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

[hwchen2017]

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