34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
"""FlagOS S2 · Task 29 gelu_and_mul — {CHIP_CN}版（FlagTree/Triton-TLE 环境）
|
||
|
||
写法严格照官方《Triton和Triton-TLE入门指南》示例：
|
||
import triton / triton.language 照常（FlagTree 提供兼容层）；
|
||
数学原语用 tle.erf（跨芯片统一，TLE-Lite 层）。
|
||
"""
|
||
import torch
|
||
import triton
|
||
import triton.language as tl
|
||
import flagtree.tle as tle
|
||
|
||
|
||
@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 + tle.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
|