`core.fusions.fused_bias_swiglu`#

Module Contents#

Classes#

`BiasSwiGLUFunction`	Custom autograd function for SwiGLU activation with bias support.
`SwiGLUFunction`	Custom autograd function for SwiGLU activation without bias.
`WeightedSwiGLUFunction`

Functions#

`swiglu`	Performs SwiGLU (Swish-Gated Linear Unit) activation function.
`bias_swiglu`	Performs SwiGLU activation with bias addition.
`weighted_swiglu`
`clamped_swiglu`	Perform SwiGLU after clamping both halves of the input.
`bias_clamped_swiglu`	Perform clamped SwiGLU after bias addition.
`clamped_weighted_swiglu`	Perform token-weighted clamped SwiGLU.
`_tanh_clamp_and_deriv`	Soft-clamps `y` and returns the clamp derivative alongside it, both in fp32.
`situ_glu`	SiTU-GLU: `s_g * tanh(y1 / s_g) * sigmoid(y1) * y2`, with `y2` optionally clamped.
`bias_situ_glu`	Bias addition followed by SiTU-GLU. The clamps act on the full pre-activation.
`weighted_situ_glu`
`swiglu_back`	Computes the gradient for the SwiGLU activation function.
`bias_swiglu_back`	Computes the gradient for the biased SwiGLU activation function.
`weighted_swiglu_back`
`clamped_swiglu_back`	Compute the input gradient for clamped SwiGLU.
`bias_clamped_swiglu_back`	Compute the input gradient for clamped SwiGLU with bias.
`clamped_weighted_swiglu_back`	Compute input and weight gradients for token-weighted clamped SwiGLU.
`_situ_glu_grads`	Gradients of SiTU-GLU w.r.t. the two halves of `y`, kept in fp32.
`situ_glu_back`
`bias_situ_glu_back`
`weighted_situ_glu_back`
`bias_swiglu_impl`	Implementation of biased SwiGLU that handles different input shapes.
`weighted_bias_swiglu_impl`	Token-wise-weighted bias swiglu fusion.

API#

core.fusions.fused_bias_swiglu.swiglu(y)#

Performs SwiGLU (Swish-Gated Linear Unit) activation function.

Parameters:: y (torch.Tensor) – Input tensor to be split into two halves along the last dimension.
Returns:: Result of SwiGLU activation: SiLU(y1) * y2, where y1, y2 are the split halves.
Return type:: torch.Tensor

core.fusions.fused_bias_swiglu.bias_swiglu(y, bias)#

Performs SwiGLU activation with bias addition.

Parameters:

y (torch.Tensor) – Input tensor.
bias (torch.Tensor) – Bias tensor to be added to input.

Returns:

Result of bias addition followed by SwiGLU activation.

Return type:

torch.Tensor

core.fusions.fused_bias_swiglu.weighted_swiglu(y, weights)#

core.fusions.fused_bias_swiglu.clamped_swiglu(y, clamp_value)#: Perform SwiGLU after clamping both halves of the input.

core.fusions.fused_bias_swiglu.bias_clamped_swiglu(y, bias, clamp_value)#: Perform clamped SwiGLU after bias addition.

core.fusions.fused_bias_swiglu.clamped_weighted_swiglu(y, weights, clamp_value)#: Perform token-weighted clamped SwiGLU.

core.fusions.fused_bias_swiglu._tanh_clamp_and_deriv(y: torch.Tensor, clamp_scale: float)#: Soft-clamps y and returns the clamp derivative alongside it, both in fp32.

core.fusions.fused_bias_swiglu.situ_glu( y, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#: SiTU-GLU: s_g * tanh(y1 / s_g) * sigmoid(y1) * y2, with y2 optionally clamped.

core.fusions.fused_bias_swiglu.bias_situ_glu( y, bias, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#: Bias addition followed by SiTU-GLU. The clamps act on the full pre-activation.

core.fusions.fused_bias_swiglu.weighted_situ_glu( y, weights, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#

core.fusions.fused_bias_swiglu.swiglu_back(g, y)#

Computes the gradient for the SwiGLU activation function.

Parameters:

g (torch.Tensor) – Gradient tensor from the subsequent layer.
y (torch.Tensor) – Input tensor that was used in the forward pass.

Returns:

Gradient with respect to the input tensor, computed using the chain rule and the derivative of the SiLU activation function.

Return type:

torch.Tensor

core.fusions.fused_bias_swiglu.bias_swiglu_back(g, y, bias)#

Computes the gradient for the biased SwiGLU activation function.

Parameters:

g (torch.Tensor) – Gradient tensor from the subsequent layer.
y (torch.Tensor) – Input tensor that was used in the forward pass.
bias (torch.Tensor) – Bias tensor that was added in the forward pass.

Returns:

Gradient with respect to the input tensor, computed after applying the bias addition.

Return type:

torch.Tensor

core.fusions.fused_bias_swiglu.weighted_swiglu_back(g, y, weights)#

core.fusions.fused_bias_swiglu.clamped_swiglu_back(g, y, clamp_value)#: Compute the input gradient for clamped SwiGLU.

core.fusions.fused_bias_swiglu.bias_clamped_swiglu_back(g, y, bias, clamp_value)#: Compute the input gradient for clamped SwiGLU with bias.

core.fusions.fused_bias_swiglu.clamped_weighted_swiglu_back(g, y, weights, clamp_value)#: Compute input and weight gradients for token-weighted clamped SwiGLU.

core.fusions.fused_bias_swiglu._situ_glu_grads( g, y, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#: Gradients of SiTU-GLU w.r.t. the two halves of y, kept in fp32.

core.fusions.fused_bias_swiglu.situ_glu_back( g, y, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#

core.fusions.fused_bias_swiglu.bias_situ_glu_back( g, y, bias, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#

core.fusions.fused_bias_swiglu.weighted_situ_glu_back( g, y, weights, gate_clamp_scale: float, linear_clamp_scale: Optional[float], )#

class core.fusions.fused_bias_swiglu.BiasSwiGLUFunction#

Bases: torch.autograd.Function

Custom autograd function for SwiGLU activation with bias support.

static forward( ctx, input, bias, fp8_input_store, cpu_offload_input, clamp_value, gate_clamp_scale, linear_clamp_scale, )#

Forward pass of biased SwiGLU activation.

Parameters:

ctx – Autograd context object for saving tensors for backward pass.
input (torch.Tensor) – Input tensor to apply SwiGLU to.
bias (torch.Tensor) – Bias tensor to be added to input before SwiGLU.
fp8_input_store (bool) – If True, stores intermediate values in FP8 format.
clamp_value (Optional[float]) – If set, hard-clamp both halves of the input before SwiGLU. Mutually exclusive with gate_clamp_scale.
gate_clamp_scale (Optional[float]) – If set, use SiTU-GLU instead of SwiGLU, with the gate’s linear factor soft-clamped to this scale.
linear_clamp_scale (Optional[float]) – If set, also soft-clamp the linear half to this scale. Requires gate_clamp_scale.

Returns:

Result of applying bias addition followed by SwiGLU activation.

Return type:

torch.Tensor

static backward(ctx, grad_output)#

Backward pass of biased SwiGLU activation.

Parameters:

ctx – Autograd context object containing saved tensors from forward pass.
grad_output (torch.Tensor) – Gradient of the loss with respect to the output.

Returns:

Tuple containing: - Gradient with respect to the input tensor - Gradient with respect to the bias tensor - None for fp8_input_store parameter

Return type:

tuple

class core.fusions.fused_bias_swiglu.SwiGLUFunction#

Bases: torch.autograd.Function

Custom autograd function for SwiGLU activation without bias.

static forward( ctx, input, fp8_input_store, cpu_offload_input, clamp_value, gate_clamp_scale, linear_clamp_scale, )#

Forward pass of SwiGLU activation.

Parameters:

ctx – Autograd context object for saving tensors for backward pass.
input (torch.Tensor) – Input tensor to apply SwiGLU to.
fp8_input_store (bool) – If True, stores intermediate values in FP8 format.
clamp_value (Optional[float]) – If set, hard-clamp both halves of the input before SwiGLU. Mutually exclusive with gate_clamp_scale.
gate_clamp_scale (Optional[float]) – If set, use SiTU-GLU instead of SwiGLU, with the gate’s linear factor soft-clamped to this scale.
linear_clamp_scale (Optional[float]) – If set, also soft-clamp the linear half to this scale. Requires gate_clamp_scale.

Returns:

Result of applying SwiGLU activation.

Return type:

torch.Tensor

static backward(ctx, grad_output)#

Backward pass of SwiGLU activation.

Parameters:

ctx – Autograd context object containing saved tensors from forward pass.
grad_output (torch.Tensor) – Gradient of the loss with respect to the output.

Returns:

Tuple containing: - Gradient with respect to the input tensor - None for fp8_input_store parameter

Return type:

tuple

class core.fusions.fused_bias_swiglu.WeightedSwiGLUFunction#

Bases: torch.autograd.Function

static forward( ctx, input, weights, fp8_input_store, clamp_value, gate_clamp_scale, linear_clamp_scale, )#

static backward(ctx, grad_output)#

core.fusions.fused_bias_swiglu.bias_swiglu_impl( input, bias, fp8_input_store=False, cpu_offload_input=False, clamp_value=None, gate_clamp_scale: Optional[float] = None, linear_clamp_scale: Optional[float] = None, )#

Implementation of biased SwiGLU that handles different input shapes.

This function reshapes the input if necessary, applies the SwiGLU activation (with or without bias), and restores the original shape.

Parameters:

input (torch.Tensor) – Input tensor to apply SwiGLU activation.
bias (torch.Tensor, optional) – Bias tensor to be added to input. If None, uses the bias-free SwiGLU variant.
fp8_input_store (bool, optional) – Whether to store intermediate values in FP8 format. Defaults to False.
cpu_offload_input (bool, optional) – Whether to mark saved activation inputs for CPU offloading. Defaults to False.
clamp_value (float, optional) – Maximum gate value and absolute linear value. When None, preserve the legacy unclamped SwiGLU behavior. Mutually exclusive with gate_clamp_scale.
gate_clamp_scale (Optional[float]) – If set, compute SiTU-GLU instead of SwiGLU: the gate becomes s_g * tanh(x / s_g) * sigmoid(x), bounded by s_g.
linear_clamp_scale (Optional[float]) – If set, also soft-clamp the linear half, bounding the output by gate_clamp_scale * linear_clamp_scale. Requires gate_clamp_scale.

Returns:

Result of biased SwiGLU activation.

Return type:

torch.Tensor

Raises:

AssertionError – If input tensor does not have 2 or 3 dimensions.

core.fusions.fused_bias_swiglu.weighted_bias_swiglu_impl( input, bias, weights, fp8_input_store=False, clamp_value=None, gate_clamp_scale: Optional[float] = None, linear_clamp_scale: Optional[float] = None, )#

Token-wise-weighted bias swiglu fusion.

Parameters:

clamp_value (float, optional) – Maximum gate value and absolute linear value. When None, preserve the legacy unclamped SwiGLU behavior. Mutually exclusive with gate_clamp_scale.
gate_clamp_scale (Optional[float]) – If set, compute SiTU-GLU instead of SwiGLU: the gate becomes s_g * tanh(x / s_g) * sigmoid(x), bounded by s_g.
linear_clamp_scale (Optional[float]) – If set, also soft-clamp the linear half. Requires gate_clamp_scale.

core.fusions.fused_bias_swiglu#

Module Contents#

Classes#

Functions#

API#

`core.fusions.fused_bias_swiglu`#