From 50c1d28365ba0b9e3850c36e78c73af31c1ac225 Mon Sep 17 00:00:00 2001 From: Alexander <208520214+AlexanderSoftCode@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:41:09 -0500 Subject: [PATCH 1/3] perf(batchnorm): size the reduction grid by device _get_shape_meta capped the reductions at 64 block rows, so at C=32 they ran 64 blocks of 256 threads: one per CU on the RX 9070 XT, which has room for 256. The grid now fills every resident block slot once, the same rule as choose_dweight_split_k: slots = mp_count * (max_threads_per_mp // 256) grid_y = max(1, min(ceil(N / 8), slots // grid_x)) 9070 XT (32 WGPs x 2048 threads): 256 blocks at every C. RTX 3050 (16 SMs x 1536 threads): 96 blocks, so C=128 drops from 256. At 256x32x32x32 on the 9070 XT the mean kernel goes from 0.159 to 0.063 ms and grad_sums from 0.153 to 0.078 ms. BatchNorm forward drops 0.19 ms per CIFAR-10 step. #pragma unroll 4 in the N loop made no difference, so it isn't used. --- aether/layers/normalization.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/aether/layers/normalization.py b/aether/layers/normalization.py index 8484bc1..59d4c2c 100644 --- a/aether/layers/normalization.py +++ b/aether/layers/normalization.py @@ -87,9 +87,13 @@ def _get_shape_meta(self, input_shape: tuple[int, ...]): N = int(np.prod(input_shape)) // C block_y = gpu_bn.BLOCK_Y + # Fill every resident block slot once along N: fewer blocks leave the GPU + # idle, more only add partials to reduce. + mp_count, max_threads_per_mp = config.get_device_multiprocessor_limits() + slots = mp_count * max(1, max_threads_per_mp // (32 * block_y)) block_dim = (32, block_y, 1) grid_x = (C + 31) // 32 - grid_y = min(64, (N + block_y - 1) // block_y) + grid_y = max(1, min((N + block_y - 1) // block_y, slots // grid_x)) grid_dim = (grid_x, grid_y, 1) static_args = ( From aebf36d78265cf0534e132a25dfb483d3927f486 Mon Sep 17 00:00:00 2001 From: Alexander <208520214+AlexanderSoftCode@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:41:34 -0500 Subject: [PATCH 2/3] perf(batchnorm): one-pass statistics with a finalize kernel The forward read x three times (mean pass, variance pass, elementwise pass) and launched about 15 kernels: two zero-inits, the divides, the sqrt, and the running-stat EMA. It now launches three: - stats: one pass accumulating d = x - x[0, c] and d*d. Shifting by a sample of the channel keeps E[d^2] - E[d]^2 from cancelling when |mean| >> std, without the per-element divide Welford needs. Each block writes its pair to partials[2, grid_y, C] instead of using atomicAdd, so there is no zero-init and the result is deterministic. grad_sums shares the template and writes partials the same way (they are summed with partials.sum here, until the next commit). - bn_forward_finalize: reduces the partials and writes stats[4, C] = (mean, var, inv_std, gamma * inv_std), then updates running_mean and running_var in place. batch_mean, batch_var, inv_std and the new scale attribute are row views of stats. - _bn_apply_ew: out = fmaf(x - mean, scale, beta), shared by training and inference. GPU inference moves from the folded x*scale + bias to this centered form, which doesn't add rounding error growing with |mean|/std. The CPU fallbacks are unchanged. get_bn_mean_kernel, get_bn_var_kernel, _bn_forward_ew and _bn_inference_ew are gone. _forward_gpu now casts running_mean/running_var to fp32. Training on NumPy under a float64 policy leaves them float64 and Model.to('cupy') keeps that dtype; GPU inference then raised TypeError, and the in-place update would have misread them. BatchNorm forward on the 9070 XT: 1.018 -> 0.399 ms GPU and 0.388 -> 0.148 ms host per CIFAR-10 step. Tests: the large-offset test now checks batch_var against the float64 variance of the same data (it fails on the naive E[x^2] - E[x]^2 and on the old two-pass kernels); the forward/backward parity loops gain a ragged (3, 37, 41, 40) shape with non-trivial gamma/beta; a new test checks two GPU runs are bit-identical; and a regression test covers float64 running statistics. --- aether/custom_kernels/batchnorm_kernel.py | 216 ++++++++++++---------- aether/layers/normalization.py | 112 ++++------- tests/layers/test_batch_normalization.py | 122 +++++++++--- 3 files changed, 255 insertions(+), 195 deletions(-) diff --git a/aether/custom_kernels/batchnorm_kernel.py b/aether/custom_kernels/batchnorm_kernel.py index 745f0e7..69d8d76 100644 --- a/aether/custom_kernels/batchnorm_kernel.py +++ b/aether/custom_kernels/batchnorm_kernel.py @@ -5,155 +5,189 @@ BLOCK_Y = 8 -_STATS_REDUCE_TEMPLATE = Template(r''' +# Shared by every kernel below: sums each thread's (a, b) pair over the block's +# BLOCK_Y rows. Returns true on the threads of row 0 that hold a channel's total. +_PRELUDE = r''' $hip_include #define WAVEFRONT 32 #define BLOCK_Y $block_y +__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]; + + sh_a[threadIdx.y][threadIdx.x] = a; + sh_b[threadIdx.y][threadIdx.x] = b; + __syncthreads(); + + if (threadIdx.y != 0 || c >= C) return false; + a = 0.0f; + b = 0.0f; + for (int j = 0; j < BLOCK_Y; ++j) { + a += sh_a[j][threadIdx.x]; + b += sh_b[j][threadIdx.x]; + } + return true; +} +''' + +# Each block reduces its rows of [N, C] and writes one (a, b) pair per channel to +# partials[2, gridDim.y, C]. No atomics, so the result is deterministic. +_STATS_REDUCE_TEMPLATE = Template(_PRELUDE + r''' extern "C" __global__ void $kernel_name( 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 - float* __restrict__ out_a, // [C], caller must zero-init - float* __restrict__ out_b, // [C], caller must zero-init + float* __restrict__ partials, // [2, gridDim.y, C] const int N, const int C ) { - __shared__ float sh_a[BLOCK_Y][WAVEFRONT]; - __shared__ float sh_b[BLOCK_Y][WAVEFRONT]; - - int c = blockIdx.x * WAVEFRONT + threadIdx.x; - int ty = threadIdx.y; + int c = blockIdx.x * WAVEFRONT + threadIdx.x; float acc_a = 0.0f; float acc_b = 0.0f; if (c < C) { - $mean_load - for (int n = blockIdx.y * BLOCK_Y + ty; n < N; n += gridDim.y * BLOCK_Y) { + $prologue + for (int n = blockIdx.y * BLOCK_Y + threadIdx.y; n < N; n += gridDim.y * BLOCK_Y) { int idx = n * C + c; - float val_p = p[idx]; - $load_q - acc_a += $accum_a_expr; - $accum_b_stmt + $accumulate } } - sh_a[ty][threadIdx.x] = acc_a; - sh_b[ty][threadIdx.x] = acc_b; - __syncthreads(); + if (!block_sum(acc_a, acc_b, c, C)) return; + partials[blockIdx.y * C + c] = acc_a; + partials[(gridDim.y + blockIdx.y) * C + c] = acc_b; +} +''') + +# Reduces the stats partials and finishes the per-channel math in one launch: +# stats[4, C] = (mean, var, inv_std, scale = gamma * inv_std), and the running +# statistics are updated in place. +_FORWARD_FINALIZE_TEMPLATE = Template(_PRELUDE + r''' +extern "C" __global__ +void $kernel_name( + const float* __restrict__ partials, // [2, G, C] sums of d = x - x[0, c] and d * d + const float* __restrict__ x, // [N, C], row 0 is the shift + const float* __restrict__ gamma, + float* __restrict__ stats, // [4, C] + float* __restrict__ running_mean, + float* __restrict__ running_var, + const int G, + const int N, + const int C, + const float eps, + const float momentum +) { + int c = blockIdx.x * WAVEFRONT + threadIdx.x; - if (ty == 0 && c < C) { - float block_a = 0.0f; - float block_b = 0.0f; - for (int j = 0; j < BLOCK_Y; ++j) { - block_a += sh_a[j][threadIdx.x]; - block_b += sh_b[j][threadIdx.x]; + float s1 = 0.0f; + float s2 = 0.0f; + if (c < C) { + for (int g = threadIdx.y; g < G; g += BLOCK_Y) { + s1 += partials[g * C + c]; + s2 += partials[(G + g) * C + c]; } - atomicAdd(&out_a[c], block_a); - $writeback_b } + + if (!block_sum(s1, s2, c, C)) return; + float d_mean = s1 / N; + float mean = x[c] + d_mean; + float var = fmaxf(s2 / N - d_mean * d_mean, 0.0f); + float inv_std = 1.0f / sqrtf(var + eps); + + stats[c] = mean; + stats[C + c] = var; + stats[2 * C + c] = inv_std; + stats[3 * C + c] = gamma[c] * inv_std; + + running_mean[c] = momentum * running_mean[c] + (1.0f - momentum) * mean; + running_var[c] = momentum * running_var[c] + (1.0f - momentum) * var; } ''') _BN_CUDA = dict(hip_include="") _BN_HIP = dict(hip_include="#include \n") -_MEAN_OP = { - "name": "mean", - "mean_load": "", - "load_q": "", - "accum_a_expr": "val_p", - "accum_b_stmt": "", - "writeback_b": "", -} - -_VAR_OP = { - "name": "var", - "mean_load": "float mu = mean[c];", - "load_q": "", - "accum_a_expr": "(val_p - mu) * (val_p - mu)", - "accum_b_stmt": "", - "writeback_b": "", +# Shifting by a sample of the channel keeps the one-pass variance from cancelling +# when |mean| >> std, without the per-element divide Welford needs. +_STATS_OP = { + "name": "stats", + "template": _STATS_REDUCE_TEMPLATE, + "prologue": "float k = p[c];", + "accumulate": "float d = p[idx] - k; acc_a += d; acc_b += d * d;", } _GRAD_SUMS_OP = { "name": "grad_sums", - "mean_load": "float mu = mean[c];", - "load_q": "float val_q = q[idx];", - "accum_a_expr": "val_p", - "accum_b_stmt": "acc_b += val_p * (val_q - mu);", - "writeback_b": "atomicAdd(&out_b[c], block_b);", + "template": _STATS_REDUCE_TEMPLATE, + "prologue": "float mu = mean[c];", + "accumulate": "float g = p[idx]; acc_a += g; acc_b += g * (q[idx] - mu);", } -_bn_reduce_kernel_cache = {} +_FORWARD_FINALIZE_OP = { + "name": "forward_finalize", + "template": _FORWARD_FINALIZE_TEMPLATE, +} +_bn_kernel_cache = {} -def _get_compiled_reduce_kernel(op_dict: dict, variant: str, block_y: int = BLOCK_Y): - """Cached, compiled reduction RawKernel for the given vendor variant. + +def _get_compiled_kernel(op_dict: dict, variant: str, block_y: int = BLOCK_Y): + """Cached, compiled RawKernel for the given vendor variant. Memoized by (op_name, variant) without baking shapes into the compile key. """ key = (op_dict["name"], variant) - if key in _bn_reduce_kernel_cache: - return _bn_reduce_kernel_cache[key] + if key in _bn_kernel_cache: + return _bn_kernel_cache[key] vendor = _BN_HIP if variant == "hip" else _BN_CUDA op_name = op_dict["name"] - kernel_name = f"bn_reduce_{op_name}_{variant}_kernel" - source = _STATS_REDUCE_TEMPLATE.substitute( + kernel_name = f"bn_{op_name}_{variant}_kernel" + source = op_dict["template"].substitute( kernel_name=kernel_name, block_y=block_y, - mean_load=op_dict["mean_load"], - load_q=op_dict["load_q"], - accum_a_expr=op_dict["accum_a_expr"], - accum_b_stmt=op_dict["accum_b_stmt"], - writeback_b=op_dict["writeback_b"], + prologue=op_dict.get("prologue", ""), + accumulate=op_dict.get("accumulate", ""), **vendor, ) kernel = config.build_kernel( lambda: config.cp.RawKernel(source, kernel_name), - name=f"bn_reduce_{op_name}_{variant}", + name=f"bn_{op_name}_{variant}", ) - _bn_reduce_kernel_cache[key] = kernel + _bn_kernel_cache[key] = kernel return kernel -def get_bn_mean_kernel(variant: str): - """Returns a compiled RawKernel for BatchNorm FP32 mean reduction.""" - return _get_compiled_reduce_kernel(_MEAN_OP, variant) +def get_bn_stats_kernel(variant: str): + """Returns a compiled RawKernel for the BatchNorm FP32 shifted sum / sum-of-squares partials.""" + return _get_compiled_kernel(_STATS_OP, variant) -def get_bn_var_kernel(variant: str): - """Returns a compiled RawKernel for BatchNorm FP32 centered variance reduction.""" - return _get_compiled_reduce_kernel(_VAR_OP, variant) +def get_bn_forward_finalize_kernel(variant: str): + """Returns a compiled RawKernel that turns the stats partials into mean, var, inv_std and scale.""" + return _get_compiled_kernel(_FORWARD_FINALIZE_OP, variant) -def get_bn_grad_sums_kernel(variant: str): - """Returns a compiled RawKernel for BatchNorm FP32 gradient sums reduction.""" - return _get_compiled_reduce_kernel(_GRAD_SUMS_OP, variant) -# --- Elementwise Finalize Kernels ----------------------------------------------- +def get_bn_grad_sums_kernel(variant: str): + """Returns a compiled RawKernel for BatchNorm FP32 gradient sums partials.""" + return _get_compiled_kernel(_GRAD_SUMS_OP, variant) -_bn_forward_ew = config.build_kernel( - lambda: config.cp.ElementwiseKernel( - 'float32 x, float32 mean, float32 inv_std, float32 gamma, float32 beta', - 'float32 out', - 'out = (x - mean) * inv_std * gamma + beta;', - 'bn_forward_ew', - ), - name='_bn_forward_ew', -) +# --- Elementwise Kernels -------------------------------------------------------- -_bn_inference_ew = config.build_kernel( +# Training and inference share it: the centered form keeps the rounding error of +# folding (mean, scale, beta) into one bias out of the output. +_bn_apply_ew = config.build_kernel( lambda: config.cp.ElementwiseKernel( - 'float32 x, float32 scale, float32 bias', + 'float32 x, float32 mean, float32 scale, float32 beta', 'float32 out', - 'out = x * scale + bias;', - 'bn_inference_ew', + 'out = fmaf(x - mean, scale, beta);', + 'bn_apply_ew', ), - name='_bn_inference_ew', + name='_bn_apply_ew', ) _bn_backward_ew = config.build_kernel( @@ -174,14 +208,10 @@ def is_gpu_bn_available(variant: str = None) -> bool: if variant is None: variant, _ = config.resolve_gpu_launch_geometry() - mean_k = get_bn_mean_kernel(variant) - var_k = get_bn_var_kernel(variant) - grad_sums_k = get_bn_grad_sums_kernel(variant) return ( - mean_k is not None - and var_k is not None - and grad_sums_k is not None - and _bn_forward_ew is not None - and _bn_inference_ew is not None + get_bn_stats_kernel(variant) is not None + and get_bn_forward_finalize_kernel(variant) is not None + and get_bn_grad_sums_kernel(variant) is not None + and _bn_apply_ew is not None and _bn_backward_ew is not None - ) \ No newline at end of file + ) diff --git a/aether/layers/normalization.py b/aether/layers/normalization.py index 59d4c2c..709c7ae 100644 --- a/aether/layers/normalization.py +++ b/aether/layers/normalization.py @@ -59,6 +59,7 @@ def __init__(self, epsilon=1e-5, momentum=0.9): self.inv_std = None self.batch_mean = None self.batch_var = None + self.scale = None self.weight_regularizer_l1 = 0 self.weight_regularizer_l2 = 0 @@ -66,8 +67,8 @@ def __init__(self, epsilon=1e-5, momentum=0.9): self.bias_regularizer_l2 = 0 self._variant = None - self._mean_kernel = None - self._var_kernel = None + self._stats_kernel = None + self._forward_finalize_kernel = None self._grad_sums_kernel = None self.forward = self._forward_fallback @@ -96,17 +97,16 @@ def _get_shape_meta(self, input_shape: tuple[int, ...]): grid_y = max(1, min((N + block_y - 1) // block_y, slots // grid_x)) grid_dim = (grid_x, grid_y, 1) - static_args = ( - np.int32(N), - np.int32(C), - ) - meta = { "N": N, "C": C, "block_dim": block_dim, "grid_dim": grid_dim, - "static_args": static_args, + # The finalize kernels reduce the grid_y partials, one block per 32 channels. + "finalize_grid_dim": (grid_x, 1, 1), + "partials_shape": (2, grid_y, C), + "static_args": (np.int32(N), np.int32(C)), + "finalize_args": (np.int32(grid_y), np.int32(N), np.int32(C)), } self._launch_cache[cache_key] = meta return meta @@ -135,25 +135,10 @@ def _compile_for_device(self, device): if device == "cupy" and config.HAS_CUPY and gpu_bn.is_gpu_bn_available(): variant, _ = config.resolve_gpu_launch_geometry() - mean_kernel = gpu_bn.get_bn_mean_kernel(variant) - var_kernel = gpu_bn.get_bn_var_kernel(variant) - grad_sums_kernel = gpu_bn.get_bn_grad_sums_kernel(variant) - - if ( - mean_kernel is None - or var_kernel is None - or grad_sums_kernel is None - or gpu_bn._bn_forward_ew is None - or gpu_bn._bn_backward_ew is None - ): - self.forward = self._forward_fallback - self.backward = self._backward_fallback - return - self._variant = variant - self._mean_kernel = mean_kernel - self._var_kernel = var_kernel - self._grad_sums_kernel = grad_sums_kernel + self._stats_kernel = gpu_bn.get_bn_stats_kernel(variant) + self._forward_finalize_kernel = gpu_bn.get_bn_forward_finalize_kernel(variant) + self._grad_sums_kernel = gpu_bn.get_bn_grad_sums_kernel(variant) self.forward = self._forward_gpu self.backward = self._backward_gpu @@ -164,65 +149,42 @@ def _compile_for_device(self, device): def _forward_gpu(self, inputs, training): xp = config.cp # The kernels read raw float* buffers; a precision policy can hand us - # fp16/fp64 from the layer before. + # fp16/fp64 from the layer before, and training on NumPy under a float64 + # policy leaves the running statistics float64. inputs = xp.ascontiguousarray(inputs, dtype=xp.float32) + self.running_mean = xp.ascontiguousarray(self.running_mean, dtype=xp.float32) + self.running_var = xp.ascontiguousarray(self.running_var, dtype=xp.float32) + output = xp.empty_like(inputs) if not training: - inv_std = 1.0 / xp.sqrt(self.running_var + self.epsilon) - scale = self.gamma * inv_std - bias = self.beta - self.running_mean * scale - - output = xp.empty_like(inputs) - gpu_bn._bn_inference_ew(inputs, scale, bias, output) + scale = self.gamma / xp.sqrt(self.running_var + self.epsilon) + gpu_bn._bn_apply_ew(inputs, self.running_mean, scale, self.beta, output) return output meta = self._get_shape_meta(inputs.shape) - N = meta["N"] C = meta["C"] + inputs_flat = inputs.reshape(meta["N"], C) - inputs_flat = inputs.reshape(N, C) - out_mean = xp.zeros((C,), dtype=xp.float32) - - kernel_args = ( - inputs_flat, - None, - None, - out_mean, - None, - ) + meta["static_args"] - self._mean_kernel(meta["grid_dim"], meta["block_dim"], kernel_args) - - batch_mean = out_mean / float(N) - - out_var = xp.zeros((C,), dtype=xp.float32) - kernel_args = ( - inputs_flat, - None, - batch_mean, - out_var, - None, - ) + meta["static_args"] - self._var_kernel(meta["grid_dim"], meta["block_dim"], kernel_args) - - batch_var = out_var / float(N) - inv_std = 1.0 / xp.sqrt(batch_var + self.epsilon) - - # Update running stats - self.running_mean = ( - self.momentum * self.running_mean + (1.0 - self.momentum) * batch_mean + 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"], ) - self.running_var = ( - self.momentum * self.running_var + (1.0 - self.momentum) * batch_var + + # Rows: mean, var, inv_std, scale. The running statistics update in place. + stats = xp.empty((4, C), dtype=xp.float32) + self._forward_finalize_kernel( + meta["finalize_grid_dim"], meta["block_dim"], + (partials, inputs_flat, self.gamma, stats, self.running_mean, self.running_var) + + meta["finalize_args"] + + (np.float32(self.epsilon), np.float32(self.momentum)), ) # Cache intermediate state self.inputs = inputs - self.batch_mean = batch_mean - self.batch_var = batch_var - self.inv_std = inv_std + self.batch_mean, self.batch_var, self.inv_std, self.scale = stats - output = xp.empty_like(inputs) - gpu_bn._bn_forward_ew(inputs, batch_mean, inv_std, self.gamma, self.beta, output) + gpu_bn._bn_apply_ew(inputs, self.batch_mean, self.scale, self.beta, output) return output def _backward_gpu(self, dvalues): @@ -236,17 +198,15 @@ def _backward_gpu(self, dvalues): dvalues_flat = dvalues.reshape(N, C) inputs_flat = self.inputs.reshape(N, C) - dbeta = xp.zeros((C,), dtype=xp.float32) - diff_sum = xp.zeros((C,), dtype=xp.float32) - + partials = xp.empty(meta["partials_shape"], dtype=xp.float32) kernel_args = ( dvalues_flat, inputs_flat, self.batch_mean, - dbeta, - diff_sum, + partials, ) + meta["static_args"] self._grad_sums_kernel(meta["grid_dim"], meta["block_dim"], kernel_args) + dbeta, diff_sum = partials.sum(axis=1) inv_std = self.inv_std gamma = self.gamma diff --git a/tests/layers/test_batch_normalization.py b/tests/layers/test_batch_normalization.py index fde519f..66e138f 100644 --- a/tests/layers/test_batch_normalization.py +++ b/tests/layers/test_batch_normalization.py @@ -10,6 +10,10 @@ class TestBatchNormLayer(base_case.AetherBaseLayerTestCase): NUM_CHANNELS_CONV = 8 CONV_SHAPE = (4, 14, 14, 8) + # N = 4551 rows is not a multiple of the block's 8, C = 40 is not a multiple + # of 32, and the reduction spans many blocks. + RAGGED_SHAPE = (3, 37, 41, 40) + EPSILON = 1e-5 MOMENTUM = 0.9 @@ -181,30 +185,45 @@ def test_gpu_launch_cache_population(self): self.layer_dense.forward(self.dense_inputs, training=True) self.assertEqual(len(self.layer_dense._launch_cache), 1) + def _make_ragged(self): + """A RAGGED_SHAPE layer with non-trivial gamma/beta, so a scale/inv_std mix-up + can't hide behind gamma = 1, plus a matching input batch.""" + C = self.RAGGED_SHAPE[-1] + layer = self.make_built_layer( + BatchNorm, input_shape=self.RAGGED_SHAPE[1:], epsilon=self.EPSILON, momentum=self.MOMENTUM + ) + layer.gamma = self.xp.random.uniform(0.5, 2.0, C).astype(self.xp.float32) + layer.beta = self.xp.random.randn(C).astype(self.xp.float32) + inputs = self.xp.random.randn(*self.RAGGED_SHAPE).astype(self.xp.float32) * 2.0 + 5.0 + return layer, inputs + def test_gpu_forward_training_parity(self): """Verify parity between fallback and GPU forward training implementations.""" if self.backend_name != "cupy": self.skipTest("GPU parity test is only applicable on CuPy backend.") - # Execute fallback path - out_fallback = self.layer_conv._forward_fallback(self.conv_inputs, training=True) - mean_fallback = self.layer_conv.batch_mean.copy() - var_fallback = self.layer_conv.batch_var.copy() - running_mean_fallback = self.layer_conv.running_mean.copy() - running_var_fallback = self.layer_conv.running_var.copy() - - # Reset parameters and running statistics to identical initial state - self.layer_conv.running_mean = self.xp.zeros(self.NUM_CHANNELS_CONV, dtype=self.xp.float32) - self.layer_conv.running_var = self.xp.ones(self.NUM_CHANNELS_CONV, dtype=self.xp.float32) - - # Execute GPU path - out_gpu = self.layer_conv._forward_gpu(self.conv_inputs, training=True) - - self.assertTrue(self.xp.allclose(out_gpu, out_fallback, atol=1e-4, rtol=1e-4)) - self.assertTrue(self.xp.allclose(self.layer_conv.batch_mean, mean_fallback.squeeze(), atol=1e-4, rtol=1e-4)) - self.assertTrue(self.xp.allclose(self.layer_conv.batch_var, var_fallback.squeeze(), atol=1e-4, rtol=1e-4)) - self.assertTrue(self.xp.allclose(self.layer_conv.running_mean, running_mean_fallback, atol=1e-4, rtol=1e-4)) - self.assertTrue(self.xp.allclose(self.layer_conv.running_var, running_var_fallback, atol=1e-4, rtol=1e-4)) + for layer, inputs in [(self.layer_conv, self.conv_inputs), self._make_ragged()]: + C = inputs.shape[-1] + with self.subTest(shape=inputs.shape): + # Execute fallback path + out_fallback = layer._forward_fallback(inputs, training=True) + mean_fallback = layer.batch_mean.copy() + var_fallback = layer.batch_var.copy() + running_mean_fallback = layer.running_mean.copy() + running_var_fallback = layer.running_var.copy() + + # Reset parameters and running statistics to identical initial state + layer.running_mean = self.xp.zeros(C, dtype=self.xp.float32) + layer.running_var = self.xp.ones(C, dtype=self.xp.float32) + + # Execute GPU path + out_gpu = layer._forward_gpu(inputs, training=True) + + self.assertTrue(self.xp.allclose(out_gpu, out_fallback, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.batch_mean, mean_fallback.squeeze(), atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.batch_var, var_fallback.squeeze(), atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.running_mean, running_mean_fallback, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.running_var, running_var_fallback, atol=1e-4, rtol=1e-4)) def test_gpu_forward_inference_parity(self): """Verify parity between fallback and GPU forward inference paths.""" @@ -226,11 +245,12 @@ def test_gpu_backward_parity_2d_and_4d(self): if self.backend_name != "cupy": self.skipTest("GPU parity test is only applicable on CuPy backend.") - for layer, inputs, shape in [ - (self.layer_dense, self.dense_inputs, self.DENSE_SHAPE), - (self.layer_conv, self.conv_inputs, self.CONV_SHAPE), + for layer, inputs in [ + (self.layer_dense, self.dense_inputs), + (self.layer_conv, self.conv_inputs), + self._make_ragged(), ]: - dvalues = self.xp.random.randn(*shape).astype(self.xp.float32) + dvalues = self.xp.random.randn(*inputs.shape).astype(self.xp.float32) # Fallback path run layer._forward_fallback(inputs, training=True) @@ -242,9 +262,55 @@ def test_gpu_backward_parity_2d_and_4d(self): layer._forward_gpu(inputs, training=True) dinputs_gpu = layer._backward_gpu(dvalues) - self.assertTrue(self.xp.allclose(dinputs_gpu, dinputs_fallback, atol=1e-4, rtol=1e-4)) - self.assertTrue(self.xp.allclose(layer.dweights, dweights_fallback, atol=1e-4, rtol=1e-4)) - self.assertTrue(self.xp.allclose(layer.dbiases, dbiases_fallback, atol=1e-4, rtol=1e-4)) + with self.subTest(shape=inputs.shape): + self.assertTrue(self.xp.allclose(dinputs_gpu, dinputs_fallback, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.dweights, dweights_fallback, atol=1e-4, rtol=1e-4)) + self.assertTrue(self.xp.allclose(layer.dbiases, dbiases_fallback, atol=1e-4, rtol=1e-4)) + + def test_gpu_training_is_deterministic(self): + """The reductions write per-block partials instead of float atomics, so two + runs over many blocks give bit-identical statistics, outputs and gradients.""" + if self.backend_name != "cupy": + self.skipTest("GPU determinism test is only applicable on CuPy backend.") + + layer, inputs = self._make_ragged() + dvalues = self.xp.random.randn(*inputs.shape).astype(self.xp.float32) + + runs = [] + for _ in range(2): + out = layer._forward_gpu(inputs, training=True) + stats = [layer.batch_mean.copy(), layer.batch_var.copy()] + dinputs = layer._backward_gpu(dvalues) + runs.append([out.copy(), *stats, dinputs.copy(), layer.dweights.copy(), layer.dbiases.copy()]) + + for name, first, second in zip( + ("output", "batch_mean", "batch_var", "dinputs", "dweights", "dbiases"), *runs + ): + with self.subTest(name=name): + self.assertTrue(self.xp.array_equal(first, second)) + + def test_gpu_float64_running_stats(self): + """Regression: training on NumPy under a float64 policy leaves the running + statistics float64, and Model.to('cupy') keeps that dtype. GPU inference then + raised TypeError, and the in-place running-stat update must not misread them.""" + if self.backend_name != "cupy": + self.skipTest("GPU dtype test is only applicable on CuPy backend.") + + layer = self.layer_dense + mean64 = self.xp.linspace(-1.0, 1.0, self.NUM_FEATURES_DENSE) + var64 = self.xp.linspace(0.5, 2.0, self.NUM_FEATURES_DENSE) + + layer.running_mean, layer.running_var = mean64.copy(), var64.copy() + layer._forward_gpu(self.dense_inputs, training=True) + expected_mean = 0.9 * mean64 + 0.1 * self.xp.mean(self.dense_inputs, axis=0) + expected_var = 0.9 * var64 + 0.1 * self.xp.var(self.dense_inputs, axis=0) + self.assertTrue(self.xp.allclose(layer.running_mean, expected_mean, atol=1e-5, rtol=1e-5)) + self.assertTrue(self.xp.allclose(layer.running_var, expected_var, atol=1e-5, rtol=1e-5)) + + layer.running_mean, layer.running_var = mean64.copy(), var64.copy() + out = layer._forward_gpu(self.dense_inputs, training=False) + expected = (self.dense_inputs - mean64) / self.xp.sqrt(var64 + self.EPSILON) + self.assertTrue(self.xp.allclose(out, expected, atol=1e-5, rtol=1e-5)) def test_gpu_non_float32_inputs_match_float32(self): """Regression: under set_precision() a Dense before BatchNorm emits fp16/fp64, @@ -288,6 +354,10 @@ def test_gpu_numerical_stability_large_offset(self): out = self.layer_conv._forward_gpu(offset_inputs, training=True) + # E[x^2] - E[x]^2 in fp32 cancels to noise here; the reference is the float64 + # variance of the same fp32 data. + expected_var = self.xp.var(offset_inputs.astype(self.xp.float64), axis=(0, 1, 2)) + self.assertTrue(self.xp.allclose(self.layer_conv.batch_var, expected_var, atol=0, rtol=1e-4)) self.assertTrue(self.xp.all(self.layer_conv.batch_var >= 0.0)) self.assertFalse(self.xp.any(self.xp.isnan(out))) self.assertFalse(self.xp.any(self.xp.isinf(out))) From d8838c5a6040025f4d1263a44f894b8cec3b4e38 Mon Sep 17 00:00:00 2001 From: Alexander <208520214+AlexanderSoftCode@users.noreply.github.com> Date: Fri, 30 Oct 2026 02:41:51 -0500 Subject: [PATCH 3/3] perf(batchnorm): fold the backward's per-channel math into a finalize kernel After grad_sums, the backward ran about ten tiny launches for dgamma, dvar, dmu and the scaled coefficients. bn_backward_finalize now reduces the grad_sums partials (sum of dy, sum of dy * (x - mean)) and writes grads[4, C] = (dgamma, dbeta, k1, k2) in one launch, with dx = dy * scale + (x - mean) * k1 + k2 k1 = -scale * inv_std^2 * sum(dy * (x - mean)) / N k2 = -scale * dbeta / N dweights and dbiases are row views of grads. The backward is three launches: grad_sums, finalize, elementwise. The two finalize kernels share the partials reduction (sum_partials) in the kernel prelude. BatchNorm backward on the 9070 XT: 1.067 -> 0.842 ms GPU and 0.358 -> 0.139 ms host per CIFAR-10 step. The elementwise pass is now at the DRAM floor (a plain cp.add with the same reads and writes takes as long). Model.manual_seed's note no longer lists BatchNorm among the kernels that accumulate with atomics; it now names MaxPool2d's backward with overlapping windows (xp.add.at) next to Conv2d. --- aether/custom_kernels/batchnorm_kernel.py | 73 ++++++++++++++++++----- aether/layers/normalization.py | 46 +++++--------- aether/model.py | 5 +- 3 files changed, 77 insertions(+), 47 deletions(-) diff --git a/aether/custom_kernels/batchnorm_kernel.py b/aether/custom_kernels/batchnorm_kernel.py index 69d8d76..0313b06 100644 --- a/aether/custom_kernels/batchnorm_kernel.py +++ b/aether/custom_kernels/batchnorm_kernel.py @@ -5,8 +5,9 @@ BLOCK_Y = 8 -# Shared by every kernel below: sums each thread's (a, b) pair over the block's -# BLOCK_Y rows. Returns true on the threads of row 0 that hold a channel's total. +# 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. _PRELUDE = r''' $hip_include #define WAVEFRONT 32 @@ -29,6 +30,20 @@ } return true; } + +__device__ __forceinline__ bool sum_partials( + const float* __restrict__ partials, int G, int c, int C, float& a, float& b +) { + a = 0.0f; + b = 0.0f; + if (c < C) { + for (int g = threadIdx.y; g < G; g += BLOCK_Y) { + a += partials[g * C + c]; + b += partials[(G + g) * C + c]; + } + } + return block_sum(a, b, c, C); +} ''' # Each block reduces its rows of [N, C] and writes one (a, b) pair per channel to @@ -81,17 +96,9 @@ const float momentum ) { int c = blockIdx.x * WAVEFRONT + threadIdx.x; + float s1, s2; + if (!sum_partials(partials, G, c, C, s1, s2)) return; - float s1 = 0.0f; - float s2 = 0.0f; - if (c < C) { - for (int g = threadIdx.y; g < G; g += BLOCK_Y) { - s1 += partials[g * C + c]; - s2 += partials[(G + g) * C + c]; - } - } - - if (!block_sum(s1, s2, c, C)) return; float d_mean = s1 / N; float mean = x[c] + d_mean; float var = fmaxf(s2 / N - d_mean * d_mean, 0.0f); @@ -107,6 +114,33 @@ } ''') +# Reduces the grad_sums partials into grads[4, C] = (dgamma, dbeta, k1, k2), so +# that dx = dy * scale + (x - mean) * k1 + k2. +_BACKWARD_FINALIZE_TEMPLATE = Template(_PRELUDE + r''' +extern "C" __global__ +void $kernel_name( + const float* __restrict__ partials, // [2, G, C] sums of dy and dy * (x - mean) + const float* __restrict__ inv_std, + const float* __restrict__ scale, + float* __restrict__ grads, // [4, C] + const int G, + const int N, + const int C +) { + int c = blockIdx.x * WAVEFRONT + threadIdx.x; + float s1, s2; + if (!sum_partials(partials, G, c, C, s1, s2)) return; + + float is = inv_std[c]; + float sc = scale[c]; + + grads[c] = is * s2; + grads[C + c] = s1; + grads[2 * C + c] = -sc * is * is * s2 / N; + grads[3 * C + c] = -sc * s1 / N; +} +''') + _BN_CUDA = dict(hip_include="") _BN_HIP = dict(hip_include="#include \n") @@ -131,6 +165,11 @@ "template": _FORWARD_FINALIZE_TEMPLATE, } +_BACKWARD_FINALIZE_OP = { + "name": "backward_finalize", + "template": _BACKWARD_FINALIZE_TEMPLATE, +} + _bn_kernel_cache = {} @@ -176,6 +215,11 @@ def get_bn_grad_sums_kernel(variant: str): """Returns a compiled RawKernel for BatchNorm FP32 gradient sums partials.""" return _get_compiled_kernel(_GRAD_SUMS_OP, variant) + +def get_bn_backward_finalize_kernel(variant: str): + """Returns a compiled RawKernel that turns the gradient partials into dgamma, dbeta and the dx coefficients.""" + return _get_compiled_kernel(_BACKWARD_FINALIZE_OP, variant) + # --- Elementwise Kernels -------------------------------------------------------- # Training and inference share it: the centered form keeps the rounding error of @@ -192,9 +236,9 @@ def get_bn_grad_sums_kernel(variant: str): _bn_backward_ew = config.build_kernel( lambda: config.cp.ElementwiseKernel( - 'float32 dvalues, float32 x, float32 mean, float32 gamma_inv_std, float32 dvar_scaled, float32 dmu_over_n', + 'float32 dvalues, float32 x, float32 mean, float32 scale, float32 k1, float32 k2', 'float32 dinputs', - 'dinputs = dvalues * gamma_inv_std + (x - mean) * dvar_scaled + dmu_over_n;', + 'dinputs = dvalues * scale + (x - mean) * k1 + k2;', 'bn_backward_ew', ), name='_bn_backward_ew', @@ -212,6 +256,7 @@ def is_gpu_bn_available(variant: str = None) -> bool: get_bn_stats_kernel(variant) is not None and get_bn_forward_finalize_kernel(variant) is not None and get_bn_grad_sums_kernel(variant) is not None + and get_bn_backward_finalize_kernel(variant) is not None and _bn_apply_ew is not None and _bn_backward_ew is not None ) diff --git a/aether/layers/normalization.py b/aether/layers/normalization.py index 709c7ae..f67b3fb 100644 --- a/aether/layers/normalization.py +++ b/aether/layers/normalization.py @@ -70,6 +70,7 @@ def __init__(self, epsilon=1e-5, momentum=0.9): self._stats_kernel = None self._forward_finalize_kernel = None self._grad_sums_kernel = None + self._backward_finalize_kernel = None self.forward = self._forward_fallback self.backward = self._backward_fallback @@ -139,6 +140,7 @@ def _compile_for_device(self, device): self._stats_kernel = gpu_bn.get_bn_stats_kernel(variant) self._forward_finalize_kernel = gpu_bn.get_bn_forward_finalize_kernel(variant) self._grad_sums_kernel = gpu_bn.get_bn_grad_sums_kernel(variant) + self._backward_finalize_kernel = gpu_bn.get_bn_backward_finalize_kernel(variant) self.forward = self._forward_gpu self.backward = self._backward_gpu @@ -195,42 +197,24 @@ def _backward_gpu(self, dvalues): N = meta["N"] C = meta["C"] - dvalues_flat = dvalues.reshape(N, C) - inputs_flat = self.inputs.reshape(N, C) - partials = xp.empty(meta["partials_shape"], dtype=xp.float32) - kernel_args = ( - dvalues_flat, - inputs_flat, - self.batch_mean, - partials, - ) + meta["static_args"] - self._grad_sums_kernel(meta["grid_dim"], meta["block_dim"], kernel_args) - dbeta, diff_sum = partials.sum(axis=1) - - inv_std = self.inv_std - gamma = self.gamma - - dgamma = inv_std * diff_sum - dvar = gamma * diff_sum * (-0.5) * (inv_std**3) - dmu = -inv_std * gamma * dbeta - - self.dbiases = dbeta - self.dweights = dgamma + 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"], + ) - gamma_inv_std = gamma * inv_std - dvar_scaled = dvar * (2.0 / N) - dmu_over_n = dmu / float(N) + # Rows: dgamma, dbeta, then the dx coefficients k1 and k2. + grads = xp.empty((4, C), dtype=xp.float32) + self._backward_finalize_kernel( + meta["finalize_grid_dim"], meta["block_dim"], + (partials, self.inv_std, self.scale, grads) + meta["finalize_args"], + ) + self.dweights, self.dbiases, k1, k2 = grads self.dinputs = xp.empty_like(dvalues) gpu_bn._bn_backward_ew( - dvalues, - self.inputs, - self.batch_mean, - gamma_inv_std, - dvar_scaled, - dmu_over_n, - self.dinputs, + dvalues, self.inputs, self.batch_mean, self.scale, k1, k2, self.dinputs ) return self.dinputs diff --git a/aether/model.py b/aether/model.py index ed6ba2a..22dfb75 100644 --- a/aether/model.py +++ b/aether/model.py @@ -312,8 +312,9 @@ def manual_seed(self, seed: int): RuntimeError: If called after `finalize()`. Note: NumPy runs are bit-reproducible. CuPy runs are not bit-identical, - because the BatchNorm and Conv2d gradient kernels accumulate with - float `atomicAdd`, whose ordering varies between runs. + because the Conv2d gradient kernels (and MaxPool2d's backward with + overlapping windows) accumulate with float atomics, whose ordering + varies between runs. """ if self.is_finalized: raise RuntimeError("Cannot set a new seed after finalize() has been called.")