29 lines
1.1 KiB
Python
29 lines
1.1 KiB
Python
# NO_CHIP_SPECIFIC
|
||
"""FlagOS S2 · Task 29 gelu_and_mul — 通用版（国际通用芯片：原生 Triton，tl.math.erf）"""
|
||
import torch
|
||
import triton
|
||
import triton.language as tl
|
||
|
||
|
||
@triton.jit
|
||
def _gelu_and_mul_kernel(x_ptr, out_ptr, d, BLOCK: tl.constexpr):
|
||
row = tl.program_id(0).to(tl.int64)
|
||
x_row = x_ptr + row * (2 * d)
|
||
out_row = out_ptr + row * d
|
||
offs = tl.arange(0, BLOCK)
|
||
for start in range(0, d, BLOCK):
|
||
idx = start + offs
|
||
mask = idx < d
|
||
x1 = tl.load(x_row + idx, mask=mask, other=0.0).to(tl.float32)
|
||
x2 = tl.load(x_row + d + idx, mask=mask, other=0.0).to(tl.float32)
|
||
g = 0.5 * x1 * (1.0 + tl.math.erf(x1 * 0.7071067811865476))
|
||
tl.store(out_row + idx, (g * x2).to(out_ptr.dtype.element_ty), mask=mask)
|
||
|
||
|
||
def gelu_and_mul(x: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
|
||
num_tokens, two_d = x.shape
|
||
d = two_d // 2
|
||
if out is None:
|
||
out = torch.empty((num_tokens, d), dtype=x.dtype, device=x.device)
|
||
_gelu_and_mul_kernel[(num_tokens,)](x, out, d, BLOCK=1024, num_warps=4)
|
||
return out
|