Skip to content

perf(nvfp4): drop the guarded expf slow path from the fused SwiGLU epilogue #194

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

Closed
MichaelDementii wants to merge 2 commits into from
Closed

perf(nvfp4): drop the guarded expf slow path from the fused SwiGLU epilogue #194

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
61250e8
perf(nvfp4): drop the guarded expf slow path from the fused SwiGLU ep…
ARVLab Sep 6, 2026
dc108d4
fix(nvfp4): keep the approximate SiLU off the zeroing range of __fdiv…
MichaelDementii Sep 9, 2026
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
29 changes: 25 additions & 4 deletions src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,27 @@

namespace ninfer::ops::detail {

// SwiGLU epilogue activation.
//
// The accurate `silu` compiles to a guarded slow path for `expf` (64
// CALL.REL.NOINC plus 64 FCHK per kernel in SASS) and a Newton-refined divide,
// and the result is then rounded to bf16 on the very next instruction. The
// hardware-approximate forms are ~1e-6 relative, four orders of magnitude below
// the bf16 quantum, so nothing that survives the rounding is affected.
//
// The exponential is folded onto the side that cannot overflow. `__fdividef`
// returns zero once the divisor reaches 2^126, which `1 + __expf(-x)` does for
// x below -87.34, and SiLU there is still a normal bf16 (-9.6e-37 at that
// edge). Written this way the divisor stays in (1, 2] for every finite x, and
// the only subnormal the form can produce, `e`, enters a multiply rather than
// the divide. It costs nothing: both forms compile to the same 40 SASS
// instructions with two MUFU and no CALL.
__device__ __forceinline__ float swiglu_silu(float x) {
const float e = __expf(-fabsf(x));
const float r = __fdividef(1.0f, 1.0f + e);
return (x >= 0.0f ? x : x * e) * r;
}

template <class Schedule>
struct Nvfp4LinearSwiGluTmaTensorStorage {
static_assert(Schedule::kBlockN == 128);
Expand Down Expand Up @@ -236,10 +257,10 @@ __global__ __launch_bounds__(
shared_output + token1 * kOutputStride + pair_row);
const auto& gate = accumulators[mma_m][mma_n];
const auto& up = accumulators[mma_m][mma_n + kGateMmaFragments];
*destination0 = __floats2bfloat162_rn(silu(gate[0] * alpha) * (up[0] * alpha),
silu(gate[1] * alpha) * (up[1] * alpha));
*destination1 = __floats2bfloat162_rn(silu(gate[2] * alpha) * (up[2] * alpha),
silu(gate[3] * alpha) * (up[3] * alpha));
*destination0 = __floats2bfloat162_rn(swiglu_silu(gate[0] * alpha) * (up[0] * alpha),
swiglu_silu(gate[1] * alpha) * (up[1] * alpha));
*destination1 = __floats2bfloat162_rn(swiglu_silu(gate[2] * alpha) * (up[2] * alpha),
swiglu_silu(gate[3] * alpha) * (up[3] * alpha));
}
}

Expand Down