[KERNELS] Save some instructions in swiglu (#8801)

Say `gelu` has n registers per thread. Currently, `exp(-alpha * gelu)`
takes 1 `sub.f32` and 2n `mul.f32` instructions since `exp(x)` gets
expanded to `exp2(log2(e) * x)` by the time we get to ptx. We can
rewrite this as scaling `gelu` by the scalar value `(-alpha * log2(e))`
which is just 1 + n `mul.f32` instructions.

I also use `ex2.approx.ftz.f32` which is a single `MUFU.EX2` in SASS,
compared to the non-ftz variant which requires multiple SASS
instructions to handle denormal values. This is fine numerically since
we add 1 to the result anyway which will round out anything below
epsilon.

Author: peterbell10

python/triton_kernels/triton_kernels/swiglu_details/_swiglu.py

@@ -35,6 +35,21 @@ def swiglu_launch_metadata(grid, kernel, args):
     return ret
 
 
+@triton.jit
+def exp2_ftz(x):
+    if tl.target_info.is_cuda():
+        return tl.inline_asm_elementwise(
+            "ex2.approx.ftz.f32 $0, $1;",
+            "=r, r",
+            [x],
+            dtype=tl.float32,
+            is_pure=True,
+            pack=1,
+        )
+    else:
+        return tl.exp2(x)
+
+
 @triton.jit
 def compute_swiglu(gelu, linear, scale, alpha, limit):
     gelu = gelu.to(tl.float32) * scale
@@ -43,7 +58,11 @@ def compute_swiglu(gelu, linear, scale, alpha, limit):
     linear = linear.to(tl.float32) * scale
     if limit is not None:
         linear = clip(linear, limit, clip_lower=True)
-    s = gelu / (1 + tl.exp(-alpha * gelu))
+
+    # exp(x) becomes exp2(log2(e) * x) in ptx. By expanding it early, we can factor
+    # (-alpha * log2_e) into a single scalar factor.
+    log2_e: tl.constexpr = 1.4426950408889634
+    s = gelu / (1 + exp2_ftz((-alpha * log2_e) * gelu))
     return tl.fma(s, linear, s)  # (s * (linear + 1))