136 lines
5.3 KiB
Python
136 lines
5.3 KiB
Python
# NO_CHIP_SPECIFIC
|
||
"""FlagOS S2 · 赛道一 · Task 29: gelu_and_mul — Triton 实现 v9（三路自适应）
|
||
|
||
冰朔的模具规矩（架构化）：包饺子前先试灶台——三条实现路线在首次调用时
|
||
各自试编译+试运行一个微型饺子，第一条能跑的路线全程使用：
|
||
路1 tle.erf ：FlagOS 官方统一原语（环境装有 FlagTree 0.6.1 且 tle.erf 可用时）
|
||
路2 tl.math.erf：Triton 标准数学库（A100/天数智芯实证可跑，v3 通过正确性）
|
||
路3 exp 拼装 ：最基础零件自拼 GELU（昆仑芯实证可跑，v4 通过；|approx-erf|<=4.8e-4）
|
||
探测只做一次并缓存；探测失败=模具塞不进，立刻明确报错，绝不带病出厨房。
|
||
"""
|
||
import torch
|
||
import triton
|
||
import triton.language as tl
|
||
|
||
try:
|
||
import flagtree.tle as _tle_mod
|
||
except Exception:
|
||
_tle_mod = None
|
||
|
||
|
||
@triton.jit
|
||
def _gelu_and_mul_erf(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)
|
||
|
||
|
||
@triton.jit
|
||
def _gelu_and_mul_tle(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)
|
||
|
||
|
||
@triton.jit
|
||
def _gelu_and_mul_exp(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)
|
||
# GELU tanh 近似；tanh(z)=1-2/(1+exp(2z))，极端值经 inf 饱和落到±1
|
||
inner = 0.7978845608028654 * (x1 + 0.044715 * x1 * x1 * x1)
|
||
g = 0.5 * x1 * (1.0 + (1.0 - 2.0 / (tl.exp(2.0 * inner) + 1.0)))
|
||
tl.store(out_row + idx, (g * x2).to(out_ptr.dtype.element_ty), mask=mask)
|
||
|
||
|
||
# tle 版仅在 flagtree 存在且 tle.erf 真实可用时纳入候选（包存在≠函数存在）
|
||
if _tle_mod is not None and hasattr(_tle_mod, "erf"):
|
||
tle = _tle_mod
|
||
_TLE_KERNEL = _gelu_and_mul_tle
|
||
else:
|
||
tle = None
|
||
_TLE_KERNEL = None
|
||
|
||
_SELECTED = None # 首次调用探测后缓存：'tle' | 'erf' | 'exp'
|
||
|
||
|
||
def _probe(x):
|
||
"""试灶台：对微型输入逐路试编译+试运行+对答案，选第一条活的。"""
|
||
global _SELECTED
|
||
n, d2 = 4, 8
|
||
d = d2 // 2
|
||
for mode, kern in (
|
||
("tle", _TLE_KERNEL),
|
||
("erf", _gelu_and_mul_erf),
|
||
("exp", _gelu_and_mul_exp),
|
||
):
|
||
if kern is None:
|
||
continue
|
||
try:
|
||
xs = torch.randn(n, d2, dtype=torch.float32, device=x.device)
|
||
ref = torch.nn.functional.gelu(xs[:, :d]) * xs[:, d:]
|
||
out = torch.empty(n, d, dtype=xs.dtype, device=xs.device)
|
||
kern[(n,)](xs, out, d, BLOCK=4, num_warps=1)
|
||
if not torch.allclose(out, ref, atol=1e-2, rtol=1e-2):
|
||
continue
|
||
_SELECTED = mode
|
||
return mode
|
||
except Exception:
|
||
continue
|
||
raise RuntimeError("gelu_and_mul: 三条实现路线（tle/erf/exp）在本环境均不可用")
|
||
|
||
|
||
def gelu_and_mul(x: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
|
||
"""x: [num_tokens, 2d] → out: [num_tokens, d]；out[i] = gelu(x[i,:d]) * x[i,d:]"""
|
||
if x.dim() != 2 or x.shape[1] % 2 != 0:
|
||
raise ValueError(f"expect x [n, 2d], got {tuple(x.shape)}")
|
||
if _SELECTED is None:
|
||
_probe(x)
|
||
kern = {"tle": _TLE_KERNEL, "erf": _gelu_and_mul_erf, "exp": _gelu_and_mul_exp}[_SELECTED]
|
||
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)
|
||
kern[(num_tokens,)](x, out, d, BLOCK=1024, num_warps=4)
|
||
return out
|
||
|
||
|
||
# ---------------- 模具：预自检 ----------------
|
||
def self_check(device: str = "cuda", seed: int = 0) -> bool:
|
||
torch.manual_seed(seed)
|
||
cases = [(1, 16), (4, 64), (33, 128), (256, 1024), (3, 4096), (17, 1536), (64, 16384)]
|
||
all_ok = True
|
||
for (n, d), dt in ((c, dt) for c in cases for dt in (torch.float32, torch.float16, torch.bfloat16)):
|
||
x = torch.randn(n, 2 * d, device=device, dtype=dt)
|
||
ref = torch.nn.functional.gelu(x[..., :d]) * x[..., d:]
|
||
got = gelu_and_mul(x)
|
||
ok = torch.allclose(got.float(), ref.float(), atol=1e-2, rtol=1e-2)
|
||
all_ok = all_ok and ok
|
||
print(f"SELF_CHECK[route={_SELECTED}]:", "ALL PASS" if all_ok else "FAIL")
|
||
return all_ok
|
||
|
||
|
||
if __name__ == "__main__":
|
||
self_check()
|