Skip to content

[Fix] [Test] Remove unused GeLU/SquaredReLU NPU kernels and add Ascend UB Manager tests #1006

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhiyuan1i merged 1 commit into from
Jul 6, 2026
Merged

[Fix] [Test] Remove unused GeLU/SquaredReLU NPU kernels and add Ascend UB Manager tests #1006

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions fla/modules/activations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -485,7 +485,6 @@ def backward(ctx, grad_output):
# this function is tanh approximation of gelu
# actual gelu is:
# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
@dispatch('modules')
@torch.compile
def gelu_fwd(x):
return (x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))).to(dtype=x.dtype)
Expand All @@ -494,7 +493,6 @@ def gelu_fwd(x):
# gradient of tanh approximation of gelu
# gradient of actual gelu is:
# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
@dispatch('modules')
@torch.compile
def gelu_bwd(g, x):
tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
Expand Down Expand Up @@ -527,14 +525,12 @@ def relu_bwd(g, x):
return torch.where(x >= 0, g, 0.0).to(dtype=x.dtype)


@dispatch('modules')
@torch.compile
def sqrelu_fwd(x):
r = F.relu(x.float())
return (r * r).to(dtype=x.dtype)


@dispatch('modules')
@torch.compile
def sqrelu_bwd(g, x):
return (2.0 * g * F.relu(x.float())).to(dtype=x.dtype)
Expand Down
16 changes: 0 additions & 16 deletions fla/modules/backends/triton_ascend/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,22 +172,6 @@ def swiglu_linear(self, x, y, weight, bias):
from fla.modules.backends.triton_ascend.activations import swiglu_linear_npu
return swiglu_linear_npu(x, y, weight, bias)

def gelu_fwd(self, x):
from fla.modules.backends.triton_ascend.activations import gelu_fwd_npu
return gelu_fwd_npu(x)

def gelu_bwd(self, g, x):
from fla.modules.backends.triton_ascend.activations import gelu_bwd_npu
return gelu_bwd_npu(g, x)

def sqrelu_fwd(self, x):
from fla.modules.backends.triton_ascend.activations import sqrelu_fwd_npu
return sqrelu_fwd_npu(x)

def sqrelu_bwd(self, g, x):
from fla.modules.backends.triton_ascend.activations import sqrelu_bwd_npu
return sqrelu_bwd_npu(g, x)

def powglu_fwd(self, x, y, power=3.0, output_contiguous=False):
from fla.modules.backends.triton_ascend.activations import powglu_fwd_npu
return powglu_fwd_npu(x, y, power=power, output_contiguous=output_contiguous)
Expand Down
168 changes: 0 additions & 168 deletions fla/modules/backends/triton_ascend/activations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,174 +352,6 @@ def swiglu_fwdbwd_kernel(
tl.store(z + z_off, z_val.to(z.dtype.element_ty), mask=mask)


@triton.heuristics(_LINEAR_HEURISTICS_XY)
@triton.jit(do_not_specialize=['T'])
def gelu_fwd_kernel(
x, y,
T,
D: tl.constexpr,
stride_x_row,
stride_y_row,
B: tl.constexpr,
X_LINEAR: tl.constexpr,
Y_LINEAR: tl.constexpr,
):
pid = tl.program_id(0)
offs = pid * B + tl.arange(0, B)
mask = offs < T
x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
t = 0.79788456 * x_val * (1.0 + 0.044715 * x_val * x_val)
tanh_out = tl.tanh(t)
y_val = x_val * 0.5 * (1.0 + tanh_out)
tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask)


@triton.heuristics(_LINEAR_HEURISTICS_BWD)
@triton.jit(do_not_specialize=['T'])
def gelu_bwd_kernel(
x, dy, dx,
T,
D: tl.constexpr,
stride_x_row,
stride_dy_row,
stride_dx_row,
B: tl.constexpr,
X_LINEAR: tl.constexpr,
DY_LINEAR: tl.constexpr,
DX_LINEAR: tl.constexpr,
):
pid = tl.program_id(0)
offs = pid * B + tl.arange(0, B)
mask = offs < T
x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
t = 0.79788456 * x_val * (1.0 + 0.044715 * x_val * x_val)
tanh_out = tl.tanh(t)
ff = 0.5 * x_val * (
(1.0 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x_val * x_val)
) + 0.5 * (1.0 + tanh_out)
dx_val = ff * g_val
tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)


@triton.heuristics(_LINEAR_HEURISTICS_XY)
@triton.jit(do_not_specialize=['T'])
def sqrelu_fwd_kernel(
x, y,
T,
D: tl.constexpr,
stride_x_row,
stride_y_row,
B: tl.constexpr,
X_LINEAR: tl.constexpr,
Y_LINEAR: tl.constexpr,
):
pid = tl.program_id(0)
offs = pid * B + tl.arange(0, B)
mask = offs < T
x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
y_off = _flat_offset(offs, D, stride_y_row, Y_LINEAR)
x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
r = tl.maximum(x_val, 0.0)
y_val = r * r
tl.store(y + y_off, y_val.to(y.dtype.element_ty), mask=mask)


@triton.heuristics(_LINEAR_HEURISTICS_BWD)
@triton.jit(do_not_specialize=['T'])
def sqrelu_bwd_kernel(
x, dy, dx,
T,
D: tl.constexpr,
stride_x_row,
stride_dy_row,
stride_dx_row,
B: tl.constexpr,
X_LINEAR: tl.constexpr,
DY_LINEAR: tl.constexpr,
DX_LINEAR: tl.constexpr,
):
pid = tl.program_id(0)
offs = pid * B + tl.arange(0, B)
mask = offs < T
x_off = _flat_offset(offs, D, stride_x_row, X_LINEAR)
dy_off = _flat_offset(offs, D, stride_dy_row, DY_LINEAR)
dx_off = _flat_offset(offs, D, stride_dx_row, DX_LINEAR)
x_val = tl.load(x + x_off, mask=mask, other=0.).to(tl.float32)
g_val = tl.load(dy + dy_off, mask=mask, other=0.).to(tl.float32)
dx_val = 2.0 * g_val * tl.maximum(x_val, 0.0)
tl.store(dx + dx_off, dx_val.to(dx.dtype.element_ty), mask=mask)


@torch.compiler.disable
def gelu_fwd_npu(x: torch.Tensor) -> torch.Tensor:
x = _ensure_inner_contiguous(x)
T, D = x.numel(), x.shape[-1]
y = _alloc_output(x)
grid, B = _activation_launch_config(T)
gelu_fwd_kernel[grid](
x, y, T=T, D=D,
stride_x_row=_get_stride(x),
stride_y_row=_get_stride(y),
BLOCK_SIZE=B,
)
return y


@torch.compiler.disable
def gelu_bwd_npu(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
x = _ensure_inner_contiguous(x)
g = _ensure_inner_contiguous(g)
T, D = x.numel(), x.shape[-1]
dx = _alloc_output(x)
grid, B = _activation_launch_config(T, is_backward=True)
gelu_bwd_kernel[grid](
x, g, dx, T=T, D=D,
stride_x_row=_get_stride(x),
stride_dy_row=_get_stride(g),
stride_dx_row=_get_stride(dx),
BLOCK_SIZE=B,
)
return dx


@torch.compiler.disable
def sqrelu_fwd_npu(x: torch.Tensor) -> torch.Tensor:
x = _ensure_inner_contiguous(x)
T, D = x.numel(), x.shape[-1]
y = _alloc_output(x)
grid, B = _activation_launch_config(T)
sqrelu_fwd_kernel[grid](
x, y, T=T, D=D,
stride_x_row=_get_stride(x),
stride_y_row=_get_stride(y),
BLOCK_SIZE=B,
)
return y


@torch.compiler.disable
def sqrelu_bwd_npu(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
x = _ensure_inner_contiguous(x)
g = _ensure_inner_contiguous(g)
T, D = x.numel(), x.shape[-1]
dx = _alloc_output(x)
grid, B = _activation_launch_config(T, is_backward=True)
sqrelu_bwd_kernel[grid](
x, g, dx, T=T, D=D,
stride_x_row=_get_stride(x),
stride_dy_row=_get_stride(g),
stride_dx_row=_get_stride(dx),
BLOCK_SIZE=B,
)
return dx


@torch.compiler.disable
def sigmoid_fwd_npu(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor:
x = _ensure_inner_contiguous(x)
Expand Down
Loading
Loading