perf(batchnorm): fuse ReLU/LeakyReLU into BatchNorm at finalize

aether/custom_kernels/batchnorm_kernel.py

@@ -5,6 +5,25 @@
 BLOCK_Y = 8
 
 
+# The activation every BatchNorm kernel applies to z = (x - mean) * scale + beta:
+# out = z > 0 ? z : alpha * z, so alpha = 1 is none, 0 is ReLU, else LeakyReLU.
+# Forward and backward compute z from this one source, so the backward's mask is
+# bit-identical to the forward's. As in ReLU's GPU backward, a NaN/inf gradient
+# at an inactive ReLU unit gives 0, not 0 * NaN.
+_ACTIVATION = r'''
+__device__ __forceinline__ float bn_z(float x, float mean, float scale, float beta) {
+    return fmaf(x - mean, scale, beta);
+}
+
+__device__ __forceinline__ float bn_act(float z, float alpha) {
+    return z > 0.0f ? z : alpha * z;
+}
+
+__device__ __forceinline__ float bn_act_grad(float z, float dy, float alpha) {
+    return z > 0.0f ? dy : (alpha != 0.0f ? alpha * dy : 0.0f);
+}
+'''
+
 # Shared by every kernel below. block_sum adds each thread's (a, b) pair over the
 # block's BLOCK_Y rows and returns true on the row-0 threads that hold a channel's
 # total; sum_partials does the same over the G rows of a partials[2, G, C] buffer.
@@ -13,6 +32,8 @@
 #define WAVEFRONT 32
 #define BLOCK_Y   $block_y
 
+$activation
+
 __device__ __forceinline__ bool block_sum(float& a, float& b, int c, int C) {
     __shared__ float sh_a[BLOCK_Y][WAVEFRONT];
     __shared__ float sh_b[BLOCK_Y][WAVEFRONT];
@@ -54,6 +75,9 @@
     const float* __restrict__ p,
     const float* __restrict__ q,      // unused (nullptr) ops still take the param
     const float* __restrict__ mean,   // channel mean [C], unused ops pass nullptr
+    const float* __restrict__ scale,  // grad_sums only, like beta and alpha
+    const float* __restrict__ beta,
+    const float alpha,
     float* __restrict__ partials,     // [2, gridDim.y, C]
     const int N,
     const int C
@@ -156,8 +180,11 @@
 _GRAD_SUMS_OP = {
     "name": "grad_sums",
     "template": _STATS_REDUCE_TEMPLATE,
-    "prologue": "float mu = mean[c];",
-    "accumulate": "float g = p[idx]; acc_a += g; acc_b += g * (q[idx] - mu);",
+    "prologue": "float mu = mean[c]; float sc = scale[c]; float bt = beta[c];",
+    "accumulate": (
+        "float xv = q[idx]; float g = bn_act_grad(bn_z(xv, mu, sc, bt), p[idx], alpha);"
+        " acc_a += g; acc_b += g * (xv - mu);"
+    ),
 }
 
 _FORWARD_FINALIZE_OP = {
@@ -190,6 +217,7 @@ def _get_compiled_kernel(op_dict: dict, variant: str, block_y: int = BLOCK_Y):
         block_y=block_y,
         prologue=op_dict.get("prologue", ""),
         accumulate=op_dict.get("accumulate", ""),
+        activation=_ACTIVATION,
         **vendor,
     )
 
@@ -226,20 +254,26 @@ def get_bn_backward_finalize_kernel(variant: str):
 # folding (mean, scale, beta) into one bias out of the output.
 _bn_apply_ew = config.build_kernel(
     lambda: config.cp.ElementwiseKernel(
-        'float32 x, float32 mean, float32 scale, float32 beta',
+        'float32 x, float32 mean, float32 scale, float32 beta, float32 alpha',
         'float32 out',
-        'out = fmaf(x - mean, scale, beta);',
+        'out = bn_act(bn_z(x, mean, scale, beta), alpha);',
         'bn_apply_ew',
+        preamble=_ACTIVATION,
     ),
     name='_bn_apply_ew',
 )
 
 _bn_backward_ew = config.build_kernel(
     lambda: config.cp.ElementwiseKernel(
-        'float32 dvalues, float32 x, float32 mean, float32 scale, float32 k1, float32 k2',
+        'float32 dvalues, float32 x, float32 mean, float32 scale, float32 beta, '
+        'float32 k1, float32 k2, float32 alpha',
         'float32 dinputs',
-        'dinputs = dvalues * scale + (x - mean) * k1 + k2;',
+        '''
+        float g = bn_act_grad(bn_z(x, mean, scale, beta), dvalues, alpha);
+        dinputs = g * scale + (x - mean) * k1 + k2;
+        ''',
         'bn_backward_ew',
+        preamble=_ACTIVATION,
     ),
     name='_bn_backward_ew',
 )

aether/layers/normalization.py

@@ -71,6 +71,9 @@ def __init__(self, epsilon=1e-5, momentum=0.9):
         self._forward_finalize_kernel = None
         self._grad_sums_kernel = None
         self._backward_finalize_kernel = None
+        # Leaky slope of an activation the GPU kernels apply after normalizing:
+        # 1 is none (see _fuse_activation).
+        self._alpha = np.float32(1.0)
 
         self.forward = self._forward_fallback
         self.backward = self._backward_fallback
@@ -133,6 +136,8 @@ def build(self, input_shape: tuple[int, ...], seed: int | None = None) -> tuple[
 
     def _compile_for_device(self, device):
         """Triggered by Model.to(device) to map low-level hardware paths."""
+        # Model re-plans activation fusion after every device sync.
+        self._alpha = np.float32(1.0)
         if device == "cupy" and config.HAS_CUPY and gpu_bn.is_gpu_bn_available():
             variant, _ = config.resolve_gpu_launch_geometry()
 
@@ -148,6 +153,15 @@ def _compile_for_device(self, device):
             self.forward = self._forward_fallback
             self.backward = self._backward_fallback
 
+    def _fuse_activation(self, alpha) -> bool:
+        """Apply `out = z > 0 ? z : alpha * z` to the output inside the GPU kernels,
+        so the activation layer that follows can be skipped: 0 is ReLU, else
+        LeakyReLU. Returns False, changing nothing, when the GPU path isn't bound."""
+        if self.forward != self._forward_gpu:
+            return False
+        self._alpha = np.float32(alpha)
+        return True
+
     def _forward_gpu(self, inputs, training):
         xp = config.cp
         # The kernels read raw float* buffers; a precision policy can hand us
@@ -160,7 +174,7 @@ def _forward_gpu(self, inputs, training):
 
         if not training:
             scale = self.gamma / xp.sqrt(self.running_var + self.epsilon)
-            gpu_bn._bn_apply_ew(inputs, self.running_mean, scale, self.beta, output)
+            gpu_bn._bn_apply_ew(inputs, self.running_mean, scale, self.beta, self._alpha, output)
             return output
 
         meta = self._get_shape_meta(inputs.shape)
@@ -170,7 +184,7 @@ def _forward_gpu(self, inputs, training):
         partials = xp.empty(meta["partials_shape"], dtype=xp.float32)
         self._stats_kernel(
             meta["grid_dim"], meta["block_dim"],
-            (inputs_flat, None, None, partials) + meta["static_args"],
+            (inputs_flat, None, None, None, None, np.float32(1.0), partials) + meta["static_args"],
         )
 
         # Rows: mean, var, inv_std, scale. The running statistics update in place.
@@ -186,7 +200,7 @@ def _forward_gpu(self, inputs, training):
         self.inputs = inputs
         self.batch_mean, self.batch_var, self.inv_std, self.scale = stats
 
-        gpu_bn._bn_apply_ew(inputs, self.batch_mean, self.scale, self.beta, output)
+        gpu_bn._bn_apply_ew(inputs, self.batch_mean, self.scale, self.beta, self._alpha, output)
         return output
 
     def _backward_gpu(self, dvalues):
@@ -200,8 +214,8 @@ def _backward_gpu(self, dvalues):
         partials = xp.empty(meta["partials_shape"], dtype=xp.float32)
         self._grad_sums_kernel(
             meta["grid_dim"], meta["block_dim"],
-            (dvalues.reshape(N, C), self.inputs.reshape(N, C), self.batch_mean, partials)
-            + meta["static_args"],
+            (dvalues.reshape(N, C), self.inputs.reshape(N, C), self.batch_mean,
+             self.scale, self.beta, self._alpha, partials) + meta["static_args"],
         )
 
         # Rows: dgamma, dbeta, then the dx coefficients k1 and k2.
@@ -214,7 +228,8 @@ def _backward_gpu(self, dvalues):
 
         self.dinputs = xp.empty_like(dvalues)
         gpu_bn._bn_backward_ew(
-            dvalues, self.inputs, self.batch_mean, self.scale, k1, k2, self.dinputs
+            dvalues, self.inputs, self.batch_mean, self.scale, self.beta, k1, k2,
+            self._alpha, self.dinputs,
         )
         return self.dinputs
 

aether/model.py

@@ -58,6 +58,15 @@
     "weights.safetensors": (zipfile.ZIP_STORED,),
 }
 
+# (normalization, activation) pairs whose activation BatchNorm can apply inside its
+# own GPU kernels, mapped to the leaky slope it takes. Exact types, so subclasses
+# that change the activation never fuse.
+_FUSIBLE_PAIRS = {
+    (layer_module.BatchNorm, layer_module.ReLU): lambda activation: 0.0,
+    (layer_module.BatchNorm, layer_module.LeakyReLU): lambda activation: activation.alpha,
+}
+
+
 def _serialize_component(obj):
     if obj is None or isinstance(obj, (NullOptimizer, NullAccuracy, NullPreprocessor)):
         return None
@@ -284,6 +293,9 @@ def __init__(self):
         self.device = None
         self.precision_policy = None
         self.trainable_layers = []
+        self._fuse_kernels = True
+        self._forward_layers = ()
+        self._backward_layers = ()
         self._seed = None
         self._predict_activation = None
         self._rng_clock = None
@@ -422,6 +434,7 @@ def to(self, device):
         self._sync_device(target_device=target_device)
 
         if self.is_finalized:
+            self._bind_fusion()
             for layer in self.trainable_layers:
                 param_dict = layer.get_parameters()
                 migrated_dict = {}
@@ -464,10 +477,19 @@ def set_precision(self, compute_dtype):
         ):
             self.preprocessor._apply_precision(self.precision_policy)
 
-    def finalize(self, input_shape: tuple[int, ...]):
+    def finalize(self, input_shape: tuple[int, ...], fuse_kernels: bool = True):
         """
         Finalizes the model architecture, propagating tensor dimensions
         and allocating weights/buffers across all registered layers.
+
+        Args:
+            input_shape (tuple[int, ...]): Shape of one sample, without the batch axis.
+            fuse_kernels (bool): On CuPy, a BatchNorm directly followed by ReLU or
+                LeakyReLU applies the activation inside its own kernels and the
+                activation layer is skipped, saving a read and a write of the
+                activations in both directions. Fusion is re-planned on every `to()`,
+                and never applies on NumPy. Not saved: `load()` fuses by default.
+                Defaults to True.
         """
         if self.is_finalized:
             raise RuntimeError("Cannot modify model after finalize() has been called.")
@@ -504,12 +526,11 @@ def finalize(self, input_shape: tuple[int, ...]):
             if getattr(layer, "weights", None) is not None or getattr(layer, "biases", None) is not None
         ]
 
-        # Nothing below the first trainable layer consumes a gradient, so backward stops
-        # there and that layer skips its dX.
-        first = self.layers.index(self.trainable_layers[0]) if self.trainable_layers else 0
-        self._backward_layers = tuple(reversed(self.layers[first:]))
+        # Nothing below the first trainable layer consumes a gradient, so that layer
+        # skips its dX (and backward stops there, see _bind_fusion).
         if self.trainable_layers:
             self.trainable_layers[0].needs_input_grad = False
+        self._fuse_kernels = fuse_kernels
 
         if self.loss is not None:
             self.loss.validate_graph(layers=self.layers)
@@ -545,8 +566,28 @@ def finalize(self, input_shape: tuple[int, ...]):
             self.device = "cupy" if config.xp.__name__ == "cupy" else "numpy"
 
         self._sync_device()
+        self._bind_fusion()
         self.is_finalized = True
 
+    def _bind_fusion(self):
+        """Fuses BatchNorm + activation pairs for the device the layers were just
+        compiled for, and rebuilds the forward/backward layer sequences without the
+        fused activations. Every _compile_for_device resets BatchNorm to no
+        activation, so a pair fused on CuPy runs unfused after `to('numpy')`."""
+        fused = set()
+        if self._fuse_kernels:
+            for norm, activation in zip(self.layers, self.layers[1:]):
+                alpha_of = _FUSIBLE_PAIRS.get((type(norm), type(activation)))
+                if alpha_of is not None and norm._fuse_activation(alpha_of(activation)):
+                    fused.add(id(activation))
+
+        self._forward_layers = tuple(layer for layer in self.layers if id(layer) not in fused)
+        # Backward stops at the first trainable layer.
+        first = self.layers.index(self.trainable_layers[0]) if self.trainable_layers else 0
+        self._backward_layers = tuple(
+            layer for layer in reversed(self.layers[first:]) if id(layer) not in fused
+        )
+
     def _check_ready(self, op_name):
         """Guard shared by train/evaluate/predict: the model must be finalized."""
         if not self.is_finalized:
@@ -658,7 +699,7 @@ def forward(self, X, training=True):
         Sequentially propagate input batch through all layers.
         """
         output = X
-        for layer in self.layers:
+        for layer in self._forward_layers:
             output = layer.forward(output, training=training)
         return output