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 3D transposed convolution, followed by two max pooling layers and a sum operation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.max_pool1 = nn.MaxPool3d(kernel_size=2)
        self.max_pool2 = nn.MaxPool3d(kernel_size=3)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.max_pool1(x)
        x = self.max_pool2(x)
        x = torch.sum(x, dim=1, keepdim=True) 
        return x

batch_size = 16
in_channels = 32
out_channels = 64
depth, height, width = 32, 32, 32
kernel_size = 5
stride = 2
padding = 2

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding]
    ```
    
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 3 of 6 for this task. Your previous solution (turn 2) 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
from torch.utils.cpp_extension import load_inline

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <float.h>

#define TILE 8
#define OC_TILE 8
#define IC_CHUNK 4
#define K 5
#define NTAP 125
#define NTHREADS 256

// ---------------------------------------------------------------------------
// Custom ConvTranspose3d (gather formulation), stride=2, kernel=5.
// out[b, oc, D, H, W] = bias[oc]
//     + sum_{ic,kd,kh,kw} X[b,ic,(D+pad-kd)/2,(H+pad-kh)/2,(W+pad-kw)/2]
//                         * W[ic,oc,kd,kh,kw]
// Only parity-matched, in-range kernel taps are evaluated (~7x FLOP savings
// vs a dense evaluation of the zero-stuffed equivalent conv).
// Weight is expected pre-packed to layout [IC][125][OC] (oc contiguous).
// Block: 256 threads, each handles 2 output voxels x 8 output channels.
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NTHREADS, 2) convt3d_s2k5_kernel(
    const float* __restrict__ X,
    const float* __restrict__ Wp,
    const float* __restrict__ bias,
    float* __restrict__ Y,
    int IC, int Di, int Hi, int Wi,
    int OC, int Do, int Ho, int Wo, int pad)
{
    __shared__ __align__(16) float sx[IC_CHUNK * 216];
    __shared__ __align__(16) float sw[NTAP * IC_CHUNK * OC_TILE];

    const int tid = threadIdx.x;                 // 0..255
    const int tx = tid & 7;
    const int ty = (tid >> 3) & 7;
    const int tz = tid >> 6;                     // 0..3

    const int ntw = (Wo + TILE - 1) / TILE;
    const int nth = (Ho + TILE - 1) / TILE;
    const int tw = blockIdx.x % ntw;
    const int th = (blockIdx.x / ntw) % nth;
    const int td = blockIdx.x / (ntw * nth);

    const int W0 = tw * TILE, H0 = th * TILE, D0 = td * TILE;
    const int oc0 = blockIdx.y * OC_TILE;
    const int b   = blockIdx.z;

    const int W = W0 + tx, H = H0 + ty;
    const int Da = D0 + tz;          // voxel 0 depth
    const int Db = D0 + tz + 4;      // voxel 1 depth
    const bool valid_a = (Da < Do) && (H < Ho) && (W < Wo);
    const bool valid_b = (Db < Do) && (H < Ho) && (W < Wo);

    // input region required by this output tile
    const int d_lo = max(0, (D0 + pad - (K - 1)) / 2);
    const int d_hi = min(Di - 1, (D0 + TILE - 1 + pad) / 2);
    const int h_lo = max(0, (H0 + pad - (K - 1)) / 2);
    const int h_hi = min(Hi - 1, (H0 + TILE - 1 + pad) / 2);
    const int w_lo = max(0, (W0 + pad - (K - 1)) / 2);
    const int w_hi = min(Wi - 1, (W0 + TILE - 1 + pad) / 2);
    const int region_d = d_hi - d_lo + 1;
    const int region_h = h_hi - h_lo + 1;
    const int region_w = w_hi - w_lo + 1;

    // per-thread local tap indices (-1 = invalid)
    int dloc_a[K], dloc_b[K], hloc[K], wloc[K];
    #pragma unroll
    for (int kd = 0; kd < K; ++kd) {
        int t = Da + pad - kd;
        int d = t >> 1;
        dloc_a[kd] = (valid_a && t >= 0 && !(t & 1) && d >= d_lo && d <= d_hi) ? (d - d_lo) : -1;
        t = Db + pad - kd;
        d = t >> 1;
        dloc_b[kd] = (valid_b && t >= 0 && !(t & 1) && d >= d_lo && d <= d_hi) ? (d - d_lo) : -1;
        t = H + pad - kd;
        d = t >> 1;
        hloc[kd] = ((valid_a || valid_b) && t >= 0 && !(t & 1) && d >= h_lo && d <= h_hi) ? (d - h_lo) : -1;
        t = W + pad - kd;
        d = t >> 1;
        wloc[kd] = ((valid_a || valid_b) && t >= 0 && !(t & 1) && d >= w_lo && d <= w_hi) ? (d - w_lo) : -1;
    }

    float acc_a[OC_TILE], acc_b[OC_TILE];
    #pragma unroll
    for (int o = 0; o < OC_TILE; ++o) { acc_a[o] = 0.f; acc_b[o] = 0.f; }

    for (int ic0 = 0; ic0 < IC; ic0 += IC_CHUNK) {
        // stage input tile into shared memory
        for (int idx = tid; idx < IC_CHUNK * 216; idx += NTHREADS) {
            const int ic = idx / 216;
            const int r  = idx % 216;
            const int dz = r / 36;
            const int dy = (r / 6) % 6;
            const int dx = r % 6;
            float v = 0.f;
            if (dz < region_d && dy < region_h && dx < region_w && (ic0 + ic) < IC) {
                const long xb = ((long)(b * IC + ic0 + ic) * Di + (d_lo + dz)) * Hi * Wi;
                v = X[xb + (h_lo + dy) * Wi + (w_lo + dx)];
            }
            sx[idx] = v;
        }
        // stage packed weights into shared memory
        for (int idx = tid; idx < NTAP * IC_CHUNK * OC_TILE; idx += NTHREADS) {
            const int oc  = idx & (OC_TILE - 1);
            const int ic  = (idx / OC_TILE) & (IC_CHUNK - 1);
            const int tap = idx / (OC_TILE * IC_CHUNK);
            float v = 0.f;
            if ((ic0 + ic) < IC && (oc0 + oc) < OC) {
                v = Wp[((long)(ic0 + ic) * NTAP + tap) * OC + (oc0 + oc)];
            }
            sw[idx] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int kd = 0; kd < K; ++kd) {
            const int dla = dloc_a[kd];
            const int dlb = dloc_b[kd];
            if (dla < 0 && dlb < 0) continue;
            #pragma unroll
            for (int kh = 0; kh < K; ++kh) {
                const int hl = hloc[kh];
                if (hl < 0) continue;
                #pragma unroll
                for (int kw = 0; kw < K; ++kw) {
                    const int wl = wloc[kw];
                    if (wl < 0) continue;
                    const int tapid = kd * 25 + kh * 5 + kw;
                    #pragma unroll
                    for (int ic = 0; ic < IC_CHUNK; ++ic) {
                        const float4* wp4 = reinterpret_cast<const float4*>(&sw[(tapid * IC_CHUNK + ic) * OC_TILE]);
                        const float4 w0 = wp4[0];
                        const float4 w1 = wp4[1];
                        if (dla >= 0) {
                            const float xv = sx[ic * 216 + (dla * 6 + hl) * 6 + wl];
                            acc_a[0] = fmaf(xv, w0.x, acc_a[0]);
                            acc_a[1] = fmaf(xv, w0.y, acc_a[1]);
                            acc_a[2] = fmaf(xv, w0.z, acc_a[2]);
                            acc_a[3] = fmaf(xv, w0.w, acc_a[3]);
                            acc_a[4] = fmaf(xv, w1.x, acc_a[4]);
                            acc_a[5] = fmaf(xv, w1.y, acc_a[5]);
                            acc_a[6] = fmaf(xv, w1.z, acc_a[6]);
                            acc_a[7] = fmaf(xv, w1.w, acc_a[7]);
                        }
                        if (dlb >= 0) {
                            const float xv = sx[ic * 216 + (dlb * 6 + hl) * 6 + wl];
                            acc_b[0] = fmaf(xv, w0.x, acc_b[0]);
                            acc_b[1] = fmaf(xv, w0.y, acc_b[1]);
                            acc_b[2] = fmaf(xv, w0.z, acc_b[2]);
                            acc_b[3] = fmaf(xv, w0.w, acc_b[3]);
                            acc_b[4] = fmaf(xv, w1.x, acc_b[4]);
                            acc_b[5] = fmaf(xv, w1.y, acc_b[5]);
                            acc_b[6] = fmaf(xv, w1.z, acc_b[6]);
                            acc_b[7] = fmaf(xv, w1.w, acc_b[7]);
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    const long ostride = (long)Do * Ho * Wo;
    const long ybase = ((long)b * OC + oc0) * ostride + (long)H * Wo + W;
    if (valid_a) {
        const long yb = ybase + (long)Da * Ho * Wo;
        #pragma unroll
        for (int o = 0; o < OC_TILE; ++o) {
            if (oc0 + o < OC) Y[yb + (long)o * ostride] = acc_a[o] + bias[oc0 + o];
        }
    }
    if (valid_b) {
        const long yb = ybase + (long)Db * Ho * Wo;
        #pragma unroll
        for (int o = 0; o < OC_TILE; ++o) {
            if (oc0 + o < OC) Y[yb + (long)o * ostride] = acc_b[o] + bias[oc0 + o];
        }
    }
}

// ---------------------------------------------------------------------------
// Fused MaxPool(2)+MaxPool(3)+sum over channels.
// Composition of non-overlapping 2x2x2 then 3x3x3 max-pools equals a single
// non-overlapping 6x6x6 max-pool. Per output element, compute the per-channel
// max over the 6x6x6 window and sum over all channels (requires C == 64).
// One block (256 threads) per output element.
// ---------------------------------------------------------------------------
__global__ void pool6_sum_kernel(
    const float* __restrict__ Y,
    float* __restrict__ O,
    int C, int D, int H, int W,
    int OD, int OH, int OW)
{
    const int oid = blockIdx.x;
    const int ow = oid % OW;
    const int oh = (oid / OW) % OH;
    const int od = oid / (OW * OH);
    const int b  = blockIdx.y;

    const int tid  = threadIdx.x;   // 0..255
    const int c    = tid & 63;      // channel
    const int part = tid >> 6;      // 0..3

    const long HW = (long)H * W;
    const float* yb = Y + ((long)b * C + c) * D * HW + (long)6 * od * HW + (long)6 * oh * W + 6 * ow;

    float m = -FLT_MAX;
    for (int v = part; v < 216; v += 4) {
        const int dz = v / 36;
        const int dy = (v / 6) % 6;
        const int dx = v % 6;
        m = fmaxf(m, yb[(long)dz * HW + (long)dy * W + dx]);
    }

    __shared__ float smax[4][64];
    __shared__ float warpsum[2];
    smax[part][c] = m;
    __syncthreads();

    if (tid < 64) {
        float mm = fmaxf(fmaxf(smax[0][tid], smax[1][tid]),
                         fmaxf(smax[2][tid], smax[3][tid]));
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            mm += __shfl_down_sync(0xffffffffu, mm, off);
        if ((tid & 31) == 0) warpsum[tid >> 5] = mm;
    }
    __syncthreads();
    if (tid == 0) {
        O[((long)b * OD + od) * OH * OW + (long)oh * OW + ow] = warpsum[0] + warpsum[1];
    }
}

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor wp, torch::Tensor bias,
                                 int64_t stride, int64_t pad)
{
    TORCH_CHECK(x.is_cuda() && wp.is_cuda() && bias.is_cuda(), "tensors must be CUDA");
    auto xc  = x.contiguous();
    auto wpc = wp.contiguous();
    auto bc  = bias.contiguous();

    const int B  = xc.size(0);
    const int IC = xc.size(1);
    const int Di = xc.size(2);
    const int Hi = xc.size(3);
    const int Wi = xc.size(4);
    const int OC = wpc.size(2);

    const int Do = (Di - 1) * (int)stride - 2 * (int)pad + K;
    const int Ho = (Hi - 1) * (int)stride - 2 * (int)pad + K;
    const int Wo = (Wi - 1) * (int)stride - 2 * (int)pad + K;

    auto Y = torch::empty({B, OC, Do, Ho, Wo}, xc.options());

    const int ntw = (Wo + TILE - 1) / TILE;
    const int nth = (Ho + TILE - 1) / TILE;
    const int ntd = (Do + TILE - 1) / TILE;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    dim3 grid(ntw * nth * ntd, (OC + OC_TILE - 1) / OC_TILE, B);
    convt3d_s2k5_kernel<<<grid, NTHREADS, 0, stream>>>(
        xc.data_ptr<float>(), wpc.data_ptr<float>(), bc.data_ptr<float>(),
        Y.data_ptr<float>(), IC, Di, Hi, Wi, OC, Do, Ho, Wo, (int)pad);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    // fused pool1 (k=2,s=2) + pool2 (k=3,s=3) + channel sum
    const int D1 = (Do - 2) / 2 + 1;
    const int H1 = (Ho - 2) / 2 + 1;
    const int W1 = (Wo - 2) / 2 + 1;
    const int OD = (D1 - 3) / 3 + 1;
    const int OH = (H1 - 3) / 3 + 1;
    const int OW = (W1 - 3) / 3 + 1;

    auto O = torch::empty({B, 1, OD, OH, OW}, xc.options());
    dim3 pgrid(OD * OH * OW, B, 1);
    pool6_sum_kernel<<<pgrid, 256, 0, stream>>>(
        Y.data_ptr<float>(), O.data_ptr<float>(), OC, Do, Ho, Wo, OD, OH, OW);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return O;
}
"""

_cpp_src = "torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor wp, torch::Tensor bias, int64_t stride, int64_t pad);"

_fused = load_inline(
    name="convt3d_pool_sum_fused_v2",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized Model:
      - ConvTranspose3d replaced by a custom tiled CUDA kernel (stride=2, k=5
        gather formulation) that evaluates only parity-matched kernel taps,
        with shared-memory blocking of inputs and pre-packed weights.
      - MaxPool3d(2) + MaxPool3d(3) + channel-sum fused into a single CUDA
        kernel exploiting that the composed non-overlapping pooling equals
        one non-overlapping 6x6x6 max-pool.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.max_pool1 = nn.MaxPool3d(kernel_size=2)
        self.max_pool2 = nn.MaxPool3d(kernel_size=3)

    def forward(self, x):
        ct = self.conv_transpose
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and ct.kernel_size == (5, 5, 5)
            and ct.stride == (2, 2, 2)
            and ct.padding == (2, 2, 2)
            and ct.dilation == (1, 1, 1)
            and ct.output_padding == (0, 0, 0)
            and ct.groups == 1
            and ct.bias is not None
            and ct.out_channels == 64
            and ct.weight.is_cuda
            and ct.weight.dtype == torch.float32
            and x.dim() == 5
            and x.size(2) >= 4
            and x.size(3) >= 4
            and x.size(4) >= 4
        ):
            # Pack weight [IC, OC, 5, 5, 5] -> [IC, 125, OC] (cheap: ~1MB copy).
            wp = ct.weight.permute(0, 2, 3, 4, 1).contiguous()
            return _fused.fused_forward_cuda(
                x, wp, ct.bias, ct.stride[0], ct.padding[0]
            )
        # generic fallback (matches reference semantics)
        x = self.conv_transpose(x)
        x = self.max_pool1(x)
        x = self.max_pool2(x)
        x = torch.sum(x, dim=1, keepdim=True)
        return x
```
Result of evaluating it:
The solution compiled and ran but its output did not match the reference. Details:
```
{"avg_difference": ["0.373057", "0.376185", "0.378482", "0.373160", "0.374262"], "correctness_issue": "Output mismatch", "correctness_trials": "(0 / 5)", "max_difference": ["1.670593", "1.449350", "1.465643", "1.465065", "1.436789"]}
```

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.
