hyperray 0.0.1__py3-none-any.whl

hyperdrone/_vendor/rl-tools/include/rl_tools/nn/layers/conv2d/operations_cuda.h

@@ -0,0 +1,809 @@
+#include "../../../version.h"
+#if (defined(RL_TOOLS_DISABLE_INCLUDE_GUARDS) || !defined(RL_TOOLS_NN_LAYERS_CONV2D_OPERATIONS_CUDA_H)) && (RL_TOOLS_USE_THIS_VERSION == 1)
+#pragma once
+#define RL_TOOLS_NN_LAYERS_CONV2D_OPERATIONS_CUDA_H
+
+#include "../../../devices/cuda.h"
+#include "../../../nn/nn.h"
+#include "../../../mode/mode.h"
+#include "../../../utils/assert/operations_cuda.h"
+#include "layer.h"
+
+#include <cudnn.h>
+
+RL_TOOLS_NAMESPACE_WRAPPER_START
+namespace rl_tools{
+    namespace nn::layers::conv2d::cuda::kernels{
+        template<typename MATH, typename SPEC, typename T>
+        __global__ void activation_forward(T* values){
+            using TI = typename SPEC::TI;
+            using ACC = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+            const TI index = blockIdx.x * blockDim.x + threadIdx.x;
+            if(index < SPEC::INTERNAL_BATCH_SIZE * SPEC::OUTPUT_HEIGHT * SPEC::OUTPUT_WIDTH * SPEC::OUTPUT_CHANNELS){
+                values[index] = static_cast<T>(activation<MATH, ACC, SPEC::ACTIVATION_FUNCTION>(static_cast<ACC>(values[index])));
+            }
+        }
+
+        template<typename MATH, typename SPEC, typename T>
+        __global__ void activation_backward(LayerGradient<SPEC> layer, const T* d_output){
+            using TI = typename SPEC::TI;
+            using ACC = typename SPEC::TYPE_POLICY::template GET<numeric_types::categories::Accumulator>;
+            const TI index = blockIdx.x * blockDim.x + threadIdx.x;
+            if(index < SPEC::INTERNAL_BATCH_SIZE * SPEC::OUTPUT_HEIGHT * SPEC::OUTPUT_WIDTH * SPEC::OUTPUT_CHANNELS){
+                ACC pre_activation = static_cast<ACC>(layer.pre_activations._data[index]);
+                if constexpr(SPEC::NORMALIZATION == Normalization::BATCH_NORM){
+                    const TI channel = index % SPEC::OUTPUT_CHANNELS;
+                    pre_activation = (pre_activation - static_cast<ACC>(layer.norm_cache.mean._data[channel])) * static_cast<ACC>(layer.norm_cache.inv_std._data[channel]);
+                    pre_activation = pre_activation * static_cast<ACC>(layer.norm.gamma.parameters._data[channel]) + static_cast<ACC>(layer.norm.beta.parameters._data[channel]);
+                }
+                layer.output._data[index] = static_cast<T>(d_activation_d_x<MATH, ACC, SPEC::ACTIVATION_FUNCTION>(pre_activation) * static_cast<ACC>(d_output[index]));
+            }
+        }
+
+        template<typename T, typename T_STAT, typename TI>
+        __global__
+        void bn_stats_training(
+            const T* pre_act,
+            T* mean,
+            T* inv_std,
+            T_STAT* running_mean,
+            T_STAT* running_var,
+            float momentum,
+            float eps,
+            TI spatial,
+            TI OC
+        ){
+            TI c = (TI)blockIdx.x;
+            if(c >= OC){
+                return;
+            }
+            __shared__ float s_sum[256];
+            __shared__ float s_sq_sum[256];
+            TI tid = (TI)threadIdx.x;
+            float local_sum = 0;
+            float local_sq_sum = 0;
+            for(TI i = tid; i < spatial; i += (TI)blockDim.x){
+                TI idx = i * OC + c;
+                float v = (float)pre_act[idx];
+                local_sum += v;
+                local_sq_sum += v * v;
+            }
+            s_sum[tid] = local_sum;
+            s_sq_sum[tid] = local_sq_sum;
+            __syncthreads();
+            for(unsigned int s = blockDim.x / 2; s > 0; s >>= 1){
+                if(threadIdx.x < s){
+                    s_sum[threadIdx.x] += s_sum[threadIdx.x + s];
+                    s_sq_sum[threadIdx.x] += s_sq_sum[threadIdx.x + s];
+                }
+                __syncthreads();
+            }
+            if(threadIdx.x == 0){
+                float m = s_sum[0] / (float)spatial;
+                float sq_m = s_sq_sum[0] / (float)spatial;
+                float var = sq_m - m * m;
+                if(var < 0.0f){
+                    var = 0.0f;
+                }
+                mean[c] = (T)m;
+                inv_std[c] = (T)(1.0f / sqrtf(var + eps));
+                float rm = (float)running_mean[c];
+                float rv = (float)running_var[c];
+                running_mean[c] = (T_STAT)((1.0f - momentum) * rm + momentum * m);
+                running_var[c] = (T_STAT)((1.0f - momentum) * rv + momentum * var);
+            }
+        }
+        template<typename T, typename T_STAT, typename TI>
+        __global__
+        void bn_stats_eval(
+            const T_STAT* running_mean,
+            const T_STAT* running_var,
+            T* mean,
+            T* inv_std,
+            float eps,
+            TI OC
+        ){
+            TI i = (TI)blockIdx.x * (TI)blockDim.x + (TI)threadIdx.x;
+            if(i < OC){
+                mean[i] = (T)(float)running_mean[i];
+                inv_std[i] = (T)(1.0f / sqrtf((float)running_var[i] + eps));
+            }
+        }
+        template<typename T, typename T_STAT, typename T_NORM_PARAM, typename TI>
+        __global__
+        void bn_forward_eval_running(
+            const T* pre_act,
+            const T_STAT* running_mean,
+            const T_STAT* running_var,
+            const T_NORM_PARAM* gamma,
+            const T_NORM_PARAM* beta,
+            T* output,
+            float eps,
+            TI total,
+            TI OC
+        ){
+            TI idx = (TI)blockIdx.x * (TI)blockDim.x + (TI)threadIdx.x;
+            if(idx < total){
+                TI c = idx % OC;
+                float inv_std = 1.0f / sqrtf((float)running_var[c] + eps);
+                float z_hat = ((float)pre_act[idx] - (float)running_mean[c]) * inv_std;
+                output[idx] = (T)((float)gamma[c] * z_hat + (float)beta[c]);
+            }
+        }
+        template<typename T, typename T_NORM_PARAM, typename TI>
+        __global__
+        void bn_forward_affine(
+            const T* pre_act,
+            const T* mean,
+            const T* inv_std,
+            const T_NORM_PARAM* gamma,
+            const T_NORM_PARAM* beta,
+            T* output,
+            TI total,
+            TI OC
+        ){
+            TI idx = (TI)blockIdx.x * (TI)blockDim.x + (TI)threadIdx.x;
+            if(idx < total){
+                TI c = idx % OC;
+                float z_hat = ((float)pre_act[idx] - (float)mean[c]) * (float)inv_std[c];
+                output[idx] = (T)((float)gamma[c] * z_hat + (float)beta[c]);
+            }
+        }
+        template<typename T, typename TI, unsigned int N_CHUNK>
+        __global__
+        void bias_backward_partial(
+            const T* __restrict__ d_conv_out,
+            float* __restrict__ partials,
+            TI spatial,
+            TI OC
+        ){
+            TI c = (TI)blockIdx.x;
+            if(c >= OC) return;
+            unsigned int chunk = blockIdx.y;
+            TI chunk_start = (spatial * (TI)chunk) / (TI)N_CHUNK;
+            TI chunk_end = (spatial * (TI)(chunk + 1)) / (TI)N_CHUNK;
+            unsigned int tid = threadIdx.x;
+            float local_sum = 0;
+            for(TI i = chunk_start + (TI)tid; i < chunk_end; i += (TI)blockDim.x){
+                local_sum += (float)d_conv_out[i * OC + c];
+            }
+            for(unsigned int offset = 16; offset > 0; offset >>= 1){
+                local_sum += __shfl_down_sync(0xffffffffu, local_sum, offset);
+            }
+            __shared__ float warp_sums[8];
+            unsigned int lane = tid & 31u;
+            unsigned int warp_id = tid >> 5;
+            if(lane == 0) warp_sums[warp_id] = local_sum;
+            __syncthreads();
+            if(warp_id == 0){
+                unsigned int num_warps = (blockDim.x + 31u) >> 5;
+                float v = (lane < num_warps) ? warp_sums[lane] : 0.0f;
+                for(unsigned int offset = 16; offset > 0; offset >>= 1){
+                    v += __shfl_down_sync(0xffffffffu, v, offset);
+                }
+                if(lane == 0) partials[c * (TI)N_CHUNK + (TI)chunk] = v;
+            }
+        }
+        template<typename T, typename TI, unsigned int N_CHUNK>
+        __global__
+        void bias_backward_finalize(
+            const float* __restrict__ partials,
+            T* __restrict__ d_bias,
+            TI OC
+        ){
+            TI c = (TI)blockIdx.x;
+            if(c >= OC) return;
+            unsigned int tid = threadIdx.x;
+            float v = (tid < N_CHUNK) ? partials[c * (TI)N_CHUNK + (TI)tid] : 0.0f;
+            for(unsigned int offset = N_CHUNK / 2u; offset > 0; offset >>= 1){
+                v += __shfl_down_sync(0xffffffffu, v, offset);
+            }
+            if(tid == 0) d_bias[c] = (T)((float)d_bias[c] + v);
+        }
+        template<typename T, typename T_NORM_PARAM, typename TI>
+        __global__
+        void bn_eval_backward(
+            const T* d_norm_out,
+            const T* pre_act,
+            const T* mean,
+            const T* inv_std,
+            const T_NORM_PARAM* gamma,
+            T* d_conv_out,
+            T* d_gamma,
+            T* d_beta,
+            TI spatial,
+            TI OC
+        ){
+            TI c = (TI)blockIdx.x;
+            if(c >= OC){
+                return;
+            }
+            __shared__ float s_d_gamma[256];
+            __shared__ float s_d_beta[256];
+            TI tid = (TI)threadIdx.x;
+            float local_d_gamma = 0;
+            float local_d_beta = 0;
+            float gamma_c = (float)gamma[c];
+            float mean_c = (float)mean[c];
+            float inv_std_c = (float)inv_std[c];
+            for(TI i = tid; i < spatial; i += (TI)blockDim.x){
+                TI idx = i * OC + c;
+                float z_hat = ((float)pre_act[idx] - mean_c) * inv_std_c;
+                float dno = (float)d_norm_out[idx];
+                d_conv_out[idx] = (T)(dno * gamma_c * inv_std_c);
+                local_d_gamma += dno * z_hat;
+                local_d_beta += dno;
+            }
+            s_d_gamma[tid] = local_d_gamma;
+            s_d_beta[tid] = local_d_beta;
+            __syncthreads();
+            for(unsigned int s = blockDim.x / 2; s > 0; s >>= 1){
+                if(threadIdx.x < s){
+                    s_d_gamma[threadIdx.x] += s_d_gamma[threadIdx.x + s];
+                    s_d_beta[threadIdx.x] += s_d_beta[threadIdx.x + s];
+                }
+                __syncthreads();
+            }
+            if(threadIdx.x == 0){
+                d_gamma[c] += (T)s_d_gamma[0];
+                d_beta[c] += (T)s_d_beta[0];
+            }
+        }
+        template<typename T, typename T_NORM_PARAM>
+        __global__
+        void bn_training_backward(
+            const T* d_norm_out,
+            const T* pre_act,
+            const T* mean,
+            const T* inv_std,
+            const T_NORM_PARAM* gamma,
+            T* d_conv_out,
+            T* d_gamma,
+            T* d_beta,
+            unsigned int spatial,
+            unsigned int OC
+        ){
+            unsigned int c = blockIdx.x;
+            if(c >= OC){
+                return;
+            }
+            __shared__ float s_sum_dz_hat[256];
+            __shared__ float s_sum_dz_hat_z_hat[256];
+            __shared__ float s_d_gamma[256];
+            __shared__ float s_d_beta[256];
+            unsigned int tid = threadIdx.x;
+            float local_sum_dz_hat = 0;
+            float local_sum_dz_hat_z_hat = 0;
+            float local_d_gamma = 0;
+            float local_d_beta = 0;
+            float gamma_c = (float)gamma[c];
+            float mean_c = (float)mean[c];
+            float inv_std_c = (float)inv_std[c];
+            for(unsigned int i = tid; i < spatial; i += blockDim.x){
+                unsigned int idx = i * OC + c;
+                float z_hat = ((float)pre_act[idx] - mean_c) * inv_std_c;
+                float dno = (float)d_norm_out[idx];
+                float d_z_hat = dno * gamma_c;
+                local_sum_dz_hat += d_z_hat;
+                local_sum_dz_hat_z_hat += d_z_hat * z_hat;
+                local_d_gamma += dno * z_hat;
+                local_d_beta += dno;
+            }
+            s_sum_dz_hat[tid] = local_sum_dz_hat;
+            s_sum_dz_hat_z_hat[tid] = local_sum_dz_hat_z_hat;
+            s_d_gamma[tid] = local_d_gamma;
+            s_d_beta[tid] = local_d_beta;
+            __syncthreads();
+            for(unsigned int s = blockDim.x / 2; s > 0; s >>= 1){
+                if(tid < s){
+                    s_sum_dz_hat[tid] += s_sum_dz_hat[tid + s];
+                    s_sum_dz_hat_z_hat[tid] += s_sum_dz_hat_z_hat[tid + s];
+                    s_d_gamma[tid] += s_d_gamma[tid + s];
+                    s_d_beta[tid] += s_d_beta[tid + s];
+                }
+                __syncthreads();
+            }
+            float sum_dz_hat = s_sum_dz_hat[0];
+            float sum_dz_hat_z_hat = s_sum_dz_hat_z_hat[0];
+            float inv_n = 1.0f / (float)spatial;
+            for(unsigned int i = tid; i < spatial; i += blockDim.x){
+                unsigned int idx = i * OC + c;
+                float z_hat = ((float)pre_act[idx] - mean_c) * inv_std_c;
+                float d_z_hat = (float)d_norm_out[idx] * gamma_c;
+                d_conv_out[idx] = (T)(inv_std_c * inv_n * ((float)spatial * d_z_hat - sum_dz_hat - z_hat * sum_dz_hat_z_hat));
+            }
+            if(tid == 0){
+                // One block owns one output channel c, so these writes are uncontended
+                // within this kernel launch and do not require atomics.
+                d_gamma[c] += (T)s_d_gamma[0];
+                d_beta[c] += (T)s_d_beta[0];
+            }
+        }
+    }
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void evaluate(devices::CUDA<DEV_SPEC>& device, const nn::layers::conv2d::LayerForward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::conv2d::Buffer<BUFFER_SPEC>&, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::conv2d::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        // cuDNN reads flat contiguous memory based on the NHWC descriptor and ignores RLtools
+        // strides. A non-contiguous input view would silently produce wrong outputs for
+        // batches > 0. Enforce dense layout at compile time to make this a build error.
+        static_assert(tensor::dense_row_major_layout<INPUT_SPEC>(), "Conv2D CUDA evaluate requires a dense row-major input; copy non-contiguous views into a dense tensor first.");
+        static_assert(tensor::dense_row_major_layout<OUTPUT_SPEC>(), "Conv2D CUDA evaluate requires a dense row-major output.");
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Activation>;
+        using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+        constexpr TI N  = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, IC = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI KH = LAYER_SPEC::KERNEL_HEIGHT, KW = LAYER_SPEC::KERNEL_WIDTH;
+        constexpr TI SH = LAYER_SPEC::STRIDE_H, SW = LAYER_SPEC::STRIDE_W;
+        constexpr TI PH = LAYER_SPEC::PADDING_H, PW = LAYER_SPEC::PADDING_W;
+        constexpr cudnnDataType_t dt = nn::cuda::get_cudnn_dtype<T>();
+        constexpr cudnnDataType_t compute_dt = (dt == CUDNN_DATA_BFLOAT16) ? CUDNN_DATA_FLOAT : dt;
+        constexpr bool HAS_BN = LAYER_SPEC::NORMALIZATION == nn::layers::conv2d::Normalization::BATCH_NORM;
+        constexpr bool HAS_RELU = LAYER_SPEC::ACTIVATION_FUNCTION == nn::activation_functions::ActivationFunction::RELU;
+        constexpr bool FUSE_RELU = HAS_RELU && !HAS_BN;
+
+        static cudnnTensorDescriptor_t xd = nullptr, yd = nullptr, bd = nullptr;
+        static cudnnFilterDescriptor_t wd = nullptr;
+        static cudnnConvolutionDescriptor_t cd = nullptr;
+        static cudnnActivationDescriptor_t fused_ad = nullptr, relu_ad = nullptr;
+        static cudnnConvolutionFwdAlgo_t cached_algo;
+        static size_t cached_ws = 0;
+        // magic static: the one-time descriptor/algorithm setup is guarded by the compiler so concurrent first calls from several host threads (one per GPU) are safe
+        static const bool initialized = [&](){
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&xd), "cudnnCreateTensorDescriptor conv2d_eval.xd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(xd, CUDNN_TENSOR_NHWC, dt, N, IC, IH, IW), "cudnnSetTensor4dDescriptor conv2d_eval.xd");
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&yd), "cudnnCreateTensorDescriptor conv2d_eval.yd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(yd, CUDNN_TENSOR_NHWC, dt, N, OC, OH, OW), "cudnnSetTensor4dDescriptor conv2d_eval.yd");
+            check_cudnn_call(device, cudnnCreateFilterDescriptor(&wd), "cudnnCreateFilterDescriptor conv2d_eval.wd");
+            check_cudnn_call(device, cudnnSetFilter4dDescriptor(wd, dt, CUDNN_TENSOR_NHWC, OC, IC, KH, KW), "cudnnSetFilter4dDescriptor conv2d_eval.wd");
+            check_cudnn_call(device, cudnnCreateConvolutionDescriptor(&cd), "cudnnCreateConvolutionDescriptor conv2d_eval.cd");
+            check_cudnn_call(device, cudnnSetConvolution2dDescriptor(cd, PH, PW, SH, SW, 1, 1, CUDNN_CROSS_CORRELATION, compute_dt), "cudnnSetConvolution2dDescriptor conv2d_eval.cd");
+            check_cudnn_call(device, cudnnSetConvolutionMathType(cd, CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION), "cudnnSetConvolutionMathType conv2d_eval.cd");
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&bd), "cudnnCreateTensorDescriptor conv2d_eval.bd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(bd, CUDNN_TENSOR_NHWC, dt, 1, OC, 1, 1), "cudnnSetTensor4dDescriptor conv2d_eval.bd");
+            check_cudnn_call(device, cudnnCreateActivationDescriptor(&fused_ad), "cudnnCreateActivationDescriptor conv2d_eval.fused_ad");
+            if constexpr(FUSE_RELU){
+                check_cudnn_call(device, cudnnSetActivationDescriptor(fused_ad, CUDNN_ACTIVATION_RELU, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_eval.fused_ad");
+                cached_algo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
+            } else {
+                check_cudnn_call(device, cudnnSetActivationDescriptor(fused_ad, CUDNN_ACTIVATION_IDENTITY, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_eval.fused_ad");
+                int ac; cudnnConvolutionFwdAlgoPerf_t ap;
+                check_cudnn_call(device, cudnnGetConvolutionForwardAlgorithm_v7(device.cudnn_handle, xd, wd, cd, yd, 1, &ac, &ap), "cudnnGetConvolutionForwardAlgorithm_v7 conv2d_eval");
+                cached_algo = ap.algo;
+            }
+            check_cudnn_call(device, cudnnGetConvolutionForwardWorkspaceSize(device.cudnn_handle, xd, wd, cd, yd, cached_algo, &cached_ws), "cudnnGetConvolutionForwardWorkspaceSize conv2d_eval");
+            if constexpr(HAS_BN && HAS_RELU){
+                check_cudnn_call(device, cudnnCreateActivationDescriptor(&relu_ad), "cudnnCreateActivationDescriptor conv2d_eval.relu_ad");
+                check_cudnn_call(device, cudnnSetActivationDescriptor(relu_ad, CUDNN_ACTIVATION_RELU, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_eval.relu_ad");
+            }
+            return true;
+        }();
+        (void)initialized;
+        if(cached_ws > 0){
+            utils::assert_exit(device, device.dynamic_memory_allocation_allowed || cached_ws <= device.cudnn_workspace_size, "Dynamic CUDA memory allocations are disabled");
+            ensure_cudnn_workspace(device, cached_ws);
+        }
+        if constexpr(HAS_BN){
+            // Keep BN paths on explicit conv+bias to avoid fragile fused-call behavior.
+            float a = 1, b = 0;
+            cudnnStatus_t stat = cudnnConvolutionForward(device.cudnn_handle, &a, xd, input._data, wd, layer.weights.parameters._data,
+                cd, cached_algo, device.cudnn_workspace, device.cudnn_workspace_size, &b, yd, output._data);
+            check_cudnn_call(device, stat, "cudnnConvolutionForward conv2d_eval");
+            float ba = 1, bb = 1;
+            stat = cudnnAddTensor(device.cudnn_handle, &ba, bd, layer.biases.parameters._data, &bb, yd, output._data);
+            check_cudnn_call(device, stat, "cudnnAddTensor conv2d_eval");
+        }
+        else{
+            float a1 = 1, a2 = 0;
+            check_cudnn_call(device, cudnnConvolutionBiasActivationForward(device.cudnn_handle, &a1, xd, input._data, wd, layer.weights.parameters._data,
+                cd, cached_algo, device.cudnn_workspace, device.cudnn_workspace_size,
+                &a2, yd, output._data, bd, layer.biases.parameters._data, fused_ad, yd, output._data), "cudnnConvolutionBiasActivationForward conv2d_eval");
+        }
+        if constexpr(HAS_BN){
+            constexpr TI BN_ELEMENT_BLOCK = 256;
+            constexpr TI TOTAL = N * OH * OW * OC;
+            constexpr TI N_BLOCKS_ELEMENT = RL_TOOLS_DEVICES_CUDA_CEIL(TOTAL, BN_ELEMENT_BLOCK);
+            using T_STAT = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormStatistics>;
+            using T_NORM_PARAM = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormParameter>;
+            nn::layers::conv2d::cuda::kernels::bn_forward_eval_running<T, T_STAT, T_NORM_PARAM, TI><<<N_BLOCKS_ELEMENT, BN_ELEMENT_BLOCK, 0, device.stream>>>(
+                output._data,
+                layer.norm.running_mean._data,
+                layer.norm.running_var._data,
+                layer.norm.gamma.parameters._data,
+                layer.norm.beta.parameters._data,
+                output._data,
+                (float)LAYER_SPEC::NORM_EPSILON,
+                TOTAL,
+                OC
+            );
+            if constexpr(HAS_RELU){
+                float ra = 1, rb = 0;
+                check_cudnn_call(device, cudnnActivationForward(device.cudnn_handle, relu_ad, &ra, yd, output._data, &rb, yd, output._data), "cudnnActivationForward conv2d_eval");
+            }
+        }
+        if constexpr(!HAS_RELU && LAYER_SPEC::ACTIVATION_FUNCTION != nn::activation_functions::IDENTITY){
+            constexpr TI BLOCK = 256;
+            constexpr TI GRID = (N * OH * OW * OC + BLOCK - 1) / BLOCK;
+            nn::layers::conv2d::cuda::kernels::activation_forward<typename DEV_SPEC::MATH, LAYER_SPEC><<<GRID, BLOCK, 0, device.stream>>>(output._data);
+        }
+        check_status(device);
+    }
+
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename OUTPUT_SPEC, typename BUFFER_SPEC, typename RNG, typename MODE = mode::Default<>>
+    void forward(devices::CUDA<DEV_SPEC>& device, nn::layers::conv2d::LayerBackward<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<OUTPUT_SPEC>& output, nn::layers::conv2d::Buffer<BUFFER_SPEC>& buffer, RNG& rng, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        static_assert(nn::layers::conv2d::check_input_output<LAYER_SPEC, INPUT_SPEC, OUTPUT_SPEC>);
+        static_assert(tensor::dense_row_major_layout<INPUT_SPEC>(), "Conv2D CUDA forward requires a dense row-major input; copy non-contiguous views into a dense tensor first.");
+        static_assert(tensor::dense_row_major_layout<OUTPUT_SPEC>(), "Conv2D CUDA forward requires a dense row-major output.");
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Activation>;
+        using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, IC = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI KH = LAYER_SPEC::KERNEL_HEIGHT, KW = LAYER_SPEC::KERNEL_WIDTH;
+        constexpr TI SH = LAYER_SPEC::STRIDE_H, SW = LAYER_SPEC::STRIDE_W;
+        constexpr TI PH = LAYER_SPEC::PADDING_H, PW = LAYER_SPEC::PADDING_W;
+        constexpr cudnnDataType_t dt = nn::cuda::get_cudnn_dtype<T>();
+        constexpr cudnnDataType_t compute_dt = (dt == CUDNN_DATA_BFLOAT16) ? CUDNN_DATA_FLOAT : dt;
+        constexpr bool HAS_BN = LAYER_SPEC::NORMALIZATION == nn::layers::conv2d::Normalization::BATCH_NORM;
+        constexpr bool HAS_RELU = LAYER_SPEC::ACTIVATION_FUNCTION == nn::activation_functions::ActivationFunction::RELU;
+        constexpr bool FUSE_RELU = HAS_RELU && !HAS_BN;
+
+        static cudnnTensorDescriptor_t xd = nullptr, yd = nullptr, bd = nullptr;
+        static cudnnFilterDescriptor_t wd = nullptr;
+        static cudnnConvolutionDescriptor_t cd = nullptr;
+        static cudnnActivationDescriptor_t fused_ad = nullptr, relu_ad = nullptr;
+        static cudnnConvolutionFwdAlgo_t cached_algo;
+        static size_t cached_ws = 0;
+        // magic static: the one-time descriptor/algorithm setup is guarded by the compiler so concurrent first calls from several host threads (one per GPU) are safe
+        static const bool initialized = [&](){
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&xd), "cudnnCreateTensorDescriptor conv2d_fwd.xd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(xd, CUDNN_TENSOR_NHWC, dt, N, IC, IH, IW), "cudnnSetTensor4dDescriptor conv2d_fwd.xd");
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&yd), "cudnnCreateTensorDescriptor conv2d_fwd.yd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(yd, CUDNN_TENSOR_NHWC, dt, N, OC, OH, OW), "cudnnSetTensor4dDescriptor conv2d_fwd.yd");
+            check_cudnn_call(device, cudnnCreateFilterDescriptor(&wd), "cudnnCreateFilterDescriptor conv2d_fwd.wd");
+            check_cudnn_call(device, cudnnSetFilter4dDescriptor(wd, dt, CUDNN_TENSOR_NHWC, OC, IC, KH, KW), "cudnnSetFilter4dDescriptor conv2d_fwd.wd");
+            check_cudnn_call(device, cudnnCreateConvolutionDescriptor(&cd), "cudnnCreateConvolutionDescriptor conv2d_fwd.cd");
+            check_cudnn_call(device, cudnnSetConvolution2dDescriptor(cd, PH, PW, SH, SW, 1, 1, CUDNN_CROSS_CORRELATION, compute_dt), "cudnnSetConvolution2dDescriptor conv2d_fwd.cd");
+            check_cudnn_call(device, cudnnSetConvolutionMathType(cd, CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION), "cudnnSetConvolutionMathType conv2d_fwd.cd");
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&bd), "cudnnCreateTensorDescriptor conv2d_fwd.bd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(bd, CUDNN_TENSOR_NHWC, dt, 1, OC, 1, 1), "cudnnSetTensor4dDescriptor conv2d_fwd.bd");
+            check_cudnn_call(device, cudnnCreateActivationDescriptor(&fused_ad), "cudnnCreateActivationDescriptor conv2d_fwd.fused_ad");
+            if constexpr(FUSE_RELU){
+                check_cudnn_call(device, cudnnSetActivationDescriptor(fused_ad, CUDNN_ACTIVATION_RELU, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_fwd.fused_ad");
+                cached_algo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;
+            } else {
+                check_cudnn_call(device, cudnnSetActivationDescriptor(fused_ad, CUDNN_ACTIVATION_IDENTITY, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_fwd.fused_ad");
+                int ac; cudnnConvolutionFwdAlgoPerf_t ap;
+                check_cudnn_call(device, cudnnGetConvolutionForwardAlgorithm_v7(device.cudnn_handle, xd, wd, cd, yd, 1, &ac, &ap), "cudnnGetConvolutionForwardAlgorithm_v7 conv2d_fwd");
+                cached_algo = ap.algo;
+            }
+            check_cudnn_call(device, cudnnGetConvolutionForwardWorkspaceSize(device.cudnn_handle, xd, wd, cd, yd, cached_algo, &cached_ws), "cudnnGetConvolutionForwardWorkspaceSize conv2d_fwd");
+            if constexpr(HAS_BN && HAS_RELU){
+                check_cudnn_call(device, cudnnCreateActivationDescriptor(&relu_ad), "cudnnCreateActivationDescriptor conv2d_fwd.relu_ad");
+                check_cudnn_call(device, cudnnSetActivationDescriptor(relu_ad, CUDNN_ACTIVATION_RELU, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_fwd.relu_ad");
+            }
+            return true;
+        }();
+        (void)initialized;
+        if(cached_ws > 0){
+            utils::assert_exit(device, device.dynamic_memory_allocation_allowed || cached_ws <= device.cudnn_workspace_size, "Dynamic CUDA memory allocations are disabled");
+            ensure_cudnn_workspace(device, cached_ws);
+        }
+        if constexpr(HAS_BN){
+            constexpr TI BN_CHANNEL_BLOCK = 256;
+            constexpr TI BN_ELEMENT_BLOCK = 256;
+            constexpr TI SPATIAL = N * OH * OW;
+            constexpr TI TOTAL = SPATIAL * OC;
+            { float a = 1, b = 0;
+              cudnnStatus_t stat = cudnnConvolutionForward(device.cudnn_handle, &a, xd, input._data, wd, layer.weights.parameters._data,
+                  cd, cached_algo, device.cudnn_workspace, device.cudnn_workspace_size, &b, yd, layer.pre_activations._data);
+              check_cudnn_call(device, stat, "cudnnConvolutionForward conv2d_fwd");
+              float ba = 1, bb = 1;
+              stat = cudnnAddTensor(device.cudnn_handle, &ba, bd, layer.biases.parameters._data, &bb, yd, layer.pre_activations._data);
+              check_cudnn_call(device, stat, "cudnnAddTensor conv2d_fwd");
+            }
+            if constexpr(mode::is<MODE, mode::Evaluation>){
+                constexpr TI N_BLOCKS_CHANNEL = RL_TOOLS_DEVICES_CUDA_CEIL(OC, BN_CHANNEL_BLOCK);
+                constexpr TI N_BLOCKS_ELEMENT = RL_TOOLS_DEVICES_CUDA_CEIL(TOTAL, BN_ELEMENT_BLOCK);
+                using T_STAT = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormStatistics>;
+                using T_NORM_PARAM = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormParameter>;
+                nn::layers::conv2d::cuda::kernels::bn_stats_eval<T, T_STAT, TI><<<N_BLOCKS_CHANNEL, BN_CHANNEL_BLOCK, 0, device.stream>>>(
+                    layer.norm.running_mean._data,
+                    layer.norm.running_var._data,
+                    layer.norm_cache.mean._data,
+                    layer.norm_cache.inv_std._data,
+                    (float)LAYER_SPEC::NORM_EPSILON,
+                    OC
+                );
+                nn::layers::conv2d::cuda::kernels::bn_forward_affine<T, T_NORM_PARAM, TI><<<N_BLOCKS_ELEMENT, BN_ELEMENT_BLOCK, 0, device.stream>>>(
+                    layer.pre_activations._data,
+                    layer.norm_cache.mean._data,
+                    layer.norm_cache.inv_std._data,
+                    layer.norm.gamma.parameters._data,
+                    layer.norm.beta.parameters._data,
+                    output._data,
+                    TOTAL,
+                    OC
+                );
+            } else {
+                using T_STAT = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormStatistics>;
+                nn::layers::conv2d::cuda::kernels::bn_stats_training<T, T_STAT, TI><<<OC, BN_CHANNEL_BLOCK, 0, device.stream>>>(
+                    layer.pre_activations._data,
+                    layer.norm_cache.mean._data,
+                    layer.norm_cache.inv_std._data,
+                    layer.norm.running_mean._data,
+                    layer.norm.running_var._data,
+                    (float)LAYER_SPEC::BN_MOMENTUM,
+                    (float)LAYER_SPEC::NORM_EPSILON,
+                    SPATIAL,
+                    OC
+                );
+                constexpr TI N_BLOCKS_ELEMENT = RL_TOOLS_DEVICES_CUDA_CEIL(TOTAL, BN_ELEMENT_BLOCK);
+                using T_NORM_PARAM_TR = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormParameter>;
+                nn::layers::conv2d::cuda::kernels::bn_forward_affine<T, T_NORM_PARAM_TR, TI><<<N_BLOCKS_ELEMENT, BN_ELEMENT_BLOCK, 0, device.stream>>>(
+                    layer.pre_activations._data,
+                    layer.norm_cache.mean._data,
+                    layer.norm_cache.inv_std._data,
+                    layer.norm.gamma.parameters._data,
+                    layer.norm.beta.parameters._data,
+                    output._data,
+                    TOTAL,
+                    OC
+                );
+            }
+            if constexpr(HAS_RELU){
+                float a = 1, b = 0;
+                check_cudnn_call(device, cudnnActivationForward(device.cudnn_handle, relu_ad, &a, yd, output._data, &b, yd, output._data), "cudnnActivationForward conv2d_fwd");
+            }
+        } else {
+            T* dst = FUSE_RELU ? output._data : layer.pre_activations._data;
+            { float a1 = 1, a2 = 0;
+              check_cudnn_call(device, cudnnConvolutionBiasActivationForward(device.cudnn_handle, &a1, xd, input._data, wd, layer.weights.parameters._data,
+                  cd, cached_algo, device.cudnn_workspace, device.cudnn_workspace_size,
+                  &a2, yd, dst, bd, layer.biases.parameters._data, fused_ad, yd, dst), "cudnnConvolutionBiasActivationForward conv2d_fwd"); }
+            if constexpr(!FUSE_RELU){
+                check_cuda_call(device, cudaMemcpyAsync(output._data, layer.pre_activations._data, N*OH*OW*OC*sizeof(T), cudaMemcpyDeviceToDevice, device.stream), "cudaMemcpyAsync conv2d_fwd preact_to_output");
+            }
+        }
+        if constexpr(!HAS_RELU && LAYER_SPEC::ACTIVATION_FUNCTION != nn::activation_functions::IDENTITY){
+            constexpr TI BLOCK = 256;
+            constexpr TI GRID = (N * OH * OW * OC + BLOCK - 1) / BLOCK;
+            nn::layers::conv2d::cuda::kernels::activation_forward<typename DEV_SPEC::MATH, LAYER_SPEC><<<GRID, BLOCK, 0, device.stream>>>(output._data);
+        }
+        check_status(device);
+    }
+
+    template<bool COMPUTE_INPUT_GRADIENT, typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    void backward_impl(devices::CUDA<DEV_SPEC>& device, nn::layers::conv2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::conv2d::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Gradient>;
+        using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+        (void)buffer;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, IC = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI KH = LAYER_SPEC::KERNEL_HEIGHT, KW = LAYER_SPEC::KERNEL_WIDTH;
+        constexpr TI SH = LAYER_SPEC::STRIDE_H, SW = LAYER_SPEC::STRIDE_W;
+        constexpr TI PH = LAYER_SPEC::PADDING_H, PW = LAYER_SPEC::PADDING_W;
+        constexpr cudnnDataType_t dt = nn::cuda::get_cudnn_dtype<T>();
+        constexpr cudnnDataType_t compute_dt = (dt == CUDNN_DATA_BFLOAT16) ? CUDNN_DATA_FLOAT : dt;
+
+        static cudnnTensorDescriptor_t xd = nullptr, yd = nullptr;
+        static cudnnConvolutionDescriptor_t cd = nullptr;
+        static cudnnActivationDescriptor_t relu_ad = nullptr;
+        static cudnnFilterDescriptor_t wd = nullptr;
+        static cudnnConvolutionBwdFilterAlgo_t cached_bf_algo;
+        static cudnnConvolutionBwdDataAlgo_t cached_bd_algo;
+        static size_t cached_bf_ws = 0, cached_bd_ws = 0;
+        static bool bf_ok = false, bd_ok = false;
+        // magic static: the one-time descriptor/algorithm setup is guarded by the compiler so concurrent first calls from several host threads (one per GPU) are safe
+        static const bool initialized = [&](){
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&xd), "cudnnCreateTensorDescriptor conv2d_bwd.xd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(xd, CUDNN_TENSOR_NHWC, dt, N, IC, IH, IW), "cudnnSetTensor4dDescriptor conv2d_bwd.xd");
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&yd), "cudnnCreateTensorDescriptor conv2d_bwd.yd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(yd, CUDNN_TENSOR_NHWC, dt, N, OC, OH, OW), "cudnnSetTensor4dDescriptor conv2d_bwd.yd");
+            check_cudnn_call(device, cudnnCreateConvolutionDescriptor(&cd), "cudnnCreateConvolutionDescriptor conv2d_bwd.cd");
+            check_cudnn_call(device, cudnnSetConvolution2dDescriptor(cd, PH, PW, SH, SW, 1, 1, CUDNN_CROSS_CORRELATION, compute_dt), "cudnnSetConvolution2dDescriptor conv2d_bwd.cd");
+            check_cudnn_call(device, cudnnSetConvolutionMathType(cd, CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION), "cudnnSetConvolutionMathType conv2d_bwd.cd");
+            if constexpr(LAYER_SPEC::ACTIVATION_FUNCTION == nn::activation_functions::ActivationFunction::RELU){
+                check_cudnn_call(device, cudnnCreateActivationDescriptor(&relu_ad), "cudnnCreateActivationDescriptor conv2d_bwd.relu_ad");
+                check_cudnn_call(device, cudnnSetActivationDescriptor(relu_ad, CUDNN_ACTIVATION_RELU, CUDNN_NOT_PROPAGATE_NAN, 0.0), "cudnnSetActivationDescriptor conv2d_bwd.relu_ad");
+            }
+            check_cudnn_call(device, cudnnCreateFilterDescriptor(&wd), "cudnnCreateFilterDescriptor conv2d_bwd.wd");
+            check_cudnn_call(device, cudnnSetFilter4dDescriptor(wd, dt, CUDNN_TENSOR_NHWC, OC, IC, KH, KW), "cudnnSetFilter4dDescriptor conv2d_bwd.wd");
+            {
+                constexpr int MA = 8;
+                int ac;
+                cudnnConvolutionBwdFilterAlgoPerf_t ap[MA];
+                check_cudnn_call(device, cudnnGetConvolutionBackwardFilterAlgorithm_v7(device.cudnn_handle, xd, yd, cd, wd, MA, &ac, ap), "cudnnGetConvolutionBackwardFilterAlgorithm_v7 conv2d_bwd");
+                for(int i = 0; i < ac; i++){
+                    if(ap[i].status == CUDNN_STATUS_SUCCESS && ap[i].determinism == CUDNN_DETERMINISTIC){
+                        cached_bf_algo = ap[i].algo;
+                        bf_ok = true;
+                        break;
+                    }
+                }
+                utils::assert_exit(device, bf_ok, "No deterministic cuDNN convolution filter-gradient algorithm available");
+                check_cudnn_call(device, cudnnGetConvolutionBackwardFilterWorkspaceSize(device.cudnn_handle, xd, yd, cd, wd, cached_bf_algo, &cached_bf_ws), "cudnnGetConvolutionBackwardFilterWorkspaceSize conv2d_bwd");
+            }
+            if constexpr(COMPUTE_INPUT_GRADIENT){
+                constexpr int MA = 8; int ac; cudnnConvolutionBwdDataAlgoPerf_t ap[MA];
+                check_cudnn_call(device, cudnnGetConvolutionBackwardDataAlgorithm_v7(device.cudnn_handle, wd, yd, cd, xd, MA, &ac, ap), "cudnnGetConvolutionBackwardDataAlgorithm_v7 conv2d_bwd");
+                for(int i = 0; i < ac; i++){
+                    if(ap[i].status == CUDNN_STATUS_SUCCESS && ap[i].determinism == CUDNN_DETERMINISTIC){
+                        cached_bd_algo = ap[i].algo;
+                        bd_ok = true;
+                        break;
+                    }
+                }
+                utils::assert_exit(device, bd_ok, "No deterministic cuDNN convolution input-gradient algorithm available");
+                check_cudnn_call(device, cudnnGetConvolutionBackwardDataWorkspaceSize(device.cudnn_handle, wd, yd, cd, xd, cached_bd_algo, &cached_bd_ws), "cudnnGetConvolutionBackwardDataWorkspaceSize conv2d_bwd");
+            }
+            return true;
+        }();
+        (void)initialized;
+
+        // Using layer.output as scratch for the pre-activation gradient preserves d_output for
+        // callers that still need the upstream gradient (e.g. resnet_block's shortcut path reuses
+        // buffer.shortcut as d_output for both conv2 and the downsample conv).
+        T* d_conv_out = layer.output._data;
+        if constexpr(LAYER_SPEC::ACTIVATION_FUNCTION == nn::activation_functions::ActivationFunction::RELU){
+            float a = 1, b = 0;
+            check_cudnn_call(device, cudnnActivationBackward(device.cudnn_handle, relu_ad, &a, yd, layer.output._data, yd, d_output._data, yd, layer.output._data, &b, yd, d_conv_out), "cudnnActivationBackward conv2d_bwd");
+        }
+        else if constexpr(LAYER_SPEC::ACTIVATION_FUNCTION != nn::activation_functions::IDENTITY){
+            constexpr TI BLOCK = 256;
+            constexpr TI GRID = (N * OH * OW * OC + BLOCK - 1) / BLOCK;
+            nn::layers::conv2d::cuda::kernels::activation_backward<typename DEV_SPEC::MATH><<<GRID, BLOCK, 0, device.stream>>>(layer, d_output._data);
+        }
+        if constexpr(LAYER_SPEC::NORMALIZATION == nn::layers::conv2d::Normalization::BATCH_NORM){
+            constexpr bool IS_EVAL = mode::is<MODE, mode::Evaluation>;
+            const T* d_norm_out_in = (LAYER_SPEC::ACTIVATION_FUNCTION != nn::activation_functions::IDENTITY) ? d_conv_out : d_output._data;
+            if constexpr(IS_EVAL){
+                constexpr TI SPATIAL = N * OH * OW;
+                constexpr TI BN_BS = 256;
+                using T_NORM_PARAM_BW = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormParameter>;
+                nn::layers::conv2d::cuda::kernels::bn_eval_backward<T, T_NORM_PARAM_BW, TI><<<OC, BN_BS, 0, device.stream>>>(
+                    d_norm_out_in, layer.pre_activations._data,
+                    layer.norm_cache.mean._data, layer.norm_cache.inv_std._data,
+                    layer.norm.gamma.parameters._data,
+                    d_conv_out,
+                    layer.norm.gamma.gradient._data, layer.norm.beta.gradient._data,
+                    SPATIAL, OC);
+            } else {
+                constexpr TI SPATIAL = N * OH * OW;
+                constexpr TI BN_BS = 256;
+                using T_NORM_PARAM_BW = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::NormParameter>;
+                nn::layers::conv2d::cuda::kernels::bn_training_backward<T, T_NORM_PARAM_BW><<<OC, BN_BS, 0, device.stream>>>(
+                    d_norm_out_in, layer.pre_activations._data,
+                    layer.norm_cache.mean._data, layer.norm_cache.inv_std._data,
+                    layer.norm.gamma.parameters._data,
+                    d_conv_out,
+                    layer.norm.gamma.gradient._data, layer.norm.beta.gradient._data,
+                    SPATIAL, OC);
+            }
+        } else if constexpr(LAYER_SPEC::ACTIVATION_FUNCTION == nn::activation_functions::IDENTITY){
+            check_cuda_call(device, cudaMemcpyAsync(d_conv_out, d_output._data, N*OH*OW*OC*sizeof(T), cudaMemcpyDeviceToDevice, device.stream), "cudaMemcpyAsync conv2d_bwd d_output_to_d_conv_out");
+        }
+        {
+            constexpr TI SPATIAL = N * OH * OW;
+            constexpr unsigned int N_CHUNK = 32;
+            constexpr TI BN_BS = 256;
+            // Phase 1: partial sums over N_CHUNK contiguous spatial chunks per channel.
+            // Phase 2: reduce N_CHUNK partials within one warp and accumulate into d_bias.
+            // Use buffer.d_input_acc as fp32 scratch (free during bias backward, size >> OC*N_CHUNK).
+            static_assert(decltype(buffer.d_input_acc)::SPEC::SIZE >= OC * N_CHUNK, "d_input_acc too small for bias scratch");
+            float* bias_scratch = (float*)buffer.d_input_acc._data;
+            dim3 p1_grid((unsigned int)OC, N_CHUNK);
+            dim3 p1_block(BN_BS);
+            nn::layers::conv2d::cuda::kernels::bias_backward_partial<T, TI, N_CHUNK><<<p1_grid, p1_block, 0, device.stream>>>(
+                d_conv_out, bias_scratch, SPATIAL, OC);
+            dim3 p2_grid((unsigned int)OC);
+            dim3 p2_block(32);
+            nn::layers::conv2d::cuda::kernels::bias_backward_finalize<T, TI, N_CHUNK><<<p2_grid, p2_block, 0, device.stream>>>(
+                bias_scratch, layer.biases.gradient._data, OC);
+        }
+
+        size_t required_ws = bf_ok ? cached_bf_ws : 0;
+        if constexpr(COMPUTE_INPUT_GRADIENT){
+            if(bd_ok && cached_bd_ws > required_ws){
+                required_ws = cached_bd_ws;
+            }
+        }
+        if(required_ws > 0){
+            utils::assert_exit(device, device.dynamic_memory_allocation_allowed || required_ws <= device.cudnn_workspace_size, "Dynamic CUDA memory allocations are disabled");
+            ensure_cudnn_workspace(device, required_ws);
+        }
+        if(bf_ok){
+            float a = 1, b = 1;
+            check_cudnn_call(device, cudnnConvolutionBackwardFilter(device.cudnn_handle, &a, xd, input._data, yd, d_conv_out,
+                cd, cached_bf_algo, device.cudnn_workspace, device.cudnn_workspace_size, &b, wd, layer.weights.gradient._data), "cudnnConvolutionBackwardFilter conv2d_bwd");
+        }
+        if constexpr(COMPUTE_INPUT_GRADIENT){
+            if(bd_ok){
+            float a = 1, b = 0;
+            check_cudnn_call(device, cudnnConvolutionBackwardData(device.cudnn_handle, &a, wd, layer.weights.parameters._data, yd, d_conv_out,
+                cd, cached_bd_algo, device.cudnn_workspace, device.cudnn_workspace_size, &b, xd, d_input._data), "cudnnConvolutionBackwardData conv2d_bwd");
+            }
+        }
+        check_status(device);
+    }
+
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    void backward_full(devices::CUDA<DEV_SPEC>& device, nn::layers::conv2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::conv2d::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        backward_impl<true>(device, layer, input, d_output, d_input, buffer, mode);
+    }
+
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename INPUT_SPEC, typename D_OUTPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    void backward(devices::CUDA<DEV_SPEC>& device, nn::layers::conv2d::LayerGradient<LAYER_SPEC>& layer, const Tensor<INPUT_SPEC>& input, Tensor<D_OUTPUT_SPEC>& d_output, nn::layers::conv2d::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        backward_impl<false>(device, layer, input, d_output, buffer.d_input_acc, buffer, mode);
+    }
+
+    template<typename DEV_SPEC, typename LAYER_SPEC, typename D_OUTPUT_SPEC, typename D_INPUT_SPEC, typename BUFFER_SPEC, typename MODE = mode::Default<>>
+    void backward_input(devices::CUDA<DEV_SPEC>& device, const nn::layers::conv2d::LayerBackward<LAYER_SPEC>& layer, const Tensor<D_OUTPUT_SPEC>& d_output, Tensor<D_INPUT_SPEC>& d_input, nn::layers::conv2d::Buffer<BUFFER_SPEC>& buffer, const Mode<MODE>& mode = Mode<mode::Default<>>{}){
+        using T = typename LAYER_SPEC::TYPE_POLICY::template GET<numeric_types::categories::Gradient>;
+        using TI = typename devices::CUDA<DEV_SPEC>::index_t;
+        constexpr TI N = LAYER_SPEC::INTERNAL_BATCH_SIZE;
+        constexpr TI IH = LAYER_SPEC::INPUT_HEIGHT, IW = LAYER_SPEC::INPUT_WIDTH, IC = LAYER_SPEC::INPUT_CHANNELS;
+        constexpr TI OH = LAYER_SPEC::OUTPUT_HEIGHT, OW = LAYER_SPEC::OUTPUT_WIDTH, OC = LAYER_SPEC::OUTPUT_CHANNELS;
+        constexpr TI KH = LAYER_SPEC::KERNEL_HEIGHT, KW = LAYER_SPEC::KERNEL_WIDTH;
+        constexpr TI SH = LAYER_SPEC::STRIDE_H, SW = LAYER_SPEC::STRIDE_W;
+        constexpr TI PH = LAYER_SPEC::PADDING_H, PW = LAYER_SPEC::PADDING_W;
+        constexpr cudnnDataType_t dt = nn::cuda::get_cudnn_dtype<T>();
+        constexpr cudnnDataType_t compute_dt = (dt == CUDNN_DATA_BFLOAT16) ? CUDNN_DATA_FLOAT : dt;
+
+        static cudnnTensorDescriptor_t xd = nullptr, yd = nullptr;
+        static cudnnConvolutionDescriptor_t cd = nullptr;
+        static cudnnFilterDescriptor_t wd = nullptr;
+        static cudnnConvolutionBwdDataAlgo_t cached_algo;
+        static size_t cached_ws = 0;
+        static bool algo_ok = false;
+        // magic static: the one-time descriptor/algorithm setup is guarded by the compiler so concurrent first calls from several host threads (one per GPU) are safe
+        static const bool initialized = [&](){
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&xd), "cudnnCreateTensorDescriptor conv2d_bwd_input.xd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(xd, CUDNN_TENSOR_NHWC, dt, N, IC, IH, IW), "cudnnSetTensor4dDescriptor conv2d_bwd_input.xd");
+            check_cudnn_call(device, cudnnCreateTensorDescriptor(&yd), "cudnnCreateTensorDescriptor conv2d_bwd_input.yd");
+            check_cudnn_call(device, cudnnSetTensor4dDescriptor(yd, CUDNN_TENSOR_NHWC, dt, N, OC, OH, OW), "cudnnSetTensor4dDescriptor conv2d_bwd_input.yd");
+            check_cudnn_call(device, cudnnCreateConvolutionDescriptor(&cd), "cudnnCreateConvolutionDescriptor conv2d_bwd_input.cd");
+            check_cudnn_call(device, cudnnSetConvolution2dDescriptor(cd, PH, PW, SH, SW, 1, 1, CUDNN_CROSS_CORRELATION, compute_dt), "cudnnSetConvolution2dDescriptor conv2d_bwd_input.cd");
+            check_cudnn_call(device, cudnnSetConvolutionMathType(cd, CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION), "cudnnSetConvolutionMathType conv2d_bwd_input.cd");
+            check_cudnn_call(device, cudnnCreateFilterDescriptor(&wd), "cudnnCreateFilterDescriptor conv2d_bwd_input.wd");
+            check_cudnn_call(device, cudnnSetFilter4dDescriptor(wd, dt, CUDNN_TENSOR_NHWC, OC, IC, KH, KW), "cudnnSetFilter4dDescriptor conv2d_bwd_input.wd");
+            constexpr int MA = 8; int ac; cudnnConvolutionBwdDataAlgoPerf_t ap[MA];
+            check_cudnn_call(device, cudnnGetConvolutionBackwardDataAlgorithm_v7(device.cudnn_handle, wd, yd, cd, xd, MA, &ac, ap), "cudnnGetConvolutionBackwardDataAlgorithm_v7 conv2d_bwd_input");
+            for(int i = 0; i < ac; i++){
+                if(ap[i].status == CUDNN_STATUS_SUCCESS && ap[i].determinism == CUDNN_DETERMINISTIC){
+                    cached_algo = ap[i].algo;
+                    algo_ok = true;
+                    break;
+                }
+            }
+            utils::assert_exit(device, algo_ok, "No deterministic cuDNN convolution input-gradient algorithm available");
+            check_cudnn_call(device, cudnnGetConvolutionBackwardDataWorkspaceSize(device.cudnn_handle, wd, yd, cd, xd, cached_algo, &cached_ws), "cudnnGetConvolutionBackwardDataWorkspaceSize conv2d_bwd_input");
+            return true;
+        }();
+        (void)initialized;
+        if(algo_ok){
+            if(cached_ws > 0){
+                utils::assert_exit(device, device.dynamic_memory_allocation_allowed || cached_ws <= device.cudnn_workspace_size, "Dynamic CUDA memory allocations are disabled");
+                ensure_cudnn_workspace(device, cached_ws);
+            }
+            float a = 1, b = 0;
+            check_cudnn_call(device, cudnnConvolutionBackwardData(device.cudnn_handle, &a, wd, layer.weights.parameters._data, yd, d_output._data,
+                cd, cached_algo, device.cudnn_workspace, device.cudnn_workspace_size, &b, xd, d_input._data), "cudnnConvolutionBackwardData conv2d_bwd_input");
+        }
+        check_status(device);
+    }
+
+    template<typename DEV_SPEC, typename SPEC>
+    void zero_gradient(devices::CUDA<DEV_SPEC>& device, nn::layers::conv2d::LayerGradient<SPEC>& layer) {
+        check_cuda_call(device, cudaMemsetAsync(layer.weights.gradient._data, 0, decltype(layer.weights.gradient)::SPEC::SIZE_BYTES, device.stream), "cudaMemsetAsync conv2d weights.gradient");
+        check_cuda_call(device, cudaMemsetAsync(layer.biases.gradient._data, 0, decltype(layer.biases.gradient)::SPEC::SIZE_BYTES, device.stream), "cudaMemsetAsync conv2d biases.gradient");
+        if constexpr(SPEC::NORMALIZATION == nn::layers::conv2d::Normalization::BATCH_NORM){
+            check_cuda_call(device, cudaMemsetAsync(layer.norm.gamma.gradient._data, 0, decltype(layer.norm.gamma.gradient)::SPEC::SIZE_BYTES, device.stream), "cudaMemsetAsync conv2d norm.gamma.gradient");
+            check_cuda_call(device, cudaMemsetAsync(layer.norm.beta.gradient._data, 0, decltype(layer.norm.beta.gradient)::SPEC::SIZE_BYTES, device.stream), "cudaMemsetAsync conv2d norm.beta.gradient");
+        }
+        check_status(device);
+    }
+}
+RL_TOOLS_NAMESPACE_WRAPPER_END
+#endif