You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

    ```
    # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed 3D convolution, applies ReLU, and then applies group normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, bias=bias)
        self.relu = nn.ReLU()
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, D, H, W).
        """
        x = self.conv_transpose(x)
        x = self.relu(x)
        x = self.group_norm(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
D, H, W = 32, 32, 32
kernel_size = 3
groups = 8
bias = False

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, groups, bias]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 4 of 6 for this task. Your previous solution (turn 3) was:
```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Optimized pipeline (all activations kept in channels_last_3d end-to-end):
#   0) Fast custom NCHW->NHWC transpose for the input (the harness feeds NCHW).
#      Tiled 32x32 shared-memory transpose, ~2x faster than the generic copy.
#   1) ConvTranspose3d rewritten as an equivalent forward Conv3d: for stride-1,
#      pad-0 transposed conv, out == conv3d(x, flip(w).swapaxes, pad=k-1).
#      The cuDNN forward implicit-GEMM NHWC path avoids the layout-conversion
#      kernels the dgrad path was emitting (~220us). The transformed weight is
#      cached, and on the first call the equivalence is verified against the
#      real conv_transpose with an automatic permanent fallback.
#   2) ReLU+GroupNorm stays fused: stats kernel (ReLU on-the-fly), a tiny
#      finalize kernel producing per-(n,c) scale/shift, and a vectorized apply
#      kernel. ReLU output is never materialized.
fastpath_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <climits>

#define STATS_BLOCK 256
#define APPLY_BLOCK 256
#define APPLY_ROWS 8

// ------------------------- NCHW -> NHWC transpose -------------------------
__global__ void nchw2nhwc_kernel(const float* __restrict__ src,
                                 float* __restrict__ dst,
                                 int C, int S) {
    __shared__ float tile[32][33];
    const int n  = blockIdx.z;
    const int c0 = blockIdx.x * 32;
    const int s0 = blockIdx.y * 32;
    const size_t base = (size_t)n * C * S;
    const int tx = threadIdx.x;  // 32
    const int ty = threadIdx.y;  // 8

    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int c = c0 + ty + i;
        const int s = s0 + tx;
        if (c < C && s < S) tile[ty + i][tx] = src[base + (size_t)c * S + s];
    }
    __syncthreads();
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int s = s0 + ty + i;
        const int c = c0 + tx;
        if (s < S && c < C) dst[base + (size_t)s * C + c] = tile[tx][ty + i];
    }
}

torch::Tensor nchw_to_nhwc(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat && x.dim() == 5,
                "bad input");
    auto xc = x.contiguous();
    const int64_t N = xc.size(0), C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    TORCH_CHECK(S <= INT_MAX && C <= INT_MAX && N <= 65535, "too big");
    auto y = at::empty_like(xc, at::MemoryFormat::ChannelsLast3d);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const dim3 block(32, 8, 1);
    const dim3 grid((unsigned)((C + 31) / 32), (unsigned)((S + 31) / 32),
                    (unsigned)N);
    nchw2nhwc_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), y.data_ptr<float>(), (int)C, (int)S);
    return y;
}

// --------------------------- fused norm pieces ---------------------------
template <int BLOCK>
__device__ __forceinline__ void block_reduce2(float fs, float fq,
                                              double* out_s, double* out_q) {
    double ds = (double)fs;
    double dq = (double)fq;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        ds += __shfl_down_sync(0xffffffffu, ds, off);
        dq += __shfl_down_sync(0xffffffffu, dq, off);
    }
    __shared__ double shs[BLOCK / 32];
    __shared__ double shq[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) { shs[wid] = ds; shq[wid] = dq; }
    __syncthreads();
    if (wid == 0) {
        constexpr int NW = BLOCK / 32;
        ds = (lane < NW) ? shs[lane] : 0.0;
        dq = (lane < NW) ? shq[lane] : 0.0;
        #pragma unroll
        for (int off = NW / 2; off > 0; off >>= 1) {
            ds += __shfl_down_sync(0xffffffffu, ds, off);
            dq += __shfl_down_sync(0xffffffffu, dq, off);
        }
        if (lane == 0) { *out_s = ds; *out_q = dq; }
    }
}

__device__ __forceinline__ void relu4_accum(const float4 v, float& fs, float& fq) {
    const float a0 = fmaxf(v.x, 0.0f);
    const float a1 = fmaxf(v.y, 0.0f);
    const float a2 = fmaxf(v.z, 0.0f);
    const float a3 = fmaxf(v.w, 0.0f);
    fs += (a0 + a1) + (a2 + a3);
    fq += (a0 * a0 + a1 * a1) + (a2 * a2 + a3 * a3);
}

// stats: grid (splits, G, N); each block reduces rows [s0, s1) of one group.
template <int VSHIFT>
__global__ void stats_pow2_kernel(const float* __restrict__ x,
                                  double* __restrict__ partials,
                                  int S, int C, int Cg, int G,
                                  int rows_per_split, int splits) {
    constexpr int V = 1 << VSHIFT;
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    const int c4 = threadIdx.x & (V - 1);
    const int r0 = threadIdx.x >> VSHIFT;
    const int rstride = STATS_BLOCK >> VSHIFT;
    float fs = 0.0f, fq = 0.0f;
    for (int r = r0; r < s1 - s0; r += rstride) {
        const float4 v = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + r) * C + (c4 << 2));
        relu4_accum(v, fs, fq);
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void stats_vec4_kernel(const float* __restrict__ x,
                                  double* __restrict__ partials,
                                  int S, int C, int Cg, int G, int V,
                                  int rows_per_split, int splits) {
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    float fs = 0.0f, fq = 0.0f;
    const int count = (s1 - s0) * V;
    for (int idx = threadIdx.x; idx < count; idx += STATS_BLOCK) {
        const int row = idx / V;
        const int c4  = idx - row * V;
        const float4 v = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + row) * C + (c4 << 2));
        relu4_accum(v, fs, fq);
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void stats_scalar_kernel(const float* __restrict__ x,
                                    double* __restrict__ partials,
                                    int S, int C, int Cg, int G,
                                    int rows_per_split, int splits) {
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    float fs = 0.0f, fq = 0.0f;
    for (int r = s0; r < s1; ++r) {
        const float* rowp = base + (size_t)r * C;
        for (int c = threadIdx.x; c < Cg; c += STATS_BLOCK) {
            const float v = fmaxf(rowp[c], 0.0f);
            fs += v;
            fq += v * v;
        }
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void finalize_affine_kernel(const double* __restrict__ partials,
                                       const float* __restrict__ gamma,
                                       const float* __restrict__ beta,
                                       float* __restrict__ a,
                                       float* __restrict__ b,
                                       int NC, int C, int Cg, int G,
                                       int splits, double inv_gs, double eps) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= NC) return;
    const int n = idx / C;
    const int c = idx - n * C;
    const int g = c / Cg;
    const double* p = partials + (size_t)(n * G + g) * splits * 2;
    double s = 0.0, q = 0.0;
    for (int j = 0; j < splits; ++j) {
        s += p[2 * j];
        q += p[2 * j + 1];
    }
    const double m = s * inv_gs;
    double var = q * inv_gs - m * m;
    if (var < 0.0) var = 0.0;
    const double rstd = 1.0 / sqrt(var + eps);
    const float av = (float)(rstd * (double)gamma[c]);
    a[idx] = av;
    b[idx] = (float)((double)beta[c] - m * (double)av);
}

// apply: grid (blocks_x, N). Each thread owns one channel-quad and walks rows.
template <int VSHIFT>
__global__ void apply_pow2_kernel(const float* __restrict__ x,
                                  const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int S, int C) {
    constexpr int V = 1 << VSHIFT;  // = C/4
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c4  = tid & (V - 1);
    const int s0  = tid >> VSHIFT;
    const int sstride = (gridDim.x * APPLY_BLOCK) >> VSHIFT;
    if (s0 >= S) return;
    const size_t img = (size_t)n * S * V;
    const float4 av = *reinterpret_cast<const float4*>(a + (size_t)n * C + (c4 << 2));
    const float4 bv = *reinterpret_cast<const float4*>(b + (size_t)n * C + (c4 << 2));
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * V + c4;
        const float4 v = *reinterpret_cast<const float4*>(x + (off << 2));
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), av.x, bv.x);
        o.y = fmaf(fmaxf(v.y, 0.0f), av.y, bv.y);
        o.z = fmaf(fmaxf(v.z, 0.0f), av.z, bv.z);
        o.w = fmaf(fmaxf(v.w, 0.0f), av.w, bv.w);
        *reinterpret_cast<float4*>(y + (off << 2)) = o;
    }
}

__global__ void apply_vec4_kernel(const float* __restrict__ x,
                                  const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int S, int C, int V) {
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c4  = tid % V;
    const int s0  = tid / V;
    const int sstride = (gridDim.x * APPLY_BLOCK) / V;
    if (s0 >= S || sstride <= 0) return;
    const size_t img = (size_t)n * S * V;
    const float4 av = *reinterpret_cast<const float4*>(a + (size_t)n * C + (c4 << 2));
    const float4 bv = *reinterpret_cast<const float4*>(b + (size_t)n * C + (c4 << 2));
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * V + c4;
        const float4 v = *reinterpret_cast<const float4*>(x + (off << 2));
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), av.x, bv.x);
        o.y = fmaf(fmaxf(v.y, 0.0f), av.y, bv.y);
        o.z = fmaf(fmaxf(v.z, 0.0f), av.z, bv.z);
        o.w = fmaf(fmaxf(v.w, 0.0f), av.w, bv.w);
        *reinterpret_cast<float4*>(y + (off << 2)) = o;
    }
}

__global__ void apply_scalar_kernel(const float* __restrict__ x,
                                    const float* __restrict__ a,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    size_t total, int C) {
    const int n = blockIdx.y;
    const size_t base = (size_t)n * total;
    const size_t ab   = (size_t)n * C;
    const size_t stride = (size_t)gridDim.x * APPLY_BLOCK;
    for (size_t p = (size_t)blockIdx.x * APPLY_BLOCK + threadIdx.x;
         p < total; p += stride) {
        const int c = (int)(p % (size_t)C);
        const float v = fmaxf(x[base + p], 0.0f);
        y[base + p] = fmaf(v, a[ab + c], b[ab + c]);
    }
}

static inline int log2_int(int v) {
    int s = 0;
    while (v > 1) { v >>= 1; ++s; }
    return s;
}

torch::Tensor fused_relu_groupnorm_nhwc(torch::Tensor x, torch::Tensor gamma,
                                        torch::Tensor beta, int64_t groups,
                                        double eps) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 5, "x must have shape (N, C, D, H, W)");

    auto xc = x.contiguous(at::MemoryFormat::ChannelsLast3d);
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    const int64_t G = groups;
    TORCH_CHECK(G > 0 && C % G == 0, "C must be divisible by groups");
    const int64_t Cg = C / G;
    TORCH_CHECK(S > 0 && S <= INT_MAX && C <= INT_MAX && Cg <= INT_MAX,
                "bad tensor dims");
    TORCH_CHECK(N <= 65535 && G <= 65535, "N/G too large for grid");
    TORCH_CHECK(gc.numel() == C && bc.numel() == C, "gamma/beta size mismatch");

    auto y = at::empty_like(xc);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t NG = N * G;
    int splits = (int)((8LL * sm_count + NG - 1) / NG);
    if (splits < 1) splits = 1;
    if (splits > 32) splits = 32;
    if ((int64_t)splits > S) splits = (int)S;
    int rows_per_split = (int)((S + splits - 1) / splits);
    splits = (int)((S + rows_per_split - 1) / rows_per_split);

    auto partials = at::empty({NG * splits * 2}, xc.options().dtype(at::kDouble));
    const dim3 sgrid(splits, (unsigned)G, (unsigned)N);
    const float* xp = xc.data_ptr<float>();
    double* pp = partials.data_ptr<double>();

    const bool vec_ok = (C % 4 == 0) && (Cg % 4 == 0);
    const int Vs = (int)(Cg / 4);
    const bool vpow2 = vec_ok && Vs > 0 && (Vs & (Vs - 1)) == 0;
    if (vec_ok && vpow2) {
        const int sh = log2_int(Vs);
        switch (sh) {
            case 0: stats_pow2_kernel<0><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 1: stats_pow2_kernel<1><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 2: stats_pow2_kernel<2><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 3: stats_pow2_kernel<3><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 4: stats_pow2_kernel<4><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 5: stats_pow2_kernel<5><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 6: stats_pow2_kernel<6><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            default: stats_vec4_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, Vs, rows_per_split, splits); break;
        }
    } else if (vec_ok) {
        stats_vec4_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, Vs, rows_per_split, splits);
    } else {
        stats_scalar_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits);
    }

    const int NC = (int)(N * C);
    auto a = at::empty({NC}, xc.options());
    auto b = at::empty({NC}, xc.options());
    finalize_affine_kernel<<<(NC + 255) / 256, 256, 0, stream>>>(
        pp, gc.data_ptr<float>(), bc.data_ptr<float>(),
        a.data_ptr<float>(), b.data_ptr<float>(),
        NC, (int)C, (int)Cg, (int)G, splits,
        1.0 / (double)(Cg * S), eps);

    float* yp = y.data_ptr<float>();
    const float* ap = a.data_ptr<float>();
    const float* bp = b.data_ptr<float>();
    if (C % 4 == 0) {
        const int V = (int)(C / 4);
        int64_t blocks_x = (S * V + (int64_t)APPLY_BLOCK * APPLY_ROWS - 1) /
                           ((int64_t)APPLY_BLOCK * APPLY_ROWS);
        const int64_t min_blocks = (V + APPLY_BLOCK - 1) / APPLY_BLOCK;
        if (blocks_x < min_blocks) blocks_x = min_blocks;
        if (blocks_x < 1) blocks_x = 1;
        const dim3 agrid((unsigned)blocks_x, (unsigned)N, 1);
        const bool vp2 = (V & (V - 1)) == 0;
        if (vp2) {
            const int sh = log2_int(V);
            switch (sh) {
                case 0: apply_pow2_kernel<0><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 1: apply_pow2_kernel<1><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 2: apply_pow2_kernel<2><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 3: apply_pow2_kernel<3><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 4: apply_pow2_kernel<4><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 5: apply_pow2_kernel<5><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 6: apply_pow2_kernel<6><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 7: apply_pow2_kernel<7><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                default: apply_vec4_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C, V); break;
            }
        } else {
            apply_vec4_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C, V);
        }
    } else {
        const size_t total = (size_t)S * (size_t)C;
        int64_t blocks_x = ((int64_t)total + (int64_t)APPLY_BLOCK * APPLY_ROWS - 1) /
                           ((int64_t)APPLY_BLOCK * APPLY_ROWS);
        if (blocks_x < 1) blocks_x = 1;
        const dim3 agrid((unsigned)blocks_x, (unsigned)N, 1);
        apply_scalar_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, total, (int)C);
    }

    return y;
}
"""

fastpath_cpp_source = (
    "torch::Tensor nchw_to_nhwc(torch::Tensor x);\n"
    "torch::Tensor fused_relu_groupnorm_nhwc(torch::Tensor x, torch::Tensor gamma, "
    "torch::Tensor beta, int64_t groups, double eps);"
)

fastpath = load_inline(
    name="convgn_fastpath_ext_v3",
    cpp_sources=fastpath_cpp_source,
    cuda_sources=fastpath_cuda_source,
    functions=["nchw_to_nhwc", "fused_relu_groupnorm_nhwc"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * runs everything in channels_last_3d (cuDNN NHWC tensor-core kernels,
        no layout conversions),
      * evaluates the stride-1 ConvTranspose3d as an equivalent forward Conv3d
        (flipped/swapped kernel, full padding) which avoids the dgrad-path
        convertTensor kernels; verified at runtime with a permanent fallback,
      * fuses ReLU + GroupNorm into custom CUDA kernels that consume the NHWC
        convolution output directly (ReLU never materialized).
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cuda.matmul.allow_tf32 = True
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, bias=bias
        )
        with torch.no_grad():
            self.conv_transpose.weight.data = (
                self.conv_transpose.weight.data.to(
                    memory_format=torch.channels_last_3d
                )
            )
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)
        self.groups = groups
        self.eps = self.group_norm.eps
        self.fused = fastpath
        if isinstance(kernel_size, (tuple, list)):
            kd, kh, kw = (int(kernel_size[0]), int(kernel_size[1]), int(kernel_size[2]))
        else:
            kd = kh = kw = int(kernel_size)
        self.pad = (kd - 1, kh - 1, kw - 1)
        self.stride = self.conv_transpose.stride
        self.dilation = self.conv_transpose.dilation
        self.conv_groups = self.conv_transpose.groups
        self._equiv_ok = (
            self.stride == (1, 1, 1)
            and self.dilation == (1, 1, 1)
            and self.conv_groups == 1
            and self.conv_transpose.padding == (0, 0, 0)
            and self.conv_transpose.output_padding == (0, 0, 0)
        )
        self._w_key = None
        self._w_conv = None
        self._use_conv3d = False
        self._checked = False

    def _get_conv_weight(self):
        w = self.conv_transpose.weight
        key = (w.data_ptr(), w._version, w.device, w.shape)
        if self._w_conv is not None and key == self._w_key:
            return self._w_conv
        with torch.no_grad():
            wc = w.detach().permute(1, 0, 2, 3, 4).flip([2, 3, 4]).contiguous(
                memory_format=torch.channels_last_3d
            )
        self._w_conv = wc
        self._w_key = key
        return wc

    def forward(self, x):
        x = self.fused.nchw_to_nhwc(x)
        if self._use_conv3d:
            x = F.conv3d(x, self._w_conv, None, (1, 1, 1), self.pad)
        else:
            if self._equiv_ok and not self._checked:
                self._checked = True
                try:
                    wc = self._get_conv_weight()
                    ref = self.conv_transpose(x)
                    alt = F.conv3d(x, wc, None, (1, 1, 1), self.pad)
                    if torch.allclose(alt, ref, atol=1e-2, rtol=1e-2):
                        self._use_conv3d = True
                        x = alt
                    else:
                        self._w_conv = None
                        self._w_key = None
                        x = ref
                except Exception:
                    self._w_conv = None
                    self._w_key = None
                    x = self.conv_transpose(x)
            else:
                x = self.conv_transpose(x)
        return self.fused.fused_relu_groupnorm_nhwc(
            x, self.group_norm.weight, self.group_norm.bias, self.groups, self.eps
        )
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 3.21 ms; the PyTorch reference (with TF32 enabled) takes 4.45 ms; speedup 1.386x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| sm80_xmma_fprop_implicit_gemm_tf32f32_tf32f32_f32_nhwckrsc_nhwc_tilesize128x128x16_stage4_ | 2024.3 | 1 |
| void apply_pow2_kernel<5>(float const*, float const*, float const*, float*, int, int) | 453.1 | 1 |
| void stats_pow2_kernel<2>(float const*, double*, int, int, int, int, int, int) | 252.5 | 1 |
| void cudnn::engines_precompiled::convertTensor_kernel<float, float, float, (cudnnKernelDat | 227.2 | 2 |
| nchw2nhwc_kernel(float const*, float*, int, int) | 200.0 | 1 |
| finalize_affine_kernel(double const*, float const*, float const*, float*, float*, int, int | 4.7 | 1 |

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.
