Add Triton grouped-GEMM for MoE experts on Ampere/Ada (#8180)
**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 |
---------
Signed-off-by: Hongwei <hongweichen@microsoft.com>
Signed-off-by: Hongwei Chen <hongweichen@microsoft.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
