Skip to content

[Common] Added an optimized gated rowwise MXFP8 SwiGLU kernel #2328

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[Common] Added an optimized gated rowwise MXFP8 SwiGLU kernel #2328

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ea79194
Added specialized Gated MXFP8 SwiGLU kernel for rowwise scaling
Oleg-Goncharov Oct 31, 2025
0eea9c8
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 31, 2025
1df9fed
Fixed issues per the review
Oleg-Goncharov Oct 31, 2025
5f9d8e7
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 31, 2025
bfa4171
Merge branch 'main' into pr_mxfp8_rowwise_swiglu_opt
Oleg-Goncharov Oct 31, 2025
25ed43e
Removed variable shadowing
Oleg-Goncharov Oct 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fixed issues per the review 
Signed-off-by: Oleg Goncharov <[email protected]>
  • Loading branch information
@Oleg-Goncharov
Oleg-Goncharov committed Oct 31, 2025
commit 1df9fed76ecb8c02be8b9744e4f8c681d569af78
18 changes: 10 additions & 8 deletions transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -699,14 +699,16 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu

// Optimized BWD/FWD SwiGLU MXFP8 Rowwise kernels for BF16/FP16 inputs
if constexpr (!std::is_same<ParamOP, ClampedSwiGLUParam>::value) {
if constexpr ((!IS_BWD && (ActOP == &silu<fp32, fp32>)) ||
(IS_BWD && (ActOP == &silu<fp32, fp32>) && (DActOP == &dsilu<fp32, fp32>))) {
if (((gated_input.dtype() == DType::kFloat16) || (gated_input.dtype() == DType::kBFloat16)) &&
(scaling_type == ScalingType::ROWWISE)) {
quantize_gated_rowwise<IS_BWD, ParamOP, ActOP, DActOP>(grad, gated_input, output, p,
stream);
return;
}
const bool is_fwd_swiglu = !IS_BWD && (ActOP == &silu<fp32, fp32>);
const bool is_bwd_swiglu = IS_BWD && (ActOP == &silu<fp32, fp32>) &&
(DActOP == &dsilu<fp32, fp32>);
const bool is_supported_data_type = (gated_input.dtype() == DType::kFloat16) ||
(gated_input.dtype() == DType::kBFloat16);
const bool is_supported_scaling_type = scaling_type == ScalingType::ROWWISE;
if (is_supported_data_type && is_supported_scaling_type && (is_fwd_swiglu || is_bwd_swiglu)) {
quantize_gated_rowwise<IS_BWD, ParamOP, ActOP, DActOP>(grad, gated_input, output, p,
stream);
return;
}
}

Expand Down
4 changes: 2 additions & 2 deletions transformer_engine/common/cast/mxfp8/specialized/gated_mxfp8_rowwise_swiglu.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,11 +251,11 @@ __device__ __forceinline__ void compute_bwd_gated_activation(
dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; // Derivative of clamp
gate_elt = min(max(-p.limit, gate_elt), p.limit) + 1.0f;
}
const float x = act_elt;
float x = act_elt;
float act_x;
float dact_x;
if constexpr (IS_CLAMPED_SWIGLU) {
const float x = min(act_elt, p.limit);
x = min(act_elt, p.limit);
const float s = sigmoidf(p.alpha * x);
act_x = x * s;
dact_x = act_elt <= p.limit ? s + s * (1 - s) * p.alpha * x : 0.0f;
Expand Down
Loading