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

torch.backends.cudnn.benchmark = True

# ---------------------------------------------------------------------------
# Key algorithmic idea (exact, not an approximation):
#   sum_{p in 4x4 window} conv[p, c]  ==  sum_k  W[c,k] * ( sum_{p} A[p,k] )
# because the 4x4 average pool is linear.  So we first compute the "reduced"
# feature map  Ab[w, tap*8+ci] = sum over the 4x4 window of the input patch
# (a 9-phase 4x4 sum-pool), then a small GEMM (M = #windows, K = 72, N = 64)
# on tensor cores, then sigmoid + per-image sum in the epilogue.
# This removes a factor 16 of conv work and never materializes the 2.4 GB
# conv output.  Everything is fused into ONE kernel.
# ---------------------------------------------------------------------------

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

namespace {

#define RW 34
#define CW 66
#define RSTRIDE 68
#define AWSTRIDE 88
#define WSM_STRIDE 82

__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, const uint32_t* b) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
#endif
}

__device__ __forceinline__ float sigf(float v) {
    return 1.0f / (1.0f + expf(-v));
}

// One block: 8x16 pooling windows (one image), i.e. input region 34x66 x 8ch.
__global__ void __launch_bounds__(256)
fused_reduced_conv_pool_sigmoid_sum_kernel(
        const float* __restrict__ x,      // [N,8,H,W] fp32
        const __half* __restrict__ wk,    // [64,80] fp16
        const float* __restrict__ bias,   // [64] fp32
        float* __restrict__ out,          // [N] fp32
        const int H, const int W,
        const int Hp, const int Wp,
        const int tilesX, const int tilesY) {
    __shared__ __half smem_in[RW * RSTRIDE * 8];   // [row][col][ci], col padded to 68
    __shared__ __half smem_a[128 * AWSTRIDE];      // [m][k], k padded to 88
    __shared__ __half smem_w[64 * WSM_STRIDE];
    __shared__ float red[64];

    const int t  = threadIdx.x;
    const int n  = blockIdx.z;
    const int ty = blockIdx.y;
    const int tx = blockIdx.x;
    const int wy0 = ty << 3;
    const int wx0 = tx << 4;

    // ---------------- weights -> smem ----------------
    {
        const uint2* wsrc = reinterpret_cast<const uint2*>(wk);
        uint32_t* wdst = reinterpret_cast<uint32_t*>(smem_w);
        #pragma unroll
        for (int i = t; i < 1280; i += 256) {
            const uint2 v = wsrc[i];
            const int h = 4 * i;
            const int row = h / 80;
            const int k = h - row * 80;
            const int d = row * (WSM_STRIDE / 2) + (k >> 1);
            wdst[d] = v.x;
            wdst[d + 1] = v.y;
        }
    }

    // ---------------- input region -> smem (fp16) ----------------
    {
        const size_t area = (size_t)H * (size_t)W;
        const float* xb = x + (size_t)n * 8 * area;
        uint32_t* dst = reinterpret_cast<uint32_t*>(smem_in);
        const int gy0 = wy0 << 2;
        const int gx0 = wx0 << 2;
        #pragma unroll
        for (int i = t; i < 4488; i += 256) {
            const int q = i & 3;
            const int pix = i >> 2;
            const int dx = pix % RW;      // note: CW(66) = 2*RW(34)
            const int dy = pix / RW;      // 0..65
            const int r  = dy >> 1;       // 0..32
            const int c  = (dy & 1) * RW + dx;  // 0..65
            const int gy = gy0 + r;
            const int gx = gx0 + c;
            float2 v = make_float2(0.f, 0.f);
            if (gy < H && gx < W) {
                const size_t off = (size_t)gy * W + gx;
                v.x = xb[(size_t)(2 * q) * area + off];
                v.y = xb[(size_t)(2 * q + 1) * area + off];
            }
            __half2 h2 = __floats2half2_rn(v.x, v.y);
            dst[(r * RSTRIDE + c) * 4 + q] = *reinterpret_cast<uint32_t*>(&h2);
        }
    }
    __syncthreads();

    // ---------------- reduced map Ab[m][k] (fp16, fp32 sums) ----------------
    {
        const int ci = t & 7;
        const int wl = t >> 3;                 // 0..31
        const int wy = wy0 + (wl >> 2);
        const int wx = wx0 + (wl & 3);
        const bool wvalid = (wy < Hp) && (wx < Wp);
        const int lwy = wl >> 2;
        const int lwx = wl & 3;
        #pragma unroll
        for (int rep = 0; rep < 4; ++rep) {
            const int xw = (wl & 3) + 4 * rep; // 0..15
            const int rbase = lwy << 2;
            const int cbase = xw << 2;
            float p0[6], p1[6], p2[6];
            uint32_t rv;
            const __half2* prow = &smem_in[(rbase * RSTRIDE + cbase) * 8 + ci];
            #define COLSUM(dst, rowoff)                                       \
                { const __half2* pp = prow + (rowoff) * (RSTRIDE * 8);        \
                  _Pragma("unroll")                                           \
                  for (int j = 0; j < 6; ++j) {                               \
                      const float2 f = __half22float2(pp[j * 8]);             \
                      dst[j] = f.x + f.y;                                     \
                  } }
            COLSUM(p0, 0)
            COLSUM(p1, 1)
            COLSUM(p2, 2)
            #pragma unroll
            for (int j = 0; j < 6; ++j) { p0[j] += p1[j]; p2[j] += p1[j]; }
            COLSUM(p1, 3)
            #pragma unroll
            for (int j = 0; j < 6; ++j) { p0[j] += p1[j]; p2[j] += p1[j]; }
            COLSUM(p1, 4)
            #pragma unroll
            for (int j = 0; j < 6; ++j) p2[j] += p1[j];
            __half* aw = &smem_a[wl * AWSTRIDE];
            const float qv0 = (p0[0]+p0[1])+(p0[2]+p0[3]);
            const float qv1 = (p0[1]+p0[2])+(p0[3]+p0[4]);
            const float qv2 = (p0[2]+p0[3])+(p0[4]+p0[5]);
            rv = wvalid ? __half2_as_u32(__floats2half2_rn(qv0, qv1)) : 0u;
            reinterpret_cast<uint32_t*>(aw)[0] = rv;
            aw[2] = __float2half_rn(wvalid ? qv2 : 0.f);
            const float rv0 = (p2[0]+p2[1])+(p2[2]+p2[3]);
            const float rv1 = (p2[1]+p2[2])+(p2[3]+p2[4]);
            const float rv2 = (p2[2]+p2[3])+(p2[4]+p2[5]);
            rv = wvalid ? __half2_as_u32(__floats2half2_rn(rv0, rv1)) : 0u;
            reinterpret_cast<uint32_t*>(aw)[8] = rv;
            aw[18] = __float2half_rn(wvalid ? rv2 : 0.f);
            // middle tap row (r=1)
            COLSUM(p1, 1)
            #pragma unroll
            for (int j = 0; j < 6; ++j) p0[j] = p1[j];
            COLSUM(p1, 2)
            #pragma unroll
            for (int j = 0; j < 6; ++j) p0[j] += p1[j];
            COLSUM(p1, 3)
            #pragma unroll
            for (int j = 0; j < 6; ++j) p0[j] += p1[j];
            COLSUM(p1, 4)
            #pragma unroll
            for (int j = 0; j < 6; ++j) p0[j] += p1[j];
            const float mv0 = (p0[0]+p0[1])+(p0[2]+p0[3]);
            const float mv1 = (p0[1]+p0[2])+(p0[3]+p0[4]);
            const float mv2 = (p0[2]+p0[3])+(p0[4]+p0[5]);
            rv = wvalid ? __half2_as_u32(__floats2half2_rn(mv0, mv1)) : 0u;
            reinterpret_cast<uint32_t*>(aw)[4] = rv;
            aw[10] = __float2half_rn(wvalid ? mv2 : 0.f);
            // zero the 8 pad columns
            reinterpret_cast<uint32_t*>(aw)[36] = 0u;
            reinterpret_cast<uint32_t*>(aw)[37] = 0u;
            reinterpret_cast<uint32_t*>(aw)[38] = 0u;
            reinterpret_cast<uint32_t*>(aw)[39] = 0u;
            #undef COLSUM
        }
    }
    __syncthreads();

    // ---------------- GEMM: S = Ab (128x80) * W^T (80x64) ----------------
    const int warp = t >> 5;
    const int lane = t & 31;
    const int g = lane >> 2;
    const int tig = lane & 3;

    float acc[8][4];
    #pragma unroll
    for (int i = 0; i < 8; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j) acc[i][j] = 0.f;

    const uint32_t* a_u32 = reinterpret_cast<const uint32_t*>(smem_a);
    #pragma unroll
    for (int ch = 0; ch < 5; ++ch) {
        uint32_t A[4];
        A[0] = a_u32[(warp * 16 + g)     * (AWSTRIDE / 2) + ch * 8 + tig];
        A[1] = a_u32[(warp * 16 + g + 8) * (AWSTRIDE / 2) + ch * 8 + tig];
        A[2] = a_u32[(warp * 16 + g)     * (AWSTRIDE / 2) + ch * 8 + tig + 4];
        A[3] = a_u32[(warp * 16 + g + 8) * (AWSTRIDE / 2) + ch * 8 + tig + 4];
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            const uint32_t* wp = reinterpret_cast<const uint32_t*>(
                &smem_w[(nt * 8 + g) * WSM_STRIDE + ch * 16]);
            uint32_t B[2];
            B[0] = wp[tig];
            B[1] = wp[tig + 4];
            mma16816(acc[nt], A, B);
        }
    }

    // -------- epilogue: sigmoid((S+bias)/16), reduce over windows --------
    float csum[2] = {0.f, 0.f};
    #pragma unroll
    for (int nt = 0; nt < 8; ++nt) {
        float* c = acc[nt];
        const float b0 = bias[nt * 8 + 2 * tig] * 0.0625f;
        const float b1 = bias[nt * 8 + 2 * tig + 1] * 0.0625f;
        csum[0] += sigf(c[0] * 0.0625f + b0) + sigf(c[2] * 0.0625f + b0);
        csum[1] += sigf(c[1] * 0.0625f + b1) + sigf(c[3] * 0.0625f + b1);
    }
    #pragma unroll
    for (int off = 4; off < 32; off <<= 1) {
        csum[0] += __shfl_xor_sync(0xffffffffu, csum[0], off);
        csum[1] += __shfl_xor_sync(0xffffffffu, csum[1], off);
    }
    // lane = g*4 + tig ; lanes 4g and 4g+1 hold channel 2g (rows g and g+8)
    if ((lane & 3) < 2) {
        const int chid = (lane >> 1) & 31;   // 0..31 per half-warp
        atomicAdd(&red[chid + ((lane & 2) ? 32 : 0)], (lane & 1) ? csum[1] : csum[0]);
    }
    __syncthreads();
    if (t < 64) {
        const float v = red[t];
        if (v != 0.0f) atomicAdd(&out[n], v);
    }
}

// ------------------------- backup kernels (v4 path) -------------------------

#define TILE_PIX 16
#define PATCH_DIM 18
#define PATCH_STRIDE 20
#define KPAD 80

__global__ void __launch_bounds__(256)
fused_conv3x3_c8_pool4_sigmoid_sum_kernel(
        const float* __restrict__ x,
        const __half* __restrict__ wk,
        const float* __restrict__ bias,
        float* __restrict__ out,
        const int H, const int W, const int Hp, const int Wp) {
    __shared__ __half smem_in[PATCH_DIM * PATCH_STRIDE * 8];
    __shared__ __half smem_w[64 * WSM_STRIDE];
    __shared__ float smem_bias[64];
    __shared__ float red[8];

    const int t = threadIdx.x;
    const int n = blockIdx.z;
    const int py0 = blockIdx.y * TILE_PIX;
    const int px0 = blockIdx.x * TILE_PIX;

    {
        const uint2* wsrc = reinterpret_cast<const uint2*>(wk);
        uint32_t* wdst = reinterpret_cast<uint32_t*>(smem_w);
        #pragma unroll
        for (int i = t; i < 1280; i += 256) {
            const uint2 v = wsrc[i];
            const int h = 4 * i;
            const int row = h / 80;
            const int k = h - row * 80;
            const int d = row * (WSM_STRIDE / 2) + (k >> 1);
            wdst[d] = v.x;
            wdst[d + 1] = v.y;
        }
        if (t < 64) smem_bias[t] = bias[t];
    }
    {
        const int area = H * W;
        const float* xb = x + (size_t)n * 8 * (size_t)area;
        uint32_t* dst = reinterpret_cast<uint32_t*>(smem_in);
        #pragma unroll
        for (int i = t; i < 1296; i += 256) {
            const int q = i & 3;
            const int pix = i >> 2;
            const int dx = pix % PATCH_DIM;
            const int dy = pix / PATCH_DIM;
            const int gy = py0 + dy;
            const int gx = px0 + dx;
            float2 v = make_float2(0.f, 0.f);
            if (gy < H && gx < W) {
                const size_t off = (size_t)gy * W + gx;
                v.x = xb[(size_t)(2 * q) * area + off];
                v.y = xb[(size_t)(2 * q + 1) * area + off];
            }
            __half2 h2 = __floats2half2_rn(v.x, v.y);
            dst[(dy * PATCH_STRIDE + dx) * 4 + q] =
                *reinterpret_cast<uint32_t*>(&h2);
        }
    }
    __syncthreads();

    const int warp = t >> 5;
    const int lane = t & 31;
    const int g = lane >> 2;
    const int tig = lane & 3;

    const int wslot0 = warp << 1;
    const int w0y = (wslot0 >> 2) << 2;
    const int w0x = (wslot0 & 3) << 2;

    const int gwy = (blockIdx.y << 2) + (wslot0 >> 2);
    const int gwx = (blockIdx.x << 2) + (wslot0 & 3);
    const bool valid0 = (gwy < Hp) && (gwx < Wp);
    const bool valid1 = (gwy < Hp) && (gwx + 1 < Wp);

    float acc[2][8][4];
    #pragma unroll
    for (int w = 0; w < 2; ++w)
        #pragma unroll
        for (int i = 0; i < 8; ++i)
            #pragma unroll
            for (int j = 0; j < 4; ++j) acc[w][i][j] = 0.f;

    const uint32_t* in_u32 = reinterpret_cast<const uint32_t*>(smem_in);
    const int ya = w0y + (g >> 2);
    const int yb = ya + 2;
    const int xa0 = w0x + (g & 3);
    const int xa1 = xa0 + 4;

    #pragma unroll
    for (int ch = 0; ch < 5; ++ch) {
        const int tap0 = 2 * ch;
        int tap1 = 2 * ch + 1;
        if (tap1 > 8) tap1 = 8;
        const int r0 = tap0 / 3, q0 = tap0 % 3;
        const int r1 = tap1 / 3, q1 = tap1 % 3;

        uint32_t A[2][4];
        A[0][0] = in_u32[((ya + r0) * PATCH_STRIDE + (xa0 + q0)) * 4 + tig];
        A[0][1] = in_u32[((yb + r0) * PATCH_STRIDE + (xa0 + q0)) * 4 + tig];
        A[0][2] = in_u32[((ya + r1) * PATCH_STRIDE + (xa0 + q1)) * 4 + tig];
        A[0][3] = in_u32[((yb + r1) * PATCH_STRIDE + (xa0 + q1)) * 4 + tig];
        A[1][0] = in_u32[((ya + r0) * PATCH_STRIDE + (xa1 + q0)) * 4 + tig];
        A[1][1] = in_u32[((yb + r0) * PATCH_STRIDE + (xa1 + q0)) * 4 + tig];
        A[1][2] = in_u32[((ya + r1) * PATCH_STRIDE + (xa1 + q1)) * 4 + tig];
        A[1][3] = in_u32[((yb + r1) * PATCH_STRIDE + (xa1 + q1)) * 4 + tig];

        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            const uint32_t* wp = reinterpret_cast<const uint32_t*>(
                &smem_w[(nt * 8 + g) * WSM_STRIDE + ch * 16]);
            uint32_t B[2];
            B[0] = wp[tig];
            B[1] = wp[tig + 4];
            mma16816(acc[0][nt], A[0], B);
            mma16816(acc[1][nt], A[1], B);
        }
    }

    float wsum = 0.f;
    #pragma unroll
    for (int w = 0; w < 2; ++w) {
        const bool valid = w ? valid1 : valid0;
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            float* c = acc[w][nt];
            const float b0 = smem_bias[nt * 8 + 2 * tig];
            const float b1 = smem_bias[nt * 8 + 2 * tig + 1];
            c[0] += b0; c[1] += b1; c[2] += b0; c[3] += b1;
            #pragma unroll
            for (int off = 4; off < 32; off <<= 1) {
                c[0] += __shfl_xor_sync(0xffffffffu, c[0], off);
                c[1] += __shfl_xor_sync(0xffffffffu, c[1], off);
                c[2] += __shfl_xor_sync(0xffffffffu, c[2], off);
                c[3] += __shfl_xor_sync(0xffffffffu, c[3], off);
            }
            float s = 1.0f / (1.0f + expf(-(c[0] + c[2]) * 0.0625f))
                    + 1.0f / (1.0f + expf(-(c[1] + c[3]) * 0.0625f));
            s = (g == 0) ? s : 0.0f;
            #pragma unroll
            for (int off = 1; off < 4; off <<= 1)
                s += __shfl_xor_sync(0xffffffffu, s, off);
            if (lane == 0 && valid) wsum += s;
        }
    }
    if (lane == 0) red[warp] = wsum;
    __syncthreads();
    if (t == 0) {
        float ssum = 0.f;
        #pragma unroll
        for (int i = 0; i < 8; ++i) ssum += red[i];
        atomicAdd(&out[n], ssum);
    }
}

__device__ __forceinline__ void block_reduce_atomic(float acc, float* out) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        acc += __shfl_down_sync(0xffffffffu, acc, off);
    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = acc;
    __syncthreads();
    if (wid == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        float v = (lane < nwarps) ? warp_sums[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) atomicAdd(out, v);
    }
}

__global__ void nchw_to_nhwc_fp16_c8_kernel(
        const float* __restrict__ in, __half* __restrict__ out,
        const size_t HW, const size_t total_pixels) {
    const size_t p = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= total_pixels) return;
    const size_t n = p / HW;
    const size_t s = p - n * HW;
    const float* base = in + n * 8 * HW + s;
    uint4 packed;
    __half2* ph = reinterpret_cast<__half2*>(&packed);
    #pragma unroll
    for (int c = 0; c < 4; ++c) {
        ph[c] = __floats2half2_rn(base[(size_t)(2 * c) * HW],
                                  base[(size_t)(2 * c + 1) * HW]);
    }
    *reinterpret_cast<uint4*>(out + p * 8) = packed;
}

__global__ void nchw_to_nhwc_fp16_generic_kernel(
        const float* __restrict__ in, __half* __restrict__ out,
        const int C, const size_t HW, const size_t total) {
    const size_t idx = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    const int c = (int)(idx % (size_t)C);
    const size_t p = idx / (size_t)C;
    const size_t n = p / HW;
    const size_t s = p - n * HW;
    out[idx] = __float2half_rn(in[(n * (size_t)C + c) * HW + s]);
}

template <int POOL>
__global__ void fused_pool_sigmoid_sum_nhwc_kernel(
        const __half* __restrict__ x, float* __restrict__ out,
        const int C, const int H, const int W,
        const int Hp, const int Wp, const int units_per_image) {
    const int n = blockIdx.y;
    const __half* __restrict__ img = x + (size_t)n * ((size_t)H * W * C);
    const int cg_count = C >> 3;

    float acc = 0.0f;
    const int stride = gridDim.x * blockDim.x;
    for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < units_per_image; v += stride) {
        const int cg = v % cg_count;
        const int tt = v / cg_count;
        const int wx = tt % Wp;
        const int wy = tt / Wp;
        const int c0 = cg << 3;

        const __half* base = img + ((size_t)(wy * POOL) * W + (size_t)(wx * POOL)) * C + c0;
        float sums[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};
        #pragma unroll
        for (int r = 0; r < POOL; ++r) {
            const __half* row = base + (size_t)r * W * C;
            #pragma unroll
            for (int q = 0; q < POOL; ++q) {
                const uint4 raw = *reinterpret_cast<const uint4*>(row + (size_t)q * C);
                const __half2* h2 = reinterpret_cast<const __half2*>(&raw);
                #pragma unroll
                for (int i = 0; i < 4; ++i) {
                    const float2 f = __half22float2(h2[i]);
                    sums[2 * i]     += f.x;
                    sums[2 * i + 1] += f.y;
                }
            }
        }
        constexpr float inv = 1.0f / (float)(POOL * POOL);
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            const float m = sums[i] * inv;
            acc += 1.0f / (1.0f + expf(-m));
        }
    }
    block_reduce_atomic(acc, &out[n]);
}

__device__ __forceinline__ float2 load_two(const float* p) {
    return *reinterpret_cast<const float2*>(p);
}
__device__ __forceinline__ float2 load_two(const __half* p) {
    return __half22float2(*reinterpret_cast<const __half2*>(p));
}

template <typename scalar_t, int POOL>
__global__ void fused_pool_sigmoid_sum_nchw_kernel(
        const scalar_t* __restrict__ x, float* __restrict__ out,
        const int C, const int H, const int W,
        const int Hp, const int Wp, const int num_windows) {
    const int n = blockIdx.y;
    const scalar_t* __restrict__ img = x + (size_t)n * C * H * W;

    float acc = 0.0f;
    const int stride = gridDim.x * blockDim.x;
    constexpr float inv = 1.0f / (float)(POOL * POOL);

    for (int v = blockIdx.x * blockDim.x + threadIdx.x; v < num_windows; v += stride) {
        const int wx = v % Wp;
        const int tt = v / Wp;
        const int wy = tt % Hp;
        const int c  = tt / Hp;

        const scalar_t* p = img + ((size_t)c * H + (size_t)wy * POOL) * W + (size_t)wx * POOL;
        float s = 0.0f;
        #pragma unroll
        for (int r = 0; r < POOL; ++r) {
            #pragma unroll
            for (int q = 0; q < POOL / 2; ++q) {
                const float2 val = load_two(p + 2 * q);
                s += val.x + val.y;
            }
            p += W;
        }
        const float m = s * inv;
        acc += 1.0f / (1.0f + expf(-m));
    }
    block_reduce_atomic(acc, &out[n]);
}

} // namespace

torch::Tensor fused_reduced_conv_pool_sigmoid_sum(
        torch::Tensor x, torch::Tensor wk, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4, "input must be CUDA NCHW");
    TORCH_CHECK(x.scalar_type() == torch::kFloat, "input must be fp32");
    auto xc = x.contiguous();
    const int N = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int H = (int)xc.size(2);
    const int W = (int)xc.size(3);
    TORCH_CHECK(C == 8, "in_channels must be 8");
    TORCH_CHECK(wk.is_cuda() && wk.scalar_type() == torch::kHalf, "wk must be CUDA fp16");
    TORCH_CHECK(wk.numel() == 64 * 80, "wk must be [64,80]");
    TORCH_CHECK(bias.is_cuda() && bias.scalar_type() == torch::kFloat, "bias must be CUDA fp32");
    TORCH_CHECK(bias.numel() == 64, "bias must be [64]");

    auto wkc = wk.contiguous();
    auto bc = bias.contiguous();
    const int Ho = H - 2;
    const int Wo = W - 2;
    TORCH_CHECK(Ho > 0 && Wo > 0, "spatial dims too small");
    const int Hp = Ho / 4;
    const int Wp = Wo / 4;

    auto out = torch::zeros({N}, xc.options().dtype(torch::kFloat32));
    if (N == 0 || Hp == 0 || Wp == 0) return out;

    const int tilesX = (Wp + 15) / 16;
    const int tilesY = (Hp + 7) / 8;
    dim3 grid(tilesX, tilesY, N);
    auto stream = at::cuda::getCurrentCUDAStream();
    fused_reduced_conv_pool_sigmoid_sum_kernel<<<grid, 256, 0, stream>>>(
        xc.data_ptr<float>(),
        reinterpret_cast<const __half*>(wkc.data_ptr<at::Half>()),
        bc.data_ptr<float>(),
        out.data_ptr<float>(),
        H, W, Hp, Wp, tilesX, tilesY);
    return out;
}

torch::Tensor fused_conv3x3_c8_pool4_sigmoid_sum(
        torch::Tensor x, torch::Tensor wk, torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 4, "input must be CUDA NCHW");
    TORCH_CHECK(x.scalar_type() == torch::kFloat, "input must be fp32");
    auto xc = x.contiguous();
    const int N = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int H = (int)xc.size(2);
    const int W = (int)xc.size(3);
    TORCH_CHECK(C == 8, "in_channels must be 8");

    auto wkc = wk.contiguous();
    auto bc = bias.contiguous();
    const int Ho = H - 2;
    const int Wo = W - 2;
    TORCH_CHECK(Ho > 0 && Wo > 0, "spatial dims too small");
    const int Hp = Ho / 4;
    const int Wp = Wo / 4;

    auto out = torch::zeros({N}, xc.options().dtype(torch::kFloat32));
    if (N == 0 || Hp == 0 || Wp == 0) return out