```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

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

#define TILE_X 8
#define TILE_Y 8
#define TILE_Z 4

__global__ void boxsum_kernel(const float* __restrict__ x,
                              float* __restrict__ box,
                              int N, int IC, int D, int H, int W) {
    __shared__ float acc[27];
    for (int i = threadIdx.x; i < 27; i += blockDim.x) acc[i] = 0.f;
    __syncthreads();

    int hw = H * W;
    int total = N * IC * D * hw;
    int dhw = D * hw;
    for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += gridDim.x * blockDim.x) {
        int t = idx % dhw;
        int c = (idx / dhw) % IC;
        int n = idx / (dhw * IC);
        int ix = t % W;
        int t2 = t / W;
        int iy = t2 % H;
        int iz = t2 / H;
        int rz = (iz == 0) ? 0 : ((iz == D - 1) ? 1 : 2);
        int ry = (iy == 0) ? 0 : ((iy == H - 1) ? 1 : 2);
        int rx = (ix == 0) ? 0 : ((ix == W - 1) ? 1 : 2);
        atomicAdd(&acc[rz * 9 + ry * 3 + rx], x[idx]);
    }
    __syncthreads();
    for (int i = threadIdx.x; i < 27; i += blockDim.x) {
        atomicAdd(&box[((long)blockIdx.x * IC + blockIdx.y) * 27 + i], acc[i]);
    }
}

__global__ void reduce_nic_kernel(const float* __restrict__ in, float* __restrict__ out, long inner, int noc) {
    __shared__ float sh[32];
    long chunk = blockIdx.x * 1024L;
    long c0 = chunk / inner;
    long rem = chunk % inner;
    long c1 = (chunk + 1023) / inner;
    if (c0 >= noc) return;
    int lane = threadIdx.x;
    float s = 0.f;
    if (c0 == c1) {
        for (long r = rem + lane; r < inner; r += 32) s += in[c0 * inner + r];
    } else {
        for (long r = rem + lane; r < inner; r += 32) s += in[c0 * inner + r];
        for (long c = c0 + 1; c < c1; ++c)
            for (long r = lane; r < inner; r += 32) s += in[c * inner + r];
        long e1 = (chunk + 1023) % inner;
        for (long r = lane; r <= e1; r += 32) s += in[c1 * inner + r];
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) s += __shfl_down_sync(0xffffffffu, s, off);
    if (lane == 0) sh[threadIdx.y] = s;
    __syncthreads();
    if (threadIdx.y == 0) {
        float v = (lane < 32) ? sh[lane] : 0.f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) atomicAdd(&out[c0], v);
    }
}

__global__ void wregion_kernel(const float* __restrict__ w, float* __restrict__ wr,
                               int IC, int OC, bool trans) {
    int ic = blockIdx.x;
    int oc = blockIdx.y;
    float lsum[3], msum[3];
    float a00 = w[((long)ic * OC + oc) * 27 + 0];
    float a01 = w[((long)ic * OC + oc) * 27 + 1];
    float a02 = w[((long)ic * OC + oc) * 27 + 2];
    float a10 = w[((long)ic * OC + oc) * 27 + 3];
    float a11 = w[((long)ic * OC + oc) * 27 + 4];
    float a12 = w[((long)ic * OC + oc) * 27 + 5];
    float a20 = w[((long)ic * OC + oc) * 27 + 6];
    float a21 = w[((long)ic * OC + oc) * 27 + 7];
    float a22 = w[((long)ic * OC + oc) * 27 + 8];
    lsum[0] = a01 + a02; lsum[1] = a00 + a01; lsum[2] = a00 + a01 + a02;
    msum[0] = a11 + a12; msum[1] = a10 + a11; msum[2] = a10 + a11 + a12;
    float h0 = a21 + a22, h1 = a20 + a21, h2 = a20 + a21 + a22;
    float hz[3]; hz[0] = h0; hz[1] = h1; hz[2] = h2;
    const float* rows[3] = {lsum, msum, hz};
    float hsum[3] = {lsum[0] + msum[0] + hz[0], lsum[1] + msum[1] + hz[1], lsum[2] + msum[2] + hz[2]};
    float l2[3], m2[3], h2_[3];
    for (int r = 0; r < 3; ++r) {
        l2[r] = a10 + a20 + (r == 2 ? a00 : 0.f);
        m2[r] = a11 + a21 + (r == 2 ? a01 : 0.f);
        h2_[r] = a12 + a22 + (r == 2 ? a02 : 0.f);
    }
    float vsum[3] = {l2[0] + m2[0] + h2_[0], l2[1] + m2[1] + h2_[1], l2[2] + m2[2] + h2_[2]};
    const float* cols[3] = {l2, m2, h2_};
    #pragma unroll
    for (int rz = 0; rz < 3; ++rz)
        #pragma unroll
        for (int ry = 0; ry < 3; ++ry)
            #pragma unroll
            for (int rx = 0; rx < 3; ++rx)
                wr[((long)(ic * OC + oc) * 27 + rz * 9 + ry * 3 + rx)] = hsum[rz] - rows[ry][rz] - cols[rx][rz] + vsum[rx];
    if (trans) {
        float* d = wr + (long)IC * OC * 27 + ((long)oc * IC + ic) * 27;
        d[0] = a00;  d[1] = a01;  d[2] = a02;
        d[3] = a10;  d[4] = a11;  d[5] = a12;
        d[6] = a20;  d[7] = a21;  d[8] = a22;
        d[9] = a00 + a10;  d[10] = a01 + a11;  d[11] = a02 + a12;
        d[12] = a20;       d[13] = a21;       d[14] = a22;
        d[15] = a20 + a10; d[16] = a21 + a11; d[17] = a22 + a12;
        d[18] = a00 + a01; d[19] = a02;       d[20] = a01 + a02;
        d[21] = a10 + a11; d[22] = a12;       d[23] = a11 + a12;
        d[24] = a20 + a21; d[25] = a22;       d[26] = a21 + a22;
    }
}

__global__ void __launch_bounds__(256) mean_kernel(const float* __restrict__ box,
                                                   const float* __restrict__ wreg,
                                                   const float* __restrict__ gamma,
                                                   const float* __restrict__ rv,
                                                   float* __restrict__ buf,
                                                   int N, int IC, int OC,
                                                   float invS, float eps) {
    int oc = blockIdx.x * blockDim.y + threadIdx.y;
    if (oc >= OC) return;
    int n = blockIdx.y;
    float s = 0.f;
    const float* bp = box + (long)n * IC * 27;
    const float* wp = wreg + (long)oc * 27;
    int stride = OC * 27;
    for (int ic = 0; ic < IC; ++ic) {
        float w0 = wp[0], w1 = wp[1], w2 = wp[2], w3 = wp[3], w4 = wp[4], w5 = wp[5], w6 = wp[6], w7 = wp[7], w8 = wp[8];
        float w9 = wp[9], w10 = wp[10], w11 = wp[11], w12 = wp[12], w13 = wp[13], w14 = wp[14], w15 = wp[15], w16 = wp[16], w17 = wp[17];
        float w18 = wp[18], w19 = wp[19], w20 = wp[20], w21 = wp[21], w22 = wp[22], w23 = wp[23], w24 = wp[24], w25 = wp[25], w26 = wp[26];
        const float* b = bp + (long)ic * 27;
        s = fmaf(b[0], w0, s); s = fmaf(b[1], w1, s); s = fmaf(b[2], w2, s);
        s = fmaf(b[3], w3, s); s = fmaf(b[4], w4, s); s = fmaf(b[5], w5, s);
        s = fmaf(b[6], w6, s); s = fmaf(b[7], w7, s); s = fmaf(b[8], w8, s);
        s = fmaf(b[9], w9, s); s = fmaf(b[10], w10, s); s = fmaf(b[11], w11, s);
        s = fmaf(b[12], w12, s); s = fmaf(b[13], w13, s); s = fmaf(b[14], w14, s);
        s = fmaf(b[15], w15, s); s = fmaf(b[16], w16, s); s = fmaf(b[17], w17, s);
        s = fmaf(b[18], w18, s); s = fmaf(b[19], w19, s); s = fmaf(b[20], w20, s);
        s = fmaf(b[21], w21, s); s = fmaf(b[22], w22, s); s = fmaf(b[23], w23, s);
        s = fmaf(b[24], w24, s); s = fmaf(b[25], w25, s); s = fmaf(b[26], w26, s);
        wp += stride;
    }
    float a = gamma[oc] * rsqrtf(rv[oc] + eps);
    buf[n * OC + oc] = a;
    buf[N * OC + n * OC + oc] = -a * s * invS;
}

__global__ void __launch_bounds__(256) final_kernel(float* __restrict__ data,
                                                    const float* __restrict__ buf,
                                                    int N, int OC, long S) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= S) return;
    int oc = blockIdx.y;
    int n = blockIdx.z;
    float a = buf[n * OC + oc];
    float b = buf[N * OC + n * OC + oc];
    long base = ((long)n * OC + oc) * S;
    data[base + i] = fmaf(a, data[base + i], b);
}

__global__ void __launch_bounds__(256, 4) fused_conv_kernel(
        const float* __restrict__ x,
        const float* __restrict__ wre,
        const float* __restrict__ buf,
        float* __restrict__ out,
        int N, int IC, int OC,
        int iD, int iH, int iW,
        int oD, int oH, int oW,
        int ozb) {
    int zb = blockIdx.z % ozb;
    int t = blockIdx.z / ozb;
    int oc = t % OC;
    int n = t / OC;

    int zp = zb * TILE_Z + threadIdx.z;
    int yp = blockIdx.y * (TILE_Y * 2) + threadIdx.y * 2;
    int xp = blockIdx.x * (TILE_X * 2) + threadIdx.x * 2;

    float acc[2][4][4];
    #pragma unroll
    for (int a = 0; a < 2; ++a)
        #pragma unroll
        for (int b = 0; b < 4; ++b)
            #pragma unroll
            for (int c = 0; c < 4; ++c) acc[a][b][c] = 0.f;

    bool z1ok = (zp + 1) < iD;
    bool y1ok = (yp + 2) < iH;
    bool x1ok = (xp + 2) < iW;
    bool interior = z1ok & y1ok & x1ok;

    long in_plane = (long)iH * iW;
    const float* xbase = x + ((long)n * IC) * (iD * in_plane);
    const float* wp = wre + (long)oc * IC * 32;

    for (int ic = 0; ic < IC; ++ic) {
        float xv[2][3][3];
        if (interior) {
            const float* p = xbase + (long)ic * (iD * in_plane) + (long)zp * in_plane + (long)yp * iW + xp;
            #pragma unroll
            for (int dz = 0; dz < 2; ++dz) {
                const float* pz = p + (long)dz * in_plane;
                #pragma unroll
                for (int dy = 0; dy < 3; ++dy) {
                    const float* py = pz + (long)dy * iW;
                    xv[dz][dy][0] = __ldg(py);
                    xv[dz][dy][1] = __ldg(py + 1);
                    xv[dz][dy][2] = __ldg(py + 2);
                }
            }
        } else {
            #pragma unroll
            for (int dz = 0; dz < 2; ++dz) {
                #pragma unroll
                for (int dy = 0; dy < 3; ++dy) {
                    #pragma unroll
                    for (int dx = 0; dx < 3; ++dx) {
                        int iz = zp + dz, iy = yp + dy, ix = xp + dx;
                        float v = 0.f;
                        if (iz < iD && iy < iH && ix < iW)
                            v = __ldg(xbase + ((long)ic * iD + iz) * in_plane + (long)iy * iW + ix);
                        xv[dz][dy][dx] = v;
                    }
                }
            }
        }

        float4 q0 = __ldg((const float4*)(wp) + 0);
        float4 q1 = __ldg((const float4*)(wp) + 1);
        float4 q2 = __ldg((const float4*)(wp) + 2);
        float4 q3 = __ldg((const float4*)(wp) + 3);
        float4 q4 = __ldg((const float4*)(wp) + 4);
        float4 q5 = __ldg((const float4*)(wp) + 5);
        float4 q6 = __ldg((const float4*)(wp) + 6);
        float A00 = q0.x, A01 = q0.y, A02 = q0.z, A10 = q0.w;
        float A11 = q1.x, A12 = q1.y, A20 = q1.z, A21 = q1.w;
        float A22 = q2.x;
        float B00 = q2.y, B01 = q2.z, B02 = q2.w;
        float B20 = q3.x, B21 = q3.y, B22 = q3.z;
        float C00 = q3.w, C01 = q4.x, C02 = q4.y;
        float C10 = q4.z, C11 = q4.w, C12 = q5.x;
        float C20 = q5.y, C21 = q5.z, C22 = q5.w;
        float D00 = q6.x, D01 = q6.y, D02 = q6.z;
        float D10 = q6.w;
        wp += 32;

        {
            float u0 = fmaf(xv[1][0][0], B00, fmaf(xv[0][0][0], C00, 0.f));
            float u1 = fmaf(xv[1][0][1], B00, fmaf(xv[0][0][1], C00, 0.f));
            float u2 = fmaf(xv[1][0][2], B00, fmaf(xv[0][0][2], C00, 0.f));
            acc[1][0][0] = fmaf(xv[0][0][0], D00, fmaf(xv[1][0][0], D10, fmaf(u0, A00, acc[1][0][0])));
            acc[1][0][1] = fmaf(xv[0][0][1], D00, fmaf(xv[1][0][1], D10, fmaf(u0, A01, fmaf(u1, A02, acc[1][0][1]))));
            acc[1][0][2] = fmaf(xv[0][0][2], D00, fmaf(xv[1][0][2], D10, fmaf(u1, A01, fmaf(u2, A02, acc[1][0][2]))));
            acc[1][0][3] = fmaf(u2, A01, acc[1][0][3]);
        }
        {
            float u0 = fmaf(xv[1][0][0], B01, fmaf(xv[0][0][0], C01, 0.f));
            float u1 = fmaf(xv[1][0][1], B01, fmaf(xv[0][0][1], C01, 0.f));
            float u2 = fmaf(xv[1][0][2], B01, fmaf(xv[0][0][2], C01, 0.f));
            float w0 = fmaf(xv[1][1][0], B00, fmaf(xv[0][1][0], C00, 0.f));
            float w1 = fmaf(xv[1][1][1], B00, fmaf(xv[0][1][1], C00, 0.f));
            float w2 = fmaf(xv[1][1][2], B00, fmaf(xv[0][1][2], C00, 0.f));
            acc[0][0][0] = fmaf(u0, A10, acc[0][0][0]);
            acc[0][0][1] = fmaf(u0, A11, fmaf(u1, A12, acc[0][0][1]));
            acc[0][0][2] = fmaf(u1, A11, fmaf(u2, A12, acc[0][0][2]));
            acc[0][0][3] = fmaf(u2, A11, acc[0][0][3]);
            acc[1][1][0] = fmaf(xv[0][1][0], D01, fmaf(u0, A20, fmaf(xv[1][1][0], D11, fmaf(w0, A00, acc[1][1][0]))));
            acc[1][1][1] = fmaf(xv[0][1][1], D01, fmaf(u0, A21, fmaf(u1, A22, fmaf(xv[1][1][1], D11, fmaf(w0, A01, fmaf(w1, A02, acc[1][1][1]))))));
            acc[1][1][2] = fmaf(xv[0][1][2], D01, fmaf(u1, A21, fmaf(u2, A22, fmaf(xv[1][1][2], D11, fmaf(w1, A01, fmaf(w2, A02, acc[1][1][2]))))));
            acc[1][1][3] = fmaf(u2, A21, fmaf(w2, A01, acc[1][1][3]));
        }
        {
            float u0 = fmaf(xv[1][0][0], B02, fmaf(xv[0][0][0], C02, 0.f));
            float u1 = fmaf(xv[1][0][1], B02, fmaf(xv[0][0][1], C02, 0.f));
            float u2 = fmaf(xv[1][0][2], B02, fmaf(xv[0][0][2], C02, 0.f));
            float w0 = fmaf(xv[1][1][0], B01, fmaf(xv[0][1][0], C01, 0.f));
            float w1 = fmaf(xv[1][1][1], B01, fmaf(xv[0][1][1], C01, 0.f));
            float w2 = fmaf(xv[1][1][2], B01, fmaf(xv[0][1][2], C01, 0.f));
            float z0 = fmaf(xv[1][2][0], B00, fmaf(xv[0][2][0], C00, 0.f));
            float z1 = fmaf(xv[1][2][1], B00, fmaf(xv[0][2][1], C00, 0.f));
            float z2 = fmaf(xv[1][2][2], B00, fmaf(xv[0][2][2], C00, 0.f));
            acc[0][1][0] = fmaf(u0, A10, fmaf(w0, A20, acc[0][1][0]));
            acc[0][1][1] = fmaf(u0, A11, fmaf(u1, A12, fmaf(w0, A21, fmaf(w1, A22, acc[0][1][1]))));
            acc[0][1][2] = fmaf(u1, A11, fmaf(u2, A12, fmaf(w1, A21, fmaf(w2, A22, acc[0][1][2]))));
            acc[0][1][3] = fmaf(u2, A11, fmaf(w2, A21, acc[0][1][3]));
            acc[1][2][0] = fmaf(xv[0][2][0], D00, fmaf(xv[1][2][0], D10, fmaf(u0, A20, fmaf(w0, A00, acc[1][2][0]))));
            acc[1][2][1] = fmaf(xv[0][2][1], D00, fmaf(xv[1][2][1], D10, fmaf(u0, A21, fmaf(u1, A22, fmaf(w0, A01, fmaf(w1, A02, acc[1][2][1]))))));
            acc[1][2][2] = fmaf(xv[0][2][2], D00, fmaf(xv[1][2][2], D10, fmaf(u1, A21, fmaf(u2, A22, fmaf(w1, A01, fmaf(w2, A02, acc[1][2][2]))))));
            acc[1][2][3] = fmaf(u2, A21, fmaf(w2, A01, acc[1][2][3]));
            acc[0][2][0] = fmaf(w0, A10, fmaf(z0, A20, acc[0][2][0]));
            acc[0][2][1] = fmaf(w0, A11, fmaf(w1, A12, fmaf(z0, A21, fmaf(z1, A22, acc[0][2][1]))));
            acc[0][2][2] = fmaf(w1, A11, fmaf(w2, A12, fmaf(z1, A21, fmaf(z2, A22, acc[0][2][2]))));
            acc[0][2][3] = fmaf(w2, A11, fmaf(z2, A21, acc[0][2][3]));
        }
        {
            float w0 = fmaf(xv[1][1][0], B02, fmaf(xv[0][1][0], C02, 0.f));
            float w1 = fmaf(xv[1][1][1], B02, fmaf(xv[0][1][1], C02, 0.f));
            float w2 = fmaf(xv[1][1][2], B02, fmaf(xv[0][1][2], C02, 0.f));
            float z0 = fmaf(xv[1][2][0], B01, fmaf(xv[0][2][0], C01, 0.f));
            float z1 = fmaf(xv[1][2][1], B01, fmaf(xv[0][2][1], C01, 0.f));
            float z2 = fmaf(xv[1][2][2], B01, fmaf(xv[0][2][2], C01, 0.f));
            acc[0][1][0] = fmaf(w0, A10, acc[0][1][0]);
            acc[0][1][1] = fmaf(w0, A11, fmaf(w1, A12, acc[0][1][1]));
            acc[0][1][2] = fmaf(w1, A11, fmaf(w2, A12, acc[0][1][2]));
            acc[0][1][3] = fmaf(w2, A11, acc[0][1][3]);
            acc[1][2][0] = fmaf(xv[0][1][0], D02, fmaf(w0, A20, acc[1][2][0]));
            acc[1][2][1] = fmaf(xv[0][1][1], D02, fmaf(w0, A21, fmaf(w1, A22, acc[1][2][1])));
            acc[1][2][2] = fmaf(xv[0][1][2], D02, fmaf(w1, A21, fmaf(w2, A22, acc[1][2][2])));
            acc[1][2][3] = fmaf(w2, A21, acc[1][2][3]);
            acc[0][2][0] = fmaf(w0, A10, fmaf(z0, A20, acc[0][2][0]));
            acc[0][2][1] = fmaf(w0, A11, fmaf(w1, A12, fmaf(z0, A21, fmaf(z1, A22, acc[0][2][1]))));
            acc[0][2][2] = fmaf(w1, A11, fmaf(w2, A12, fmaf(z1, A21, fmaf(z2, A22, acc[0][2][2]))));
            acc[0][2][3] = fmaf(w2, A11, fmaf(z2, A21, acc[0][2][3]));
            acc[0][3][0] = fmaf(z0, A10, acc[0][3][0]);
            acc[0][3][1] = fmaf(z0, A11, fmaf(z1, A12, acc[0][3][1]));
            acc[0][3][2] = fmaf(z1, A11, fmaf(z2, A12, acc[0][3][2]));
            acc[0][3][3] = fmaf(z2, A11, acc[0][3][3]);
        }
    }

    float sc = buf[n * OC + oc];
    float sh = buf[N * OC + n * OC + oc];

    long out_plane = (long)oH * oW;
    float* obase = out + ((long)n * OC + oc) * (oD * out_plane);
    int oz0 = 2 * zp;
    #pragma unroll
    for (int a = 0; a < 2; ++a) {
        int oz = oz0 + a;
        if (oz >= oD) continue;
        #pragma unroll
        for (int b = 0; b < 4; ++b) {
            int oy = 2 * yp + b;
            if (oy >= oH) continue;
            int ox = 2 * xp;
            float* op = obase + (long)oz * out_plane + (long)oy * oW + ox;
            if (ox + 3 < oW) {
                float2 v01, v23;
                v01.x = fmaf(sc, acc[a][b][0], sh);
                v01.y = fmaf(sc, acc[a][b][1], sh);
                v23.x = fmaf(sc, acc[a][b][2], sh);
                v23.y = fmaf(sc, acc[a][b][3], sh);
                *(float2*)(op) = v01;
                *(float2*)(op + 2) = v23;
            } else {
                #pragma unroll
                for (int c = 0; c < 4; ++c) {
                    if (ox + c < oW) op[c] = fmaf(sc, acc[a][b][c], sh);
                }
            }
        }
    }
}

std::vector<torch::Tensor> fused_forward(
        torch::Tensor x, torch::Tensor weight, torch::Tensor gamma,
        torch::Tensor rv, torch::Tensor buf, torch::Tensor box,
        torch::Tensor wreg, double eps,
        int64_t N, int64_t IC, int64_t OC,
        int64_t iD, int64_t iH, int64_t iW,
        int64_t oD, int64_t oH, int64_t oW,
        int64_t ngroups) {
    auto stream = at::cuda::getCurrentCUDAStream();
    long S = (long)oD * oH * oW;

    boxsum_kernel<<<(int)ngroups, (int)IC, 0, stream>>>(
        x.data_ptr<float>(), box.data_ptr<float>(), (int)N, (int)IC, (int)iD, (int)iH, (int)iW);

    wregion_kernel<<<dim3((int)IC, (int)OC), 1, 0, stream>>>(
        weight.data_ptr<float>(), wreg.data_ptr<float>(), (int)IC, (int)OC, true);

    {
        long inner = (long)iD * iH * iW;
        long tot = (long)N * IC * inner;
        int nb = (int)((tot + 1023) / 1024);
        reduce_nic_kernel<<<nb, dim3(32, 32), 0, stream>>>(
            box.data_ptr<float>(), box.data_ptr<float>() + (long)N * IC * 27, inner, (int)(N * IC));
    }

    float invS = 1.0f / (float)S;
    mean_kernel<<<dim3((int)((OC + 15) / 16), (int)N), dim3(1, 16), 0, stream>>>(
        box.data_ptr<float>(), wreg.data_ptr<float>(), gamma.data_ptr<float>(), rv.data_ptr<float>(),
        buf.data_ptr<float>(), (int)N, (int)IC, (int)OC, invS, (float)eps);

    auto out = torch::empty({N, OC, oD, oH, oW}, x.options());
    int ozb = (int)((oD / 2 + TILE_Z - 1) / TILE_Z);
    dim3 grid((int)((oW / 2 + TILE_X - 1) / TILE_X),
              (int)((oH / 2 + TILE_Y - 1) / TILE_Y),
              (int)(N * OC * ozb));
    fused_conv_kernel<<<grid, dim3(TILE_X, TILE_Y, TILE_Z), 0, stream>>>(
        x.data_ptr<float>(), wreg.data_ptr<float>() + (long)IC * OC * 27, buf.data_ptr<float>(),
        out.data_ptr<float>(), (int)N, (int)IC, (int)OC,
        (int)iD, (int)iH, (int)iW, (int)oD, (int)oH, (int)oW, ozb);
    return {out};
}

std::vector<torch::Tensor> finalize(torch::Tensor data, torch::Tensor buf,
                                    int64_t N, int64_t OC, int64_t S) {
    auto stream = at::cuda::getCurrentCUDAStream();
    int threads = 256;
    long blocks = (S + threads - 1) / threads;
    final_kernel<<<dim3((int)blocks, (int)OC, (int)N), threads, 0, stream>>>(
        data.data_ptr<float>(), buf.data_ptr<float>(), (int)N, (int)OC, (long)S);
    return {data};
