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 convolution, global average pooling, adds a bias, applies log-sum-exp, sum, and multiplication.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)  # Global average pooling
        x = x + self.bias
        x = torch.logsumexp(x, dim=1, keepdim=True)  # Log-sum-exp
        x = torch.sum(x, dim=(2, 3))  # Sum
        x = x * 10.0  # Multiplication
        return x

batch_size = 16
in_channels = 64
out_channels = 128
height = width = 512
kernel_size = 3
bias_shape = (out_channels, 1, 1)

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, bias_shape]
    ```
    
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 6 of 6 for this task. Your previous solution (turn 5) 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

# Key algorithmic insight (exact in FP32):
# For ConvTranspose2d with stride=1, padding=0, dilation=1, groups=1 (a "full"
# convolution), the global sum of each output channel factorizes exactly:
#     sum_{h,w} out[n,c,h,w] = sum_k (sum_{h,w} x[n,k,h,w]) * (sum_{p,q} w[k,c,p,q])
#                              + conv_bias[c] * H_out * W_out
# So the whole network collapses to:
#   1) per-channel sum of the input            (memory-bound, streams ~1 GiB)
#   2) per-(in,out)-channel sum of the weights (tiny)
#   3) small matmul + bias adds + logsumexp + scale (fully fused)
#
# Performance notes:
#   * The dominant reduction kernel streams the whole input tensor as one flat
#     range using float4 loads with streaming (evict-first) cache hints. The
#     grid size is chosen as an exact multiple of the GPU's resident-block
#     capacity (queried via the occupancy API), which eliminates the wave-
#     quantization tail on any architecture (A100/H100/etc.). Blocks that
#     straddle a channel boundary handle each channel segment separately.
#   * The tiny weight-tap-sum work is fused into the same kernel launch (a few
#     extra blocks at the front of the grid), removing one kernel launch. The
#     per-channel accumulator is zeroed with a cheap async memset.
#   * The final matmul+bias+logsumexp kernel uses warp-shuffle reductions
#     instead of shared-memory trees (fewer __syncthreads).

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math.h>
#include <cstdint>

#define FULL_MASK 0xffffffffu

// Warp + cross-warp block reduction (512 threads = 16 warps). Result valid on thread 0.
__device__ __forceinline__ float block_reduce(float s) {
    __shared__ float sh[16];
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_down_sync(FULL_MASK, s, off);
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = s;
    __syncthreads();
    float r = 0.0f;
    if (wid == 0) {
        r = (lane < 16) ? sh[lane] : 0.0f;
        for (int off = 8; off > 0; off >>= 1)
            r += __shfl_down_sync(FULL_MASK, r, off);
    }
    return r;
}

// Fused kernel:
//   * The first `wblocks` blocks compute the per-(in,out)-channel sums of the
//     K convolution taps (tiny; runs concurrently with the streaming blocks).
//   * The remaining blocks stream the whole input tensor as one flat range,
//     split into equal chunks. Each block accumulates float4-vectorized
//     partial sums per channel segment and atomically adds them into
//     Sx[channel].
// Requires: HW % 4 == 0, chunk % 4 == 0, x 16-byte aligned, Sx pre-zeroed.
__global__ void __launch_bounds__(512)
fused_stream_wsum_kernel(const float* __restrict__ x,
                         const float* __restrict__ w,
                         float* __restrict__ Sx,
                         float* __restrict__ Sw,
                         long HW, long total, long chunk,
                         int wtotal, int K, int wblocks) {
    if ((int)blockIdx.x < wblocks) {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < wtotal) {
            const float* p = w + (long)idx * K;
            float s = 0.0f;
            for (int i = 0; i < K; ++i) s += p[i];
            Sw[idx] = s;
        }
        return;
    }
    long start = ((long)blockIdx.x - wblocks) * chunk;
    long end = start + chunk;
    if (end > total) end = total;
    const long stride = blockDim.x;
    while (start < end) {
        long channel = start / HW;
        long seg_end = (channel + 1) * HW;
        if (seg_end > end) seg_end = end;
        const float4* p4 = reinterpret_cast<const float4*>(x + start);
        long n4 = (seg_end - start) >> 2;
        float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
        long i = threadIdx.x;
        long limit = n4 - 3 * stride;
        for (; i < limit; i += 4 * stride) {
            float4 a = __ldcs(p4 + i);
            float4 b = __ldcs(p4 + i + stride);
            float4 c = __ldcs(p4 + i + 2 * stride);
            float4 d = __ldcs(p4 + i + 3 * stride);
            acc.x += a.x + b.x + c.x + d.x;
            acc.y += a.y + b.y + c.y + d.y;
            acc.z += a.z + b.z + c.z + d.z;
            acc.w += a.w + b.w + c.w + d.w;
        }
        for (; i < n4; i += stride) {
            float4 a = __ldcs(p4 + i);
            acc.x += a.x; acc.y += a.y; acc.z += a.z; acc.w += a.w;
        }
        float s = (acc.x + acc.y) + (acc.z + acc.w);
        s = block_reduce(s);
        if (threadIdx.x == 0) atomicAdd(Sx + channel, s);
        start = seg_end;
    }
}

// Fused init for the scalar fallback path: zero Sx AND compute tap sums.
__global__ void init_ws_kernel(const float* __restrict__ w,
                               float* __restrict__ Sw,
                               float* __restrict__ Sx,
                               int total, int NC, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < NC) Sx[idx] = 0.0f;
    if (idx < total) {
        const float* ptr = w + (long)idx * K;
        float s = 0.0f;
        for (int i = 0; i < K; ++i) s += ptr[i];
        Sw[idx] = s;
    }
}

// Scalar fallback for the (rare) case where planes are not 16-byte aligned.
__global__ void __launch_bounds__(512)
channel_sum_scalar_kernel(const float* __restrict__ x,
                          float* __restrict__ out,
                          long HW, long chunk_floats, int chunks) {
    long channel = blockIdx.x / chunks;
    int cid = blockIdx.x % chunks;
    long start = (long)cid * chunk_floats;
    long len = HW - start;
    if (len > chunk_floats) len = chunk_floats;

    const float* p = x + channel * HW + start;
    float s = 0.0f;
    for (long j = threadIdx.x; j < len; j += blockDim.x) {
        s += __ldcs(p + j);
    }
    s = block_reduce(s);
    if (threadIdx.x == 0) atomicAdd(out + channel, s);
}

// One block per batch element: fused matmul (Sx @ Sw), scale, bias adds,
// log-sum-exp over channels and final multiplication by 10. Uses warp-shuffle
// reductions (block size is always a multiple of 32).
__global__ void matmul_bias_lse_kernel(const float* __restrict__ Sx,
                                       const float* __restrict__ Sw,
                                       const float* __restrict__ conv_bias,
                                       const float* __restrict__ bias,
                                       float* __restrict__ out,
                                       int C_in, int C_out,
                                       float inv_hw) {
    int n = blockIdx.x;
    int c = threadIdx.x;
    bool valid = c < C_out;

    float y = -3.402823466e38f;
    if (valid) {
        float acc = 0.0f;
        const float* sx = Sx + (long)n * C_in;
        const float* sw = Sw + c;
        #pragma unroll 4
        for (int k = 0; k < C_in; ++k) {
            acc += sx[k] * sw[(long)k * C_out];
        }
        y = acc * inv_hw + conv_bias[c] + bias[c];
    }

    int lane = c & 31;
    int wid = c >> 5;
    int nw = blockDim.x >> 5;
    __shared__ float sh[32];

    // Block reduction: max over channels.
    float m = y;
    for (int off = 16; off > 0; off >>= 1)
        m = fmaxf(m, __shfl_xor_sync(FULL_MASK, m, off));
    if (lane == 0) sh[wid] = m;
    __syncthreads();
    if (wid == 0) {
        m = (lane < nw) ? sh[lane] : -3.402823466e38f;
        for (int off = 16; off > 0; off >>= 1)
            m = fmaxf(m, __shfl_xor_sync(FULL_MASK, m, off));
        if (lane == 0) sh[0] = m;
    }
    __syncthreads();
    m = sh[0];

    // Block reduction: sum of exp(y - max).
    float e = valid ? expf(y - m) : 0.0f;
    for (int off = 16; off > 0; off >>= 1)
        e += __shfl_xor_sync(FULL_MASK, e, off);
    __syncthreads();  // ensure everyone has read sh[0] before reuse
    if (lane == 0) sh[wid] = e;
    __syncthreads();
    if (wid == 0) {
        float s = (lane < nw) ? sh[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            s += __shfl_xor_sync(FULL_MASK, s, off);
        if (lane == 0) out[n] = (m + logf(s)) * 10.0f;
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias,
                            int64_t kernel_h, int64_t kernel_w) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "fused_forward expects a CUDA float32 input");
    TORCH_CHECK(x.dim() == 4, "fused_forward expects a 4D input");
    int64_t N = x.size(0);
    int64_t C_in = x.size(1);
    int64_t H = x.size(2);
    int64_t W = x.size(3);
    int64_t C_out = weight.size(1);
    TORCH_CHECK(C_out <= 1024, "C_out must be <= 1024");

    int64_t H_out = H + kernel_h - 1;
    int64_t W_out = W + kernel_w - 1;
    float inv_hw = 1.0f / (float)(H_out * W_out);

    // Single workspace: Sx (N*C_in) followed by Sw (C_in*C_out).
    auto workspace = torch::empty({N * C_in + C_in * C_out}, x.options());
    float* Sx = workspace.data_ptr<float>();
    float* Sw = Sx + N * C_in;
    auto out = torch::empty({N, 1}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* xptr = x.data_ptr<float>();
    const float* wptr = weight.data_ptr<float>();

    long HW = (long)H * (long)W;
    long NC = (long)N * (long)C_in;
    long total = NC * HW;
    int wtotal = (int)(C_in * C_out);
    int K = (int)(kernel_h * kernel_w);

    bool vec = ((HW & 3) == 0) &&
               ((reinterpret_cast<uintptr_t>(xptr) & 15) == 0);

    if (vec) {
        // Number of blocks that can be resident simultaneously (cached).
        int dev = 0;
        cudaGetDevice(&dev);
        static int wave_cache[64] = {0};
        int wave = (dev >= 0 && dev < 64) ? wave_cache[dev] : 0;
        if (wave == 0) {
            int sms = 0;
            cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev);
            int nb = 0;
            cudaOccupancyMaxActiveBlocksPerMultiprocessor(
                &nb, fused_stream_wsum_kernel, 512, 0);
            if (nb <= 0) nb = 4;
            if (sms <= 0) sms = 108;
            wave = sms * nb;
            if (dev >= 0 && dev < 64) wave_cache[dev] = wave;
        }
        // Pick a chunk size so the grid is (nearly) an exact multiple of the
        // resident-block capacity: no wave-quantization tail on any GPU.
        const long TARGET = 32768;  // ~128 KB streamed per block
        long denom = (long)wave * TARGET;
        long k = (total + denom - 1) / denom;
        if (k < 1) k = 1;
        long grid_target = (long)wave * k;
        long chunk = (total + grid_target - 1) / grid_target;
        chunk = ((chunk + 3) / 4) * 4;  // keep float4 alignment
        if (chunk > HW) chunk = HW;     // HW % 4 == 0 in this path
        if (chunk < 4) chunk = 4;
        long main_blocks = (total + chunk - 1) / chunk;
        int wblocks = (wtotal + 511) / 512;

        cudaMemsetAsync(Sx, 0, NC * sizeof(float), stream);
        fused_stream_wsum_kernel<<<(int)(main_blocks + wblocks), 512, 0, stream>>>(
            xptr, wptr, Sx, Sw, HW, total, chunk, wtotal, K, wblocks);
    } else {
        int init_threads = 256;
        int NCi = (int)NC;
        int init_max = wtotal > NCi ? wtotal : NCi;
        init_ws_kernel<<<(init_max + init_threads - 1) / init_threads, init_threads, 0, stream>>>(
            wptr, Sw, Sx, wtotal, NCi, K);

        const long CHUNK = 32768;
        long chunks = (HW + CHUNK - 1) / CHUNK;
        if (chunks * NC < 2048) {
            long want = (2048 + NC - 1) / NC;
            long max_chunks = (HW + 2047) / 2048;
            if (want > max_chunks) want = max_chunks;
            if (want > chunks) chunks = want;
            if (chunks < 1) chunks = 1;
        }
        long chunk_floats = (HW + chunks - 1) / chunks;
        chunks = (HW + chunk_floats - 1) / chunk_floats;
        channel_sum_scalar_kernel<<<(int)(chunks * NC), 512, 0, stream>>>(
            xptr, Sx, HW, chunk_floats, (int)chunks);
    }

    int block = 32;
    while (block < C_out && block < 1024) block <<= 1;
    matmul_bias_lse_kernel<<<(int)N, block, 0, stream>>>(
        Sx, Sw, conv_bias.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), (int)C_in, (int)C_out, inv_hw);

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias, int64_t kernel_h, int64_t kernel_w);"
)

fused_module = load_inline(
    name="fused_convtranspose_gap_lse_v5",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Replaces the expensive transposed convolution +
    global average pooling + bias + logsumexp + sum + scale pipeline with two
    fused custom CUDA kernel launches, exploiting the exact factorization of the
    global sum of a full (stride-1) transposed convolution. The dominant input
    reduction kernel streams the input with vectorized float4 loads over a flat
    grid sized to an exact multiple of the GPU's resident-block capacity (found
    via the occupancy API), eliminating wave-quantization tails on any
    architecture; the tiny weight-reduction is fused into the same launch. The
    final matmul + bias + logsumexp + scale kernel uses warp-shuffle
    reductions. All math is done in FP32.
    """

    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        if isinstance(kernel_size, (tuple, list)):
            self.kernel_h = int(kernel_size[0])
            self.kernel_w = int(kernel_size[1])
        else:
            self.kernel_h = int(kernel_size)
            self.kernel_w = int(kernel_size)

        def _is_one(v):
            return v == 1 or tuple(v) == (1, 1)

        def _is_zero(v):
            return v == 0 or tuple(v) == (0, 0)

        ct = self.conv_transpose
        # The sum factorization is exact only for a "full" transposed conv.
        self._fast = (
            out_channels <= 1024
            and ct.groups == 1
            and _is_one(ct.stride)
            and _is_zero(ct.padding)
            and _is_zero(ct.output_padding)
            and _is_one(ct.dilation)
        )
        self._fn = fused_module.fused_forward

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and self._fast and x.dim() == 4:
            ct = self.conv_transpose
            return self._fn(
                x.contiguous(),
                ct.weight.contiguous(),
                ct.bias.contiguous(),
                self.bias.reshape(-1).contiguous(),
                self.kernel_h,
                self.kernel_w,
            )
        # Fallback to the reference implementation for non-CUDA / non-FP32 inputs.
        x = self.conv_transpose(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)
        x = x + self.bias
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = torch.sum(x, dim=(2, 3))
        x = x * 10.0
        return x
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 0.781 ms; the PyTorch reference (with TF32 enabled) takes 14.8 ms; speedup 18.950x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| fused_stream_wsum_kernel(float const*, float const*, float*, float*, long, long, long, int | 740.3 | 1 |
| matmul_bias_lse_kernel(float const*, float const*, float const*, float const*, float*, int | 8.0 | 1 |
| Memset (Device) | 1.3 | 1 |

The fastest correct solution so far (19.271x speedup) is:
```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

# Key algorithmic insight (unchanged, exact in FP32):
# For ConvTranspose2d with stride=1, padding=0 (a "full" convolution), the global
# sum of each output channel factorizes exactly:
#     sum_{h,w} out[n,c,h,w] = sum_k (sum_{h,w} x[n,k,h,w]) * (sum_{p,q} w[k,c,p,q])
#                              + conv_bias[c] * H_out * W_out
# So the whole network collapses to:
#   1) per-channel sum of the input            (memory-bound, streams ~1 GiB)
#   2) per-(in,out)-channel sum of the weights (tiny)
#   3) small matmul + bias adds + logsumexp + scale (fully fused)
#
# Improvements over the previous version:
#   * The dominant channel-sum kernel now splits each channel into several
#     chunks (8192 blocks instead of 1024). This removes the large wave-
#     quantization tail (1024 blocks ~= 2.4 waves on A100/H100) and pushes the
#     kernel much closer to peak DRAM bandwidth.
#   * Chunks are combined with atomicAdd, and the zero-init of the accumulator
#     buffer is fused into the tiny weight-sum kernel (removes one memset).
#   * Input-sum and weight-sum share one workspace allocation.

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math.h>
#include <cstdint>

#define FULL_MASK 0xffffffffu

// Warp + cross-warp block reduction (512 threads = 16 warps). Result valid on thread 0.
__device__ __forceinline__ float block_reduce(float s) {
    __shared__ float sh[16];
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_down_sync(FULL_MASK, s, off);
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = s;
    __syncthreads();
    float r = 0.0f;
    if (wid == 0) {
        r = (lane < 16) ? sh[lane] : 0.0f;
        for (int off = 8; off > 0; off >>= 1)
            r += __shfl_down_sync(FULL_MASK, r, off);
    }
    return r;
}

// Fused init: zero the per-channel input-sum accumulator AND compute the
// per-(in_channel, out_channel) sums of the K convolution taps.
__global__ void init_ws_kernel(const float* __restrict__ w,
                               float* __restrict__ Sw,
                               float* __restrict__ Sx,
                               int total, int NC, int K) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < NC) Sx[idx] = 0.0f;
    if (idx < total) {
        const float* ptr = w + (long)idx * K;
        float s = 0.0f;
        for (int i = 0; i < K; ++i) s += ptr[i];
        Sw[idx] = s;
    }
}

// Each block sums one contiguous chunk of one (batch, channel) plane using
// vectorized float4 streaming loads with 4x unrolling, then atomically
// accumulates the chunk result into the per-channel sum.
__global__ void __launch_bounds__(512)
channel_sum_kernel(const float* __restrict__ x,
                   float* __restrict__ out,
                   long HW, long chunk_floats, int chunks) {
    long channel = blockIdx.x / chunks;
    int cid = blockIdx.x % chunks;
    long start = (long)cid * chunk_floats;
    long len = HW - start;
    if (len > chunk_floats) len = chunk_floats;

    const float* p = x + channel * HW + start;
    long len4 = len >> 2;
    const float4* p4 = reinterpret_cast<const float4*>(p);

    const long stride = blockDim.x;
    const long limit4 = len4 - 3 * stride;

    float4 acc = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
    long i = threadIdx.x;
    for (; i < limit4; i += 4 * stride) {
        float4 a = __ldcs(p4 + i);
        float4 b = __ldcs(p4 + i + stride);
        float4 c = __ldcs(p4 + i + 2 * stride);
        float4 d = __ldcs(p4 + i + 3 * stride);
        acc.x += a.x + b.x + c.x + d.x;
        acc.y += a.y + b.y + c.y + d.y;
        acc.z += a.z + b.z + c.z + d.z;
        acc.w += a.w + b.w + c.w + d.w;
    }
    for (; i < len4; i += stride) {
        float4 a = __ldcs(p4 + i);
        acc.x += a.x; acc.y += a.y; acc.z += a.z; acc.w += a.w;
    }
    float s = (acc.x + acc.y) + (acc.z + acc.w);
    for (long j = (len4 << 2) + threadIdx.x; j < len; j += stride) {
        s += p[j];
    }

    s = block_reduce(s);
    if (threadIdx.x == 0) atomicAdd(out + channel, s);
}

// Scalar fallback for the (rare) case where planes are not 16-byte aligned.
__global__ void __launch_bounds__(512)
channel_sum_scalar_kernel(const float* __restrict__ x,
                          float* __restrict__ out,
                          long HW, long chunk_floats, int chunks) {
    long channel = blockIdx.x / chunks;
    int cid = blockIdx.x % chunks;
    long start = (long)cid * chunk_floats;
    long len = HW - start;
    if (len > chunk_floats) len = chunk_floats;

    const float* p = x + channel * HW + start;
    float s = 0.0f;
    for (long j = threadIdx.x; j < len; j += blockDim.x) {
        s += __ldcs(p + j);
    }
    s = block_reduce(s);
    if (threadIdx.x == 0) atomicAdd(out + channel, s);
}

// One block per batch element: fused matmul (Sx @ Sw), scale, bias adds,
// log-sum-exp over channels and final multiplication by 10.
__global__ void matmul_bias_lse_kernel(const float* __restrict__ Sx,
                                       const float* __restrict__ Sw,
                                       const float* __restrict__ conv_bias,
                                       const float* __restrict__ bias,
                                       float* __restrict__ out,
                                       int C_in, int C_out,
                                       float inv_hw) {
    int n = blockIdx.x;
    int c = threadIdx.x;
    bool valid = c < C_out;

    float y = 0.0f;
    if (valid) {
        float acc = 0.0f;
        const float* sx = Sx + n * C_in;
        const float* sw = Sw + c;
        #pragma unroll 4
        for (int k = 0; k < C_in; ++k) {
            acc += sx[k] * sw[k * C_out];
        }
        y = acc * inv_hw + conv_bias[c] + bias[c];
    }

    __shared__ float sh[1024];
    // Block reduction: max over channels.
    sh[c] = valid ? y : -3.402823466e38f;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (c < off) sh[c] = fmaxf(sh[c], sh[c + off]);
        __syncthreads();
    }
    float m = sh[0];
    __syncthreads();
    // Block reduction: sum of exp(y - max).
    sh[c] = valid ? expf(y - m) : 0.0f;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (c < off) sh[c] += sh[c + off];
        __syncthreads();
    }
    if (c == 0) {
        out[n] = (m + logf(sh[0])) * 10.0f;
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight,
                            torch::Tensor conv_bias, torch::Tensor bias,
                            int64_t kernel_h, int64_t kernel_w) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "fused_forward expects a CUDA float32 input");
    int64_t N = x.size(0);
    int64_t C_in = x.size(1);
    int64_t H = x.size(2);
    int64_t W = x.size(3);
    int64_t C_out = weight.size(1);
    TORCH_CHECK(C_out <= 1024, "C_out must be <= 1024");

    int64_t H_out = H + kernel_h - 1;
    int64_t W_out = W + kernel_w - 1;
    float inv_hw = 1.0f / (float)(H_out * W_out);

    // Single workspace: Sx (N*C_in) followed by Sw (C_in*C_out).
    auto workspace = torch::empty({N * C_in + C_in * C_out}, x.options());
    float* Sx = workspace.data_ptr<float>();
    float* Sw = Sx + N * C_in;
    auto out = torch::empty({N, 1}, x.options());

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

    int total = (int)(C_in * C_out);
    int NC = (int)(N * C_in);
    int K = (int)(kernel_h * kernel_w);
    int init_threads = 256;
    int init_max = total > NC ? total : NC;
    init_ws_kernel<<<(init_max + init_threads - 1) / init_threads, init_threads, 0, stream>>>(
        weight.data_ptr<float>(), Sw, Sx, total, NC, K);

    long HW = (long)H * (long)W;
    // Split each channel plane into chunks so the grid has many blocks; this
    // removes wave-quantization tails and keeps every SM fed.
    const long CHUNK = 32768;  // floats (128 KB) per block
    long chunks = (HW + CHUNK - 1) / CHUNK;
    if (chunks * NC < 2048) {
        long want = (2048 + NC - 1) / NC;
        long max_chunks = (HW + 2047) / 2048;
        if (want > max_chunks) want = max_chunks;
        if (want > chunks) chunks = want;
        if (chunks < 1) chunks = 1;
    }
    long chunk_floats = (HW + chunks - 1) / chunks;
    chunk_floats = ((chunk_floats + 3) / 4) * 4;  // keep float4 alignment
    chunks = (HW + chunk_floats - 1) / chunk_floats;

    int grid = (int)(chunks * NC);
    const float* xptr = x.data_ptr<float>();
    bool vec = ((HW & 3) == 0) &&
               ((reinterpret_cast<uintptr_t>(xptr) & 15) == 0);
    if (vec) {
        channel_sum_kernel<<<grid, 512, 0, stream>>>(
            xptr, Sx, HW, chunk_floats, (int)chunks);
    } else {
        channel_sum_scalar_kernel<<<grid, 512, 0, stream>>>(
            xptr, Sx, HW, chunk_floats, (int)chunks);
    }

    int block = 1;
    while (block < C_out && block < 1024) block <<= 1;
    matmul_bias_lse_kernel<<<(int)N, block, 0, stream>>>(
        Sx, Sw, conv_bias.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), (int)C_in, (int)C_out, inv_hw);

    return out;
}
"""

cpp_source = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor conv_bias, torch::Tensor bias, int64_t kernel_h, int64_t kernel_w);"
)

fused_module = load_inline(
    name="fused_convtranspose_gap_lse_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Replaces the expensive transposed convolution +
    global average pooling + bias + logsumexp + sum + scale pipeline with three
    fused custom CUDA kernels, exploiting the exact factorization of the global
    sum of a full (stride-1) transposed convolution. The dominant input
    reduction kernel streams the input with vectorized float4 loads over a
    finely chunked grid (atomically combining chunk partial sums) to reach
    near-peak DRAM bandwidth. All math is done in FP32.
    """

    def __init__(self, in_channels, out_channels, kernel_size, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        if isinstance(kernel_size, (tuple, list)):
            self.kernel_h = int(kernel_size[0])
            self.kernel_w = int(kernel_size[1])
        else:
            self.kernel_h = int(kernel_size)
            self.kernel_w = int(kernel_size)

    def forward(self, x):
        if (x.is_cuda and x.dtype == torch.float32
                and self.conv_transpose.out_channels <= 1024):
            return fused_module.fused_forward(
                x.contiguous(),
                self.conv_transpose.weight.contiguous(),
                self.conv_transpose.bias.contiguous(),
                self.bias.reshape(-1).contiguous(),
                self.kernel_h,
                self.kernel_w,
            )
        # Fallback to the reference implementation for non-CUDA / non-FP32 inputs.
        x = self.conv_transpose(x)
        x = torch.mean(x, dim=(2, 3), keepdim=True)
        x = x + self.bias
        x = torch.logsumexp(x, dim=1, keepdim=True)
        x = torch.sum(x, dim=(2, 3))
        x = x * 10.0
        return x
```

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.


Precision rule: the timed forward must compute in FP32 (TF32 tensor-core math is allowed, as in the PyTorch reference). Do not cast inputs, weights or intermediate results to float16 or bfloat16 and do not call half-precision kernels in forward().
