| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm.triton_utils import tl, triton |
| 5 | |
| 6 | |
| 7 | @triton.jit |
| 8 | def fast_exp(x): |
| 9 | """Faster alternative to tl.exp() using the hardware exp2 instruction. |
| 10 | |
| 11 | tl.math.exp2 maps directly to a single ex2.approx.f32 PTX instruction, |
| 12 | while tl.exp goes through libdevice __nv_expf which adds function call |
| 13 | overhead and extra range checking. |
| 14 | """ |
| 15 | # exp(x) = exp2(x * log2(e)), where log2(e) = 1/ln(2) = 1.4426950408889634 |
| 16 | LOG2E = tl.constexpr(1.4426950408889634) |
| 17 | return tl.math.exp2(LOG2E * x) |