diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh index a54ef62d69..9dc612a5f7 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh @@ -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 struct Nvfp4LinearSwiGluTmaTensorStorage { static_assert(Schedule::kBlockN == 128); @@ -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)); } }