Add fused triton kernel for swiglu - #8244
Conversation
Signed-off-by: Hongwei Chen <[email protected]>
Signed-off-by: Hongwei Chen <[email protected]>
Signed-off-by: Hongwei Chen <[email protected]>
Signed-off-by: Hongwei Chen <[email protected]>
![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: 570831af74
ℹ️ 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".
| w3.to(cast_dtype).transpose(-2, -1), | ||
| offs=offsets, | ||
| ) | ||
| h = swiglu(gate, up) |
There was a problem hiding this comment.
When GroupedExperts selects the native torch._grouped_mm path, including when callers pass disable_triton_grouped_mm=True to avoid the Triton grouped-GEMM path, this line still dispatches the new Triton SwiGLU kernel whenever the triton package is importable. In environments where Triton imports but kernels are unsupported or intentionally disabled, the supposed native path can now fail or JIT Triton anyway; keep this path on the eager F.silu(gate) * up expression or plumb an explicit opt-in for the fused activation.
Useful? React with 👍 / 👎.
| _TRITON_AVAILABLE = False | ||
|
|
||
| __all__ = ["group_gemm_triton", "is_available"] | ||
| from deepspeed.ops.triton_ops._triton import _TRITON_AVAILABLE, triton, tl |
There was a problem hiding this comment.
Moving the module without a compatibility shim breaks existing users that import the previously exposed deepspeed.moe.group_gemm_triton module or its is_available helper, even though the implementation still exists under the new package. Please leave a thin deepspeed/moe/group_gemm_triton.py re-export, including an is_available alias to is_triton_available, so downstream code does not start failing at import time.
Useful? React with 👍 / 👎.
1 participant
Performance:
SwiGLU Benchmark
Configuration:
hidden=1024,dtype=bf16triton_fusedeagertriton_fusedeagertriton_fusedeager