Repository navigation
AutoEP: fused Triton kernel for the clamped expert activations - #8652
pengdurice wants to merge 2 commits into
Conversation
f0db347 to
90c9660
Compare
Signed-off-by: pengdurice <pengduhit@gmail.com>
90c9660 to
b5febe5
Compare
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b5febe562a
ℹ️ 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".
| g = tl.minimum(gate, limit) | ||
| u = tl.minimum(tl.maximum(up, -limit), limit) |
There was a problem hiding this comment.
When a preceding grouped GEMM yields a NaN in gate or up, Triton's elementwise minimum/maximum use their default non-NaN-propagating mode, unlike the eager torch.clamp reference, so these lines replace the NaN with a finite clamp bound. The backward kernel repeats that clamp and its comparisons then zero the corresponding gradient, allowing a diverged AutoEP step to produce finite outputs and gradients; request NaN propagation (or use comparisons that preserve NaNs) in both kernels.
Useful? React with 👍 / 👎.
| n_elements = gate.numel() | ||
| if n_elements > 0: | ||
| grid = (triton.cdiv(n_elements, _BLOCK_SIZE), ) | ||
| _swiglu_clamped_fwd_kernel[grid](gate, |
There was a problem hiding this comment.
When gate and up reside on a CUDA device that is not the process's current device (for example, tensors on cuda:1 while the caller remains on cuda:0), this launch and the analogous backward launch run on Triton's current device/stream rather than the tensor device, so the kernel sees inaccessible pointers or launches on the wrong stream. The packed implementation already enters an input-device context around both launches; the separate implementation needs the same protection.
AGENTS.md reference: AGENTS.md:L40-L40
Useful? React with 👍 / 👎.
| return out | ||
|
|
||
| @staticmethod | ||
| def backward(ctx, grad_out: torch.Tensor): |
There was a problem hiding this comment.
When a caller requests create_graph=True from the separate swiglu_clamped path, this backward writes gradients into fresh empty_like buffers via Triton, so the returned first derivatives are detached and a subsequent derivative fails. The packed class explicitly detects and rejects this unsupported mode, but this new separate class neither rejects it nor provides a differentiable fallback; add the same guard or a differentiable backward here.
Useful? React with 👍 / 👎.
AutoEP: fused Triton kernel for the clamped expert activations
Branch:
peng-add-fuse-fn-clamped-form-v1, rebased onmasterat5f266e26d(2026-10-15). The expert activation registry itregisters into is merged (#8650), and so is the fused gate+up projection (#8733), which this PR now also covers.
Part of: #8672 (Q4 2026 roadmap), "Mixture-of-Expert (MoE) Training Scaling and Kernel Development": "AutoEP
Performance Optimization & Parallel Folding (AutoEP + AutoTP)" (PoC @tohtana @pengdurice). Follows the merged expert
activation registry (#8650) and MiniMax-M3 preset (#8651).
Summary
The clamped expert activations
swiglu_clamped(DeepSeek-V4) andswiglu_oai(GPT-OSS, MiniMax-M3) run as plainPyTorch expressions that keep about four
[rows, ffn]temporaries for backward. Plainswigluhas fused Tritonkernels that keep only
gateandup. This PR adds the same kernels for the clamped forms, for both expert layouts:gate_up_impl="separate"(two grouped GEMMs):swiglu_clamped(gate, up, ...).gate_up_impl="fused"(one grouped GEMM, Add an opt-in fused gate+up grouped GEMM for AutoEP experts #8733):swiglu_clamped_packed(gate_up, ...), which reads the two halvesof the packed GEMM output in place, as master's
swiglu_packeddoes.At the busiest rank of a MiniMax-M3 run that ran out of memory (1,045,081 rows x 3,072, bf16), one
swiglu_oaicall goes from 41.86 GiB extra peak memory and 77.8 ms to 17.94 GiB and 11.9 ms, the same as plain
swiglu.Duplicate check
Open PRs and issues searched on 2026-10-12 for "swiglu", "clamp", "expert activation" and "gpt-oss": no other work on
fused clamped activations. #8748 (another author, open) adds a different kernel,
swiglu_weighted(routing weightinside the activation, plain
swigluonly), to the same two files (swiglu_triton.pyand the registry inep_experts.py). The two do not overlap in function; whichever merges second needs a small rebase.Problem
Each clamped form is two clamps, a sigmoid, three products and an add in plain PyTorch, each a full
[rows, ffn]temporary, most of them saved for backward. On an expert-parallel rank this is where the memory goes, because
AutoEP has no capacity limit and real routers send some ranks several times the mean number of rows.
MiniMax-M3 at 65,536 tokens with expert parallel 32 on 4 nodes of 8 H200 ran out of memory there: the busiest rank
received 1,045,081 rows (3.99 times its own 262,144) and failed inside
apply_expert_activationon one 11.96 GiBtemporary, with 129.50 GiB already in use. (Setup: all 60 layers, real checkpoint,
minimax_m3preset, ZeRO-3 withthe optimizer state on the host, activation checkpointing; torch 2.8.0+cu126, transformers 5.15.1, DeepSpeed master
4a5856c25plus the then-open registry PR.) Over the 57 MoE layers, the busiest of 32 ranks received a median 4.20times the mean number of rows (maximum 7.84).
The change
deepspeed/ops/triton_ops/swiglu_triton.py_swiglu_clamped_fwd_kernel/_swiglu_clamped_bwd_kernel, modelled on master's swiglu kernels: int64 offsets, float32 math, one rounding on store,OAI: tl.constexprselects the form,PACKED: tl.constexprreadsgateandupfrom one[rows, 2 * INTER]tensor with master's packed indexing. The clamp gradient passes at the limit, astorch.clamp's does. Publicswiglu_clamped(gate, up, alpha, limit, oai)andswiglu_clamped_packed(gate_up, alpha, limit, oai); plain-PyTorch fallback without Triton or on the CPU. The packed backward raises forcreate_graph=True, as master's packed swiglu doesdeepspeed/moe/ep_experts.pyfused_fnandpacked_fused_fnofswiglu_clampedandswiglu_oai. Nothing else changes: the grouped-GEMM expert paths reach them throughapply_expert_activation(..., fused=True)andapply_packed_expert_activation; the for-loop reference path keepsfused=Falsedocs/_pages/config-json.mdtests/unit/v1/ops/triton_ops/test_swiglu_triton.py,tests/unit/v1/moe/test_autoep_expert_activation.pySize: 264 added and 2 removed lines under
deepspeed/(232 added inswiglu_triton.py, 32 inep_experts.py).Measured
One H200, torch 2.8.0+cu126, Triton 3.4.0, through
apply_expert_activation/apply_packed_expert_activation(
fused=True). 1,045,081 rows x 3,072 in bf16 (one[rows, ffn]tensor = 5.98 GiB), alpha 1.702, limit7.0, inputs at 6 x standard normal so most values are past a clamp. Extra peak =
max_memory_allocatedduringforward + backward minus the memory allocated before (inputs and upstream gradient). Kept = memory still allocated
after the forward, minus the output. Time = median of 5 forward + backward passes after 3 warm-ups.
1c862f7df: extra peak / kept / timeswiglu(existing kernels, for comparison)swigluswiglu_clampedswiglu_clampedswiglu_oaiswiglu_oaiPer MoE layer on that rank,
swiglu_oaikeeps 23.92 GiB less for backward and runs 6.5 times faster (separate) or9.0 times faster (packed). The packed clamped kernels are 1.2-1.3 ms slower than master's packed
swiglukernel,which does less arithmetic per element. Bench script:
expert_activation_bench.py(TODO(user): attach as a gist ordrop this line).
Tests
Command, from
tests/, on 1 node of 8 H200 (143,771 MiB each, driver 580.126.09), torch 2.8.0+cu126, Triton 3.4.0,transformers 5.15.1, pytest 9.1.1:
DEEPSPEED_RUN_H100_TESTS=1 python -m pytest unit/v1/ops/triton_ops/test_swiglu_triton.py \ unit/v1/moe/test_autoep_expert_activation.py -v \ -k "not (skewed_and_empty or accumulated_updates_and_checkpoint)"Result: 204 passed, 12 deselected, in 31 s. On master
5f266e26d(2026-10-15): 204 passed, 12 deselected, in 34 s on torch 2.8.0; the two empty-expert tests from #8733 on torch 2.9.1: 12 passed on this branch and 12 passed on master.The two deselected tests are master's (#8733). On torch 2.8.0+cu126 on H200 they do not finish, on master as well
as on this branch: they pass an empty expert group (
counts = [37, 0, 291, 72]and[16, 0, 32, 16]) totorch._grouped_mm, whose backward never returns for some shapes when a group is empty on that torch version (GPU at100% with no progress; known from MiniMax-M3 work, one-GPU reproduction available). AutoEP itself never passes an
empty group, because
generate_permute_indicesgives every local expert at least 8 rows. On torch 2.8.0 master'stest_matches_separate_with_skewed_and_empty_experts[swiglu]did not finish in 300 s; on torch 2.9.1+cu128 (Triton3.5.1) both tests pass, 12 passed on this branch and 12 passed on master.
New and changed tests:
test_clamped_forward_matches_reference(2 forms x 3 dtypes x 6 shapes)test_clamped_backward_matches_reference(2 x 3 x 5)test_clamped_alpha_and_limit_are_honouredtest_clamped_gradient_at_the_limit_is_inclusivegate == limitorup == -limit, wheretorch.clamp's passestest_clamped_empty_input,test_clamped_non_contiguous_input,test_clamped_shape_and_dtype_mismatch_rejectedtest_clamped_packed_matches_the_separate_kernel_exactly(2 forms x 2 dtypes x 4 shapes, including 3-D and empty)test_fused_kernel_matches_plain_pytorch(every form with afused_fn)apply_expert_activation(fused=True)differing from the plain form evaluated in float32 and rounded once (bf16, 256 x 3,072)test_packed_kernels_are_registered(wastest_only_swiglu_has_a_packed_kernel)test_grouped_gemm_paths_match_the_for_loop,test_packed_activation_matches_the_separate_halves(master's, now running the new kernels)pre-commit run --filespasses on every changed file.Compatibility
The change is an elementwise kernel behind AutoEP's existing activation registry; it does not depend on the ZeRO
stage, offload or parallel layout.
gate_up_impl"separate"and"fused";torch._grouped_mmand Triton grouped-GEMM paths (unit tests)torch.cuda.device(#ignore-cuda), as master's packed swiglu doesAI assistance
This PR was written with AI assistance (Claude).
🤖 Generated with Claude Code