```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_columns_cuda(
    torch::Tensor x, int64_t kh, int64_t kw, int64_t ph, int64_t pw);

torch::Tensor sigmoid_sum_cuda(
    torch::Tensor x, torch::Tensor bias);
"""

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

__global__ void pooled_columns_3x3_4x4(
    const float* __restrict__ x,
    float* __restrict__ columns,
    int64_t total,
    int channels,
    int height,
    int width,
    int out_h,
    int out_w)
{
    const int spatial = out_h * out_w;

    for (int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         index < total;
         index += int64_t(blockDim.x) * gridDim.x)
    {
        const int position = index % spatial;
        const int64_t nc = index / spatial;
        const int py = position / out_w;
        const int px = position % out_w;

        const float* input =
            x + nc * height * width + int64_t(py * 4) * width + px * 4;

        float horizontal[6][3];

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

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

        #pragma unroll
        for (int ky = 0; ky < 3; ++ky) {
            #pragma unroll
            for (int kx = 0; kx < 3; ++kx) {
                const float value =
                    ((horizontal[ky][kx] + horizontal[ky + 1][kx])
                     + horizontal[ky + 2][kx])
                    + horizontal[ky + 3][kx];

                columns[(nc * 9 + ky * 3 + kx) * spatial + position] =
                    value * 0.0625f;
            }
        }
    }
}

__global__ void pooled_columns_generic(
    const float* __restrict__ x,
    float* __restrict__ columns,
    int64_t total,
    int channels,
    int height,
    int width,
    int out_h,
    int out_w,
    int kh,
    int kw,
    int ph,
    int pw)
{
    const int spatial = out_h * out_w;
    const int kernel_area = kh * kw;

    for (int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         index < total;
         index += int64_t(blockDim.x) * gridDim.x)
    {
        const int position = index % spatial;
        const int64_t column = index / spatial;
        const int offset = column % kernel_area;
        const int64_t nc = column / kernel_area;

        const int y = (position / out_w) * ph + offset / kw;
        const int xx = (position % out_w) * pw + offset % kw;

        const float* input =
            x + nc * height * width + int64_t(y) * width + xx;

        float sum = 0.0f;
        for (int iy = 0; iy < ph; ++iy) {
            for (int ix = 0; ix < pw; ++ix) {
                sum += input[int64_t(iy) * width + ix];
            }
        }

        columns[index] = sum / float(ph * pw);
    }
}

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

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

    if (lane == 0) {
        warp_sums[warp] = value;
    }
    __syncthreads();

    value = threadIdx.x < 8 ? warp_sums[lane] : 0.0f;

    if (warp == 0) {
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            value += __shfl_down_sync(0xffffffff, value, offset);
        }
    }
    return value;
}

__global__ void sigmoid_sum_partial(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    float* __restrict__ partial,
    int64_t elements,
    int64_t spatial,
    int chunks)
{
    const int batch = blockIdx.y;
    const int chunk = blockIdx.x;
    const int64_t start = int64_t(chunk) * 2048 + threadIdx.x;
    const float* input = x + int64_t(batch) * elements;

    float sum = 0.0f;

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const int64_t index = start + i * 256;
        if (index < elements) {
            const float value = input[index] + bias[index / spatial];
            sum += 1.0f / (1.0f + __expf(-value));
        }
    }

    sum = block_sum(sum);
    if (threadIdx.x == 0) {
        partial[int64_t(batch) * chunks + chunk] = sum;
    }
}

__global__ void finish_sum(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int chunks)
{
    float sum = 0.0f;
    for (int i = threadIdx.x; i < chunks; i += blockDim.x) {
        sum += partial[int64_t(blockIdx.x) * chunks + i];
    }

    sum = block_sum(sum);
    if (threadIdx.x == 0) {
        output[blockIdx.x] = sum;
    }
}

torch::Tensor pooled_columns_cuda(
    torch::Tensor x, int64_t kh, int64_t kw, int64_t ph, int64_t pw)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32);
    TORCH_CHECK(x.dim() == 4 && x.is_contiguous());
    TORCH_CHECK(kh > 0 && kw > 0 && ph > 0 && pw > 0);

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

    const int64_t batch = x.size(0);
    const int64_t channels = x.size(1);
    const int64_t height = x.size(2);
    const int64_t width = x.size(3);
    const int64_t out_h = (height - kh + 1) / ph;
    const int64_t out_w = (width - kw + 1) / pw;

    TORCH_CHECK(height >= kh && width >= kw && out_h > 0 && out_w > 0);

    auto columns = torch::empty(
        {batch, channels * kh * kw, out_h * out_w}, x.options());

    if (batch == 0) {
        return columns;
    }

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

    if (kh == 3 && kw == 3 && ph == 4 && pw == 4) {
        const int64_t total = batch * channels * out_h * out_w;
        const int blocks = int(std::min<int64_t>((total + 255) / 256, 65535));

        pooled_columns_3x3_4x4<<<blocks, 256, 0, stream>>>(
            x.data_ptr<float>(), columns.data_ptr<float>(), total,
            channels, height, width, out_h, out_w);
    } else {
        const int64_t total = columns.numel();
        const int blocks = int(std::min<int64_t>((total + 255) / 256, 65535));

        pooled_columns_generic<<<blocks, 256, 0, stream>>>(
            x.data_ptr<float>(), columns.data_ptr<float>(), total,
            channels, height, width, out_h, out_w, kh, kw, ph, pw);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return columns;
}

torch::Tensor sigmoid_sum_cuda(torch::Tensor x, torch::Tensor bias)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32);
    TORCH_CHECK(x.dim() == 3 && x.is_contiguous());
    TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kFloat32);
    TORCH_CHECK(bias.is_contiguous() && bias.numel() == x.size(1));
    TORCH_CHECK(x.device() == bias.device());
    TORCH_CHECK(x.size(0) <= 65535);

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

    const int batch = x.size(0);
    const int64_t spatial = x.size(2);
    const int64_t elements = x.size(1) * spatial;
    const int chunks = int((elements + 2047) / 2048);

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

    auto partial = torch::empty({batch, chunks}, x.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    sigmoid_sum_partial<<<dim3(chunks, batch), 256, 0, stream>>>(
        x.data_ptr<float>(), bias.data_ptr<float>(),
        partial.data_ptr<float>(), elements, spatial, chunks);

    C10_CUDA_KERNEL_LAUNCH_CHECK();

    finish_sum<<<batch, 256, 0, stream>>>(
        partial.data_ptr<float>(), output.data_ptr<float>(), chunks);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_ops = load_inline(
    name="pooled_conv_sigmoid_sum_cuda_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["pooled_columns_cuda", "sigmoid_sum_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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)

    def forward(self, x):
        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or self.conv.weight.dtype != torch.float32
            or torch.is_autocast_enabled()
            or x.shape[0] > 65535
        ):
            return torch.sigmoid(self.avg_pool(self.conv(x))).sum((1, 2, 3))

        kh, kw = self.conv.kernel_size
        pool = self.avg_pool.kernel_size
        ph, pw = (pool, pool) if isinstance(pool, int) else pool

        # Average pooling commutes with the linear convolution.
        # Construct averaged patches directly at the pooled resolution.
        columns = _ops.pooled_columns_cuda(
            x.contiguous(), kh, kw, ph, pw
        )
        weight = self.conv.weight.reshape(self.conv.out_channels, -1)
        values = torch.matmul(weight, columns)

        return _ops.sigmoid_sum_cuda(values, self.conv.bias.contiguous())
```