```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor pooled_conv_sum_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t pool_h,
    int64_t pool_w);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>

constexpr int THREADS = 256;
constexpr int CHUNK = 4096;

// Pool each convolution input patch before performing the channel mixing.
// The resulting matrix has shape [batch, in_channels * kh * kw, pitch].
__global__ void lower_pool_3x3_4x4(
    const float* __restrict__ x,
    float* __restrict__ lowered,
    int channels, int height, int width,
    int out_width, int positions, int pitch)
{
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    const int channel = blockIdx.y;
    const int batch = blockIdx.z;
    if (p >= pitch)
        return;

    float* dst = lowered +
        (static_cast<int64_t>(batch) * channels + channel) * 9 * pitch + p;

    if (p >= positions) {
        #pragma unroll
        for (int k = 0; k < 9; ++k)
            dst[static_cast<int64_t>(k) * pitch] = 0.0f;
        return;
    }

    const int iy = (p / out_width) * 4;
    const int ix = (p % out_width) * 4;
    const float* src = x +
        ((static_cast<int64_t>(batch) * channels + channel) * height + iy)
        * width + ix;

    float horizontal[6][3];

    #pragma unroll
    for (int r = 0; r < 6; ++r) {
        const float a = src[r * width];
        const float b = src[r * width + 1];
        const float c = src[r * width + 2];
        const float d = src[r * width + 3];
        const float e = src[r * width + 4];
        const float f = src[r * width + 5];

        horizontal[r][0] = ((a + b) + c) + d;
        horizontal[r][1] = ((b + c) + d) + e;
        horizontal[r][2] = ((c + d) + e) + f;
    }

    #pragma unroll
    for (int r = 0; r < 3; ++r) {
        #pragma unroll
        for (int c = 0; c < 3; ++c) {
            const float value =
                ((horizontal[r][c] + horizontal[r + 1][c]) +
                  horizontal[r + 2][c]) + horizontal[r + 3][c];
            dst[static_cast<int64_t>(r * 3 + c) * pitch] =
                value * 0.0625f;
        }
    }
}

__global__ void lower_pool_generic(
    const float* __restrict__ x,
    float* __restrict__ lowered,
    int channels, int height, int width,
    int kh, int kw, int ph, int pw,
    int out_width, int positions, int pitch)
{
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    const int channel = blockIdx.y;
    const int batch = blockIdx.z;
    if (p >= pitch)
        return;

    float* dst = lowered +
        (static_cast<int64_t>(batch) * channels + channel) * kh * kw * pitch
        + p;

    const int iy = (p / out_width) * ph;
    const int ix = (p % out_width) * pw;
    const float scale = 1.0f / static_cast<float>(ph * pw);

    for (int r = 0; r < kh; ++r) {
        for (int c = 0; c < kw; ++c) {
            float value = 0.0f;
            if (p < positions) {
                const float* src = x +
                    ((static_cast<int64_t>(batch) * channels + channel)
                     * height + iy + r) * width + ix + c;

                for (int pr = 0; pr < ph; ++pr)
                    for (int pc = 0; pc < pw; ++pc)
                        value += src[pr * width + pc];
            }
            dst[static_cast<int64_t>(r * kw + c) * pitch] = value * scale;
        }
    }
}

__device__ __forceinline__ float warp_sum(float value)
{
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        value += __shfl_down_sync(0xffffffff, value, offset);
    return value;
}

__device__ __forceinline__ float block_sum(float value)
{
    __shared__ float sums[8];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    value = warp_sum(value);
    if (lane == 0)
        sums[warp] = value;
    __syncthreads();

    value = threadIdx.x < 8 ? sums[lane] : 0.0f;
    if (warp == 0)
        value = warp_sum(value);
    return value;
}

__global__ void bias_sigmoid_partial(
    const float* __restrict__ input,
    const float* __restrict__ bias,
    float* __restrict__ partial,
    int positions, int pitch, int channels, int parts)
{
    const int batch = blockIdx.y;
    const int part = blockIdx.x;
    const int64_t sample_size = static_cast<int64_t>(channels) * pitch;
    const int64_t begin = static_cast<int64_t>(part) * CHUNK;
    const float* src = input + static_cast<int64_t>(batch) * sample_size;

    float total = 0.0f;
    #pragma unroll
    for (int j = 0; j < CHUNK / THREADS; ++j) {
        const int64_t i = begin + threadIdx.x + j * THREADS;
        if (i < sample_size) {
            const int channel = i / pitch;
            const int p = i - static_cast<int64_t>(channel) * pitch;
            if (p < positions) {
                const float value = src[i] + bias[channel];
                total += 1.0f / (1.0f + expf(-value));
            }
        }
    }

    total = block_sum(total);
    if (threadIdx.x == 0)
        partial[static_cast<int64_t>(batch) * parts + part] = total;
}

__global__ void finish_sum(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int parts)
{
    const int batch = blockIdx.x;
    float total = 0.0f;
    for (int i = threadIdx.x; i < parts; i += blockDim.x)
        total += partial[static_cast<int64_t>(batch) * parts + i];

    total = block_sum(total);
    if (threadIdx.x == 0)
        output[batch] = total;
}

torch::Tensor pooled_conv_sum_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t pool_h,
    int64_t pool_w)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
                "All tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat &&
                weight.scalar_type() == at::kFloat &&
                bias.scalar_type() == at::kFloat,
                "Expected float32 tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous() &&
                bias.is_contiguous(), "Expected contiguous tensors");
    TORCH_CHECK(x.dim() == 4 && weight.dim() == 4,
                "Expected NCHW input and OIHW weights");
    TORCH_CHECK(x.device() == weight.device() &&
                x.device() == bias.device(), "Device mismatch");
    TORCH_CHECK(pool_h > 0 && pool_w > 0, "Invalid pooling size");

    const c10::cuda::CUDAGuard guard(x.device());

    const int batch = x.size(0);
    const int channels = x.size(1);
    const int height = x.size(2);
    const int width = x.size(3);
    const int outputs = weight.size(0);
    const int kh = weight.size(2);
    const int kw = weight.size(3);
    const int ph = pool_h;
    const int pw = pool_w;

    TORCH_CHECK(weight.size(1) == channels && bias.numel() == outputs,
                "Channel mismatch");
    TORCH_CHECK(height >= kh + ph - 1 && width >= kw + pw - 1,
                "Input is too small");

    auto result = torch::empty({batch}, x.options());
    if (batch == 0)
        return result;

    const int oh = (height - kh + 1) / ph;
    const int ow = (width - kw + 1) / pw;
    const int positions = oh * ow;
    const int pitch = ((positions + 31) / 32) * 32;
    const int reduction_size = channels * kh * kw;

    auto lowered = torch::empty(
        {batch, reduction_size, pitch}, x.options());
    auto products = torch::empty(
        {batch, outputs, pitch}, x.options());

    const auto stream = at::cuda::getCurrentCUDAStream();
    const dim3 grid((pitch + THREADS - 1) / THREADS, channels, batch);

    if (kh == 3 && kw == 3 && ph == 4 && pw == 4) {
        lower_pool_3x3_4x4<<<grid, THREADS, 0, stream>>>(
            x.data_ptr<float>(), lowered.data_ptr<float>(),
            channels, height, width, ow, positions, pitch);
    } else {
        lower_pool_generic<<<grid, THREADS, 0, stream>>>(
            x.data_ptr<float>(), lowered.data_ptr<float>(),
            channels, height, width, kh, kw, ph, pw,
            ow, positions, pitch);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    // Column-major GEMM:
    // [pitch, reduction_size] @ [reduction_size, outputs].
    // The weights are shared across all batch elements.
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    const float alpha = 1.0f;
    const float beta = 0.0f;
    const auto compute_type =
        at::cuda::getCurrentDeviceProperties()->major >= 8
        ? CUBLAS_COMPUTE_32F_FAST_TF32
        : CUBLAS_COMPUTE_32F;

    const cublasStatus_t status = cublasGemmStridedBatchedEx(
        handle,
        CUBLAS_OP_N, CUBLAS_OP_N,
        pitch, outputs, reduction_size,
        &alpha,
        lowered.data_ptr<float>(), CUDA_R_32F, pitch,
        static_cast<long long>(pitch) * reduction_size,
        weight.data_ptr<float>(), CUDA_R_32F, reduction_size,
        0,
        &beta,
        products.data_ptr<float>(), CUDA_R_32F, pitch,
        static_cast<long long>(pitch) * outputs,
        batch,
        compute_type,
        CUBLAS_GEMM_DEFAULT);

    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
                "cublasGemmStridedBatchedEx failed: ",
                static_cast<int>(status));

    const int64_t sample_size = static_cast<int64_t>(outputs) * pitch;
    const int parts = (sample_size + CHUNK - 1) / CHUNK;
    auto partial = torch::empty({batch, parts}, x.options());

    bias_sigmoid_partial<<<dim3(parts, batch), THREADS, 0, stream>>>(
        products.data_ptr<float>(), bias.data_ptr<float>(),
        partial.data_ptr<float>(), positions, pitch, outputs, parts);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_sum<<<batch, THREADS, 0, stream>>>(
        partial.data_ptr<float>(), result.data_ptr<float>(), parts);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return result;
}
"""

_ops = load_inline(
    name="pooled_patch_conv_sigmoid_sum_v2",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["pooled_conv_sum_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, pool_kernel_size):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
        self.avg_pool = nn.AvgPool2d(pool_kernel_size)
        self.pool_size = (
            (pool_kernel_size, pool_kernel_size)
            if isinstance(pool_kernel_size, int)
            else tuple(pool_kernel_size)
        )

    def forward(self, x):
        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or self.conv.weight.dtype != torch.float32
            or self.conv.bias.dtype != torch.float32
            or torch.is_autocast_enabled()
        ):
            return torch.sigmoid(self.avg_pool(self.conv(x))).sum(
                dim=(1, 2, 3)
            )

        return _ops.pooled_conv_sum_cuda(
            x.contiguous(),
            self.conv.weight.contiguous(),
            self.conv.bias.contiguous(),
            self.pool_size[0],
            self.pool_size[1],
        )
```