Is this page helpful?

core.fusions.fused_weighted_squared_relu#

Module Contents#

Classes#

WeightedSquaredReLUFunction

Autograd wrapper around the (optionally weighted, optionally clamped) Squared-ReLU fused kernels.

Functions#

weighted_squared_relu

Element-wise weight applied after Squared-ReLU.

_squared_relu_back

Gradient of Squared-ReLU.

weighted_squared_relu_back

Backward for weighted Squared-ReLU.

_tanh_relu_over_scale

tanh(ReLU(x) / clamp_scale) in fp32, the shared term of the clamped forward/backward.

weighted_clamped_squared_relu

Element-wise weight applied after tanh soft-clamped Squared-ReLU.

weighted_clamped_squared_relu_back

Backward for weighted tanh soft-clamped Squared-ReLU.

clamped_squared_relu

Tanh-soft-clamped squared-ReLU without token weights.

clamped_squared_relu_back

Backward for tanh-soft-clamped squared-ReLU, recomputed from the raw input.

weighted_squared_relu_impl

Squared-ReLU fusion with optional per-token weights and optional tanh soft clamping.

API#

core.fusions.fused_weighted_squared_relu.weighted_squared_relu(
x: torch.Tensor,
weights: torch.Tensor,
) → torch.Tensor#

Element-wise weight applied after Squared-ReLU.

Parameters:

  • x (torch.Tensor) – Input tensor.

  • weights (torch.Tensor) – Weight tensor that will be broadcast-multiplied with the activation result. Typically of shape (B, 1) so it can be broadcast across the hidden dimension.

Returns:

squared_relu(x) * weights with original dtype preserved.

Return type:

torch.Tensor

core.fusions.fused_weighted_squared_relu._squared_relu_back(g: torch.Tensor, x: torch.Tensor) → torch.Tensor#

Gradient of Squared-ReLU.

The derivative of (ReLU(x))^2 w.r.t x is 2 * ReLU(x).

core.fusions.fused_weighted_squared_relu.weighted_squared_relu_back(
g: torch.Tensor,
x: torch.Tensor,
weights: torch.Tensor,
)#

Backward for weighted Squared-ReLU.

Returns gradients w.r.t x and weights.

core.fusions.fused_weighted_squared_relu._tanh_relu_over_scale(
x: torch.Tensor,
clamp_scale: float,
) → torch.Tensor#

tanh(ReLU(x) / clamp_scale) in fp32, the shared term of the clamped forward/backward.

core.fusions.fused_weighted_squared_relu.weighted_clamped_squared_relu(
x: torch.Tensor,
weights: torch.Tensor,
clamp_scale: float,
) → torch.Tensor#

Element-wise weight applied after tanh soft-clamped Squared-ReLU.

Parameters:

  • x (torch.Tensor) – Input tensor.

  • weights (torch.Tensor) – Weight tensor that will be broadcast-multiplied with the activation result.

  • clamp_scale (float) – The soft-clamp scale s.

Returns:

The weighted activation with original dtype preserved.

Return type:

torch.Tensor

core.fusions.fused_weighted_squared_relu.weighted_clamped_squared_relu_back(
g: torch.Tensor,
x: torch.Tensor,
weights: torch.Tensor,
clamp_scale: float,
)#

Backward for weighted tanh soft-clamped Squared-ReLU.

Returns gradients w.r.t x and weights.

core.fusions.fused_weighted_squared_relu.clamped_squared_relu(
x: torch.Tensor,
clamp_scale: float,
) → torch.Tensor#

Tanh-soft-clamped squared-ReLU without token weights.

This matches squared_relu(tanh_soft_clamp(x, clamp_scale)). The clamped pre-activation and the square stay in FP32 and only the result is rounded back to the input dtype.

core.fusions.fused_weighted_squared_relu.clamped_squared_relu_back(
g: torch.Tensor,
x: torch.Tensor,
clamp_scale: float,
)#

Backward for tanh-soft-clamped squared-ReLU, recomputed from the raw input.

class core.fusions.fused_weighted_squared_relu.WeightedSquaredReLUFunction#

Bases: torch.autograd.Function

Autograd wrapper around the (optionally weighted, optionally clamped) Squared-ReLU fused kernels.

Only the raw input (and the weights, when given) is saved for backward. The unfused squared_relu(tanh_soft_clamp(x)) ordering saves both x and the clamped intermediate, doubling the activation memory kept alive for this op.

static forward(
ctx,
input: torch.Tensor,
weights: Optional[torch.Tensor],
clamp_scale: Optional[float],
)#

forward method for WeightedSquaredReLUFunction

Parameters:

  • ctx – context object to store intermediate tensors.

  • input (torch.Tensor) – input tensor.

  • weights (Optional[torch.Tensor]) – optional per-token weight tensor.

  • clamp_scale (Optional[float]) – if set, soft-clamp the input with clamp_scale * tanh(input / clamp_scale) before the activation.

static backward(ctx, grad_output: torch.Tensor)#

backward method for WeightedSquaredReLUFunction

Parameters:

  • ctx – context object to store intermediate tensors.

  • grad_output (torch.Tensor) – gradient of the output of the forward function.

core.fusions.fused_weighted_squared_relu.weighted_squared_relu_impl(
input: torch.Tensor,
weights: Optional[torch.Tensor] = None,
clamp_scale: Optional[float] = None,
) → torch.Tensor#

Squared-ReLU fusion with optional per-token weights and optional tanh soft clamping.

Parameters:

  • input (torch.Tensor) – Input tensor of shape (B, *, hidden_size) where * can be the sequence dimension.

  • weights (Optional[torch.Tensor]) – Optional per-token weights broadcastable to the output of squared_relu once input is flattened to (-1, hidden_size). When None, the clamped squared-ReLU is applied and only input is saved for backward.

  • clamp_scale (Optional[float]) – if set, precondition the input with the tanh soft-clamp clamp_scale * tanh(input / clamp_scale). At least one of weights and clamp_scale must be given.

Returns:

Output tensor with the same shape as input except that the hidden dimension remains unchanged.

Return type:

torch.Tensor