module 'triton.language.math' has no attribute 'tanh'

[NoahBPeterson]

autokernel/kernels/fused_mlp.py

Line 92 in 7843582

gate_activated = 0.5 * acc_gate * (1.0 + tl.math.tanh(0.7978845608 * (acc_gate + 0.044715 * acc_gate * acc_gate * acc_gate)))

Here's a minimal reproduction of this bug, using fused_mlp.py as an example:

import triton
import triton.language as tl
from triton.backends.compiler import GPUTarget
from triton.compiler import ASTSource


@triton.jit
def gelu_kernel(x_ptr, out_ptr, N, BLOCK: tl.constexpr):
    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    x = tl.load(x_ptr + offs, mask=mask)
    y = 0.5 * x * (1.0 + tl.math.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))
    tl.store(out_ptr + offs, y, mask=mask)


src = ASTSource(
    fn=gelu_kernel,
    signature={"x_ptr": "*fp32", "out_ptr": "*fp32", "N": "i32", "BLOCK": "constexpr"},
    constexprs={"BLOCK": 256},
)
triton.compile(src, target=GPUTarget("cuda", 89, 32))

Traceback (most recent call last):
  File "/workspace/repro_tanh.py", line 27, in <module>
    triton.compile(src, target=GPUTarget("cuda", 89, 32))
  File ".../triton/compiler/compiler.py", line 307, in compile
    module = src.make_ir(target, options, codegen_fns, module_map, context)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../triton/compiler/compiler.py", line 80, in make_ir
    return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
triton.compiler.errors.CompilationError: at 5:25:
def gelu_kernel(x_ptr, out_ptr, N, BLOCK: tl.constexpr):
    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    x = tl.load(x_ptr + offs, mask=mask)
    y = 0.5 * x * (1.0 + tl.math.tanh(0.7978845608 * (x + 0.044715 * x * x * x)))
                         ^
AttributeError("module 'triton.language.math' has no attribute 'tanh'")