```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

# Custom CUDA kernels for Conv3d with kernel (K, K, 1).
#
# The operator is an independent 2D convolution over dims (2, 3) for every
# index of the innermost dim 4 ("depth"), i.e. a GEMM
#   out[(n,h,w,d), oc] = sum_{c,ky,kx} x[n,c,h+ky,w+kx,d] * W[oc,c,ky,kx]
# with M = N*HO*WO*D rows, N = OC columns, K = C*KS*KS.
#
# Fast path (sm_80+, K <= 32, OC % 8 == 0, M % 128 == 0, output contiguous
# rows): implicit-GEMM tensor-core kernel using mma.m16n8k8.tf32 with the
# 3xTF32 decomposition (hi*hi + hi*lo + lo*hi), which delivers near-fp32
# accuracy at tensor-core throughput:
#   * 128x64 block tile, 4 warps (2x2), warp tile 64x32,
#   * A staged through shared memory with im2col gather, coalesced global
#     reads and an XOR swizzle that makes both staging and mma fragment
#     loads bank-conflict-free,
#   * B (weights, K <= 32 x OC <= 256) is split hi/lo once per kernel into a
#     persistent shared buffer (weights fit in L2/L1 easily),
#   * fp32 accumulators stay in registers for all 8 n8 tiles (single k-step
#     when K <= 8) or are drained through a padded shared staging buffer so
#     global stores are fully coalesced 128B transactions.
#
# A direct scalar fallback kernel handles every other configuration, and any
# unsupported case falls back to the reference nn.Conv3d (cuDNN).
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 128
#define PADQ 8
#define MAXOC 256
#define MAXK  32
#define MAXK0 32

__device__ __forceinline__ unsigned cvt_tf32(float x) {
    unsigned r;
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
}

__device__ __forceinline__ void mma_tf32(
    float &c0, float &c1, float &c2, float &c3,
    unsigned a0, unsigned a1, unsigned a2, unsigned a3,
    unsigned b0, unsigned b1)
{
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// Split W[oc, k] into tf32 hi/lo parts once per kernel launch.
// threads.x must be >= OC; every block replicates the result.
__global__ void splitw_kernel(const float* __restrict__ w,
                              unsigned* __restrict__ wh,
                              unsigned* __restrict__ wl,
                              int OC, int K)
{
    __shared__ float wf[MAXK0];
    int oc = threadIdx.x;
    for (int k0 = 0; k0 < K; k0 += MAXK0) {
        int kmax = min(MAXK0, K - k0);
        __syncthreads();
        if (oc == 0) {
            for (int k = 0; k < kmax; ++k) wf[k] = w[(size_t)k * OC + k0 + k];
        }
        __syncthreads();
        if (oc < OC) {
            const float* wp = w + (size_t)oc * K + k0;
            unsigned* hp = wh + (size_t)oc * K + k0;
            unsigned* lp = wl + (size_t)oc * K + k0;
            for (int k = 0; k < kmax; ++k) {
                float x = wp[k];
                unsigned h = cvt_tf32(x);
                float l = x - __uint_as_float(h);
                hp[k] = h;
                lp[k] = cvt_tf32(l);
            }
        }
    }
}

template <bool SINGLE_K>
__global__ void __launch_bounds__(128, 3) conv_tc_kernel(
    const float* __restrict__ x,
    const unsigned* __restrict__ wh,
    const unsigned* __restrict__ wl,
    float* __restrict__ out,
    int M, int OC, int K, int Q,
    int HO, int WO, int D, int C, int H, int W,
    int KS, int pad)
{
    extern __shared__ char smem_raw[];
    // layout: Ahi | Alo | Bhi/Blo (SINGLE_K ? K : 32)*OC each | q-values
    unsigned* Ahi = (unsigned*)smem_raw;
    unsigned* Alo = Ahi + BM * MAXK;
    unsigned* Bh  = Alo + BM * MAXK;
    unsigned* Bl  = Bh + (size_t)(SINGLE_K ? MAXK0 : MAXK) * MAXOC;
    float*    qv  = (float*)(Bl + (size_t)(SINGLE_K ? MAXK0 : MAXK) * MAXOC);
    float*    Cs  = qv + (SINGLE_K ? 0 : (BM * PADQ));

    const int tx  = threadIdx.x;
    const int wr  = tx >> 5;          // warp row (2)
    const int wc  = wr & 1;           // warp col (2)
    const int g   = (tx >> 2) & 7;    // mma group id
    const int t   = tx & 3;           // mma thread id in group
    const int swz = t << 3;

    const int m0  = blockIdx.x * BM;

    // ---------------- stage B (weights) ----------------
    {
        const int cols = SINGLE_K ? OC : MAXOC;
        const int nB   = SINGLE_K ? K * OC : (MAXK * OC);
        for (int i = tx; i < nB; i += 128) {
            int k = i / OC, oc = i - k * OC;
            unsigned hv = 0, lv = 0;
            if (k < K && oc < OC) {
                size_t a = (size_t)k * OC + oc;
                hv = wh[a];
                lv = wl[a];
            }
            int nph = (oc ^ ((k & 3) << 3)) & (cols - 1);
            Bh[k * cols + nph] = hv;
            Bl[k * cols + nph] = lv;
        }
    }

    // ---------------- stage A (im2col gather) ----------------
    const long x_stride_c = (long)H * W * D;
    {
        #pragma unroll
        for (int i = 0; i < MAXK; ++i) {   // MAXK/BM = 32/128 -> 1 row each
            int r = tx + i * 128;          // r < BM always (i == 0)
            if (r < BM) {
                int m  = m0 + r;
                int rw = m / D;   int d  = m - rw * D;
                int rh = rw / WO; int w  = rw - rh * WO;
                int rn = rh / HO; int h  = rh - rn * HO;
                const float* xn = x + (size_t)rn * C * x_stride_c;
                for (int k = 0; k < K; ++k) {
                    int c  = k / (KS * KS);
                    int kk = k - c * (KS * KS);
                    int ky = kk / KS, kx = kk - ky * KS;
                    int ih = h - pad + ky;
                    int iw = w - pad + kx;
                    float v = 0.0f;
                    if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W) {
                        v = xn[((size_t)c * H + ih) * W * D + (size_t)iw * D + d];
                    }
                    unsigned h32 = cvt_tf32(v);
                    unsigned l32 = cvt_tf32(v - __uint_as_float(h32));
                    int rp = r ^ ((k & 3) << 3);
                    Ahi[k * BM + rp] = h32;
                    Alo[k * BM + rp] = l32;
                }
            }
        }
    }
    __syncthreads();

    // ---------------- shared pointers for this warp ----------------
    const int rows_w = wr << 6;                    // 64-row slab
    const int ksw    = wc ? 4 : 0;
    const unsigned* AwH = Ahi + ksw * BM + rows_w + g + ((t & 2) << 2);
    const unsigned* AwL = Alo + ksw * BM + rows_w + g + ((t & 2) << 2);
    const int cols = SINGLE_K ? OC : MAXOC;
    const unsigned* BwH = Bh + ksw * cols + (wc << 5) + ((g ^ t) << 3) + t;
    const unsigned* BwL = Bl + ksw * cols + (wc << 5) + ((g ^ t) << 3) + t;

    const int nreg  = SINGLE_K ? (OC >> 3) : 4;    // live n8 tiles
    const int nof   = SINGLE_K ? 0 : (wc << 5);
    const int mof   = SINGLE_K ? (wr << 6) : 0;
    const int npass = SINGLE_K ? 1 : ((OC + 31) >> 5);

    // ---------------- precompute coalesced store indices ----------------
    // (computed unconditionally so every thread reaches all barriers)
    int  rlim = 0, wmax = 0, hleft = 0, nout = 0, hout = 0;
    long yrow = 0;
    int  wq = 0, dq = 0;
    bool kvalid[PADQ];
    int  ksrc[PADQ], kdst[PADQ];
    if (!SINGLE_K) {
        long m_base = (long)blockIdx.x * BM + (wr << 6);
        rlim  = (int)min((long)64, (long)M - m_base);
        int n0  = (int)(m_base / Q);
        int rh0 = (int)(m_base - (long)n0 * Q);
        nout = n0; hout = rh0 / WO; wq = rh0 - hout * WO;
        wmax  = min(WO - wq, 7);
        hleft = HO - hout;
        yrow  = ((long)n0 * OC * HO + hout) * Q + (long)wq * D;
        dq = 0;
        #pragma unroll
        for (int k = 0; k < PADQ; ++k) {
            int r = tx + k * 128;
            kvalid[k] = (r < rlim);
            ksrc[k]   = kvalid[k] ? (r << 3) : 0;
            kdst[k]   = kvalid[k] ? r : 0;
        }
    }

    for (int np = 0; np < npass; ++np) {
        const unsigned* BHp = BwH + (SINGLE_K ? 0 : (np << 8));
        const unsigned* BLp = BwL + (SINGLE_K ? 0 : (np << 8));

        float c[16][4];
        #pragma unroll
        for (int i = 0; i < 16; ++i) {
            c[i][0] = 0.0f; c[i][1] = 0.0f; c[i][2] = 0.0f; c[i][3] = 0.0f;
        }

        if (SINGLE_K) {
            unsigned aH[4][4], aL[4][4];
            #pragma unroll
            for (int mt = 0; mt < 4; ++mt) {
                const unsigned* apH = AwH + (mt << 4);
                const unsigned* apL = AwL + (mt << 4);
                aH[mt][0] = apH[0];             aH[mt][1] = apH[8];
                aH[mt][2] = apH[4 * BM];        aH[mt][3] = apH[4 * BM + 8];
                aL[mt][0] = apL[0];             aL[mt][1] = apL[8];
                aL[mt][2] = apL[4 * BM];        aL[mt][3] = apL[4 * BM + 8];
            }
            for (int nt = 0; nt < nreg; ++nt) {
                unsigned bH0 = BwH[nt << 3], bH1 = BwH[(nt << 3) + 4 * cols];
                unsigned bL0 = BwL[nt << 3], bL1 = BwL[(nt << 3) + 4 * cols];
                #pragma unroll
                for (int mt = 0; mt < 4; ++mt) {
                    float* cc = c[(mt << 2) + (nt & 3)];
                    mma_tf32(cc[0], cc[1], cc[2], cc[3],
                             aL[mt][0], aL[mt][1], aL[mt][2], aL[mt][3],
                             bH0, bH1);
                    mma_tf32(cc[0], cc[1], cc[2], cc[3],
                             aH[mt][0], aH[mt][1], aH[mt][2], aH[mt][3],
                             bL0, bL1);
                    mma_tf32(cc[0], cc[1], cc[2], cc[3],
                             aH[mt][0], aH[mt][1], aH[mt][2], aH[mt][3],
                             bH0, bH1);
                }
            }
        } else {
            for (int k0 = 0; k0 < K; k0 += 8) {
                unsigned aH[4][4], aL[4][4];
                #pragma unroll
                for (int mt = 0; mt < 4; ++mt) {
                    const unsigned* apH = AwH + k0 * BM + (mt << 4);
                    const unsigned* apL = AwL + k0 * BM + (mt << 4);
                    aH[mt][0] = apH[0];             aH[mt][1] = apH[8];
                    aH[mt][2] = apH[4 * BM];        aH[mt][3] = apH[4 * BM + 8];
                    aL[mt][0] = apL[0];             aL[mt][1] = apL[8];
                    aL[mt][2] = apL[4 * BM];        aL[mt][3] = apL[4 * BM + 8];
                }
                #pragma unroll
                for (int nt = 0; nt < 4; ++nt) {
                    unsigned bH0 = BHp[k0 * cols + (nt << 3)];
                    unsigned bH1 = BHp[k0 * cols + (nt << 3) + 4 * cols];
                    unsigned bL0 = BLp[k0 * cols + (nt << 3)];
                    unsigned bL1 = BLp[k0 * cols + (nt << 3) + 4 * cols];
                    #pragma unroll
                    for (int mt = 0; mt < 4; ++mt) {
                        float* cc = c[(mt << 2) + nt];
                        mma_tf32(cc[0], cc[1], cc[2], cc[3],
                                 aL[mt][0], aL[mt][1], aL[mt][2], aL[mt][3],
                                 bH0, bH1);
                        mma_tf32(cc[0], cc[1], cc[2], cc[3],
                                 aH[mt][0], aH[mt][1], aH[mt][2], aH[mt][3],
                                 bL0, bL1);
                        mma_tf32(cc[0], cc[1], cc[2], cc[3],
                                 aH[mt][0], aH[mt][1], aH[mt][2], aH[mt][3],
                                 bH0, bH1);
                    }
                }
            }
        }

        if (SINGLE_K) {
            // ---------------- staged coalesced stores ----------------
            const int wq0 = wq, dq0 = dq, nout0 = nout, hleft0 = hleft;
            long ybase = yrow + (long)mof * D + (long)nof * Q;
            for (int nt = 0; nt < nreg; ++nt) {
                float* q0 = qv + (g << 3) + (t << 1);
                float* q1 = qv + ((g + 8) << 3) + (t << 1);
                long yw = ybase + (long)(nt << 3) * Q;
                #pragma unroll
                for (int mt = 0; mt < 4; ++mt) {
                    float* cc = c[(mt << 2) + (nt & 3)];
                    q0[mt << 6]     = cc[0];
                    q0[(mt << 6)+1] = cc[1];
                    q1[mt << 6]     = cc[2];
                    q1[(mt << 6)+1] = cc[3];
                }
                __syncthreads();
                int wqq = wq0, dqq = dq0, nn = nout0, hl = hleft0;
                long yl = yw;
                int rem = rlim;
                while (rem > 0) {
                    int wseg = min(WO - wqq, 7);
                    int rwc  = wseg * D;
                    int cnt  = min(rwc, rem);
                    const float* src = qv + (((int)(yl - yw) / (int)D) << 3);
                    float* dst = out + yl;
                    for (int r = tx; r < cnt; r += 128) {
                        int dd  = r % D;
                        int off = r - dd;
                        __stcs(dst + off + dd, src[(off >> 3) + dd]);
                    }
                    int adv = (rwc >> 3) + 1;
                    src += adv << 3;
                    yl  += (long)adv * D;
                    rem -= cnt;
                    wqq += wseg;
                    if (wqq >= WO) {
                        wqq = 0; ++hl;
                        if (--hl == 0) { ++nn; hl = HO; }
                        // recompute row base for the new (n, h)
                        long mrow = (yl - yw) / D;   // rows consumed beyond start
                        (void)mrow;
                        long mm = yl / D;            // flattened (n*h+h)*WO + w
                        nn = (int)(mm / Q);
                        int rh = (int)(mm - (long)nn * Q);
                        int hh = rh / WO;
                        wqq = rh - hh * WO;
                        yl  = (((long)nn * OC * HO + hh) * Q) + (long)wqq * D
                              + (long)(nt << 3) * Q + (long)(nof) * Q
                              + (long)(mof) * D + yw - yw;  // rebuilt below
                        yl  = yw; // fallback (never used for in-row tiles)
                        break;
                    }
                    // still inside same output row: rebase dst pointer
                    dst = out + yl;
                }
                __syncthreads();
            }
        } else {
            // 32-wide tiles: stage rows through qv, then copy
            const int wq0 = wq, dq0 = dq, nout0 = nout, hleft0 = hleft;
            long ybase = yrow + (long)(np << 5) * Q;
            float* dstb = out + ybase + (long)(wc << 5) * Q;
            #pragma unroll
            for (int mt = 0; mt < 4; ++mt) {
                float* q0 = qv + (wr << 9) + (mt << 7) + (g << 3) + (t << 1);
                float* q1 = q0 + 64;
                #pragma unroll
                for (int nt = 0; nt < 4; ++nt) {
                    float* cc = c[(mt << 2) + nt];
                    q0[nt << 3]       = cc[0];
                    q0[(nt << 3) + 1] = cc[1];
                    q1[nt << 3]       = cc[2];
                    q1[(nt << 3) + 1] = cc[3];
                }
                __syncthreads();
                float* d2 = dstb + (long)(mt << 4) * D;
                #pragma unroll
                for (int k = 0; k < PADQ; ++k) {
                    if (kvalid[k]) {
                        __stcs(d2 + kdst[k], qv[(wr << 9) + (mt << 7) + ksrc[k]]);
                    }
                }
                __syncthreads();
            }
            (void)wq0; (void)dq0; (void)nout0; (void)hleft0;
        }
    }
}

// ---------------------------------------------------------------------------
// Generic direct-convolution fallback (any K, stride, pad, dilation).
// ---------------------------------------------------------------------------
__global__ void conv3d_generic_fwd_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    long total,
    int C, int H, int W, int D,
    int OC, int HO, int WO,
    int KS, int stride, int pad, int dil)
{
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    int d = idx % D;
    long t = idx / D;
    int ww = t % WO; t /= WO;
    int h  = t % HO; t /= HO;
    int oc = t % OC;
    int n  = t / OC;

    float sum = bias ? bias[oc] : 0.0f;
    for (int c = 0; c < C; ++c) {
        for (int ky = 0; ky < KS; ++ky) {
            int ih = h * stride - pad + ky * dil;
            if (ih < 0 || ih >= H) continue;
            for (int kx = 0; kx < KS; ++kx) {
                int iw = ww * stride - pad + kx * dil;
                if (iw < 0 || iw >= W) continue;
                sum += x[(((long)(n * C + c) * H + ih) * W + iw) * D + d]
                     * w[(((long)oc * C + c) * KS + ky) * KS + kx];
            }
        }
    }
    out[idx] = sum;
}

// ---------------------------------------------------------------------------
torch::Tensor conv3d_fwd(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias_opt,
    int64_t stride, int64_t padding, int64_t dilation)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(weight.size(4) == 1, "kernel depth must be 1");

    auto xc = x.contiguous();
    auto wc = weight.contiguous();

    int N  = xc.size(0);
    int C  = xc.size(1);
    int H  = xc.size(2);
    int W  = xc.size(3);
    int D  = xc.size(4);
    int OC = wc.size(0);
    int KS = wc.size(2);

    int HO = (H + 2 * (int)padding - (int)dilation * (KS - 1) - 1) / (int)stride + 1;
    int WO = (W + 2 * (int)padding - (int)dilation * (KS - 1) - 1) / (int)stride + 1;

    auto out = torch::empty({N, OC, HO, WO, D}, xc.options());
    long total = (long)N * OC * HO * WO * D;
    if (total <= 0) return out;

    const float* bias_ptr =
        (bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0)
            ? bias_opt->data_ptr<float>()
            : nullptr;

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

    int Q = HO * WO * D;
    long M = (long)N * Q;
    int K = C * KS * KS;

    bool has_bias = (bias_ptr != nullptr);
    bool tc = (!has_bias && stride == 1 && dilation == 1 &&
               K >= 1 && K <= MAXK && OC % 8 == 0 && OC <= MAXOC &&
               D >= 1 && M % BM == 0 && Q % BM == 0 &&
               M <= 2000000000L / 4);

    static int configured_smem = -1;
    if (tc) {
        int dev = 0;
        cudaGetDevice(&dev);
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        tc = (prop.major > 8) || (prop.major == 8 && prop.minor >= 0);
        tc = tc && (prop.major >= 8);
        if (tc) {
            size_t smem_single = ((size_t)2 * BM * MAXK
                                  + 2 * (size_t)MAXK0 * OC) * 4;
            size_t smem_multi  = ((size_t)2 * BM * MAXK
                                  + 2 * (size_t)MAXK * MAXOC
                                  + BM * PADQ) * 4;
            size_t need = smem_multi > smem_single ? smem_multi : smem_single;
            if (need > (size_t)prop.sharedMemPerBlockOptin) {
                tc = false;
            } else {
                if (configured_smem != (int)need) {
                    cudaFuncSetAttribute(conv_tc_kernel<true>,
                        cudaFuncAttributeMaxDynamicSharedMemorySize,
                        (int)smem_single);
                    cudaFuncSetAttribute(conv_tc_kernel<false>,
                        cudaFuncAttributeMaxDynamicSharedMemorySize,
                        (int)smem_multi);
                    configured_smem = (int)need;
                }
            }
        }
    }

    if (tc) {
        auto opts = xc.options().dtype(at::kInt);
        auto wh_i = torch::empty({(long)OC * K}, opts);
        auto wl_i = torch::empty({(long)OC * K}, opts);
        auto wT = wc.transpose(0, 4).contiguous();  // [1, C, KS, KS, OC]
        auto wk = wT.view({(long)K, OC});

        int wt = OC < 128 ? 128 : (OC < 1024 ? ((OC + 31) / 32) * 32 : 1024);
        splitw_kernel<<<1, wt, 0, stream>>>(
            wk.data_ptr<float>(),
            (unsigned*)wh_i.data_ptr<int>(),
            (unsigned*)wl_i.data_ptr<int>(), OC, K);

        dim3 grid((unsigned)(M / BM), 1, 1);
        if (K <= 8) {
            size_t smem = ((size_t)2 * BM * MAXK + 2 * (size_t)MAXK0 * OC) * 4;
            conv_tc_kernel<true><<<grid, 128, smem, stream>>>(
                xc.data_ptr<float>(),
                (unsigned*)wh_i.data_ptr<int>(),
                (unsigned*)wl_i.data_ptr<int>(),
                out.data_ptr<float>(),
                (int)M, OC, K, Q, HO, WO, D, C, H, W, KS, (int)padding);
        } else {
            size_t smem = ((size_t)2 * BM * MAXK + 2 * (size_t)MAXK * MAXOC
                           + BM * PADQ) * 4;
            conv_tc_kernel<false><<<grid, 128, smem, stream>>>(
                xc.data_ptr<float>(),
                (unsigned*)wh_i.data_ptr<int>(),
                (unsigned*)wl_i.data_ptr<int>(),
                out.data_ptr<float>(),
                (int)M, OC, K, Q, HO, WO, D, C, H, W, KS, (int)padding);
        }
    } else {
        int block = 256;
        long grid = (total + block - 1) / block;
        conv3d_generic_fwd_kernel<<<(unsigned)grid, block, 0, stream>>>(
            xc.data_ptr<float>(), wc.data_ptr<float>(), bias_ptr,
            out.data_ptr<float>(), total,
            C, H, W, D, OC, HO, WO, KS,
            (int)stride, (int)padding, (int)dilation);
    }

    return out;
}
"""

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor conv3d_fwd(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias,
    int64_t stride, int64_t padding, int64_t dilation);
"""

_conv3d_mod = load_inline(
    name="custom_conv3d_tc_v5",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_fwd"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Conv3d with kernel (kernel_size, kernel_size, 1) implemented with custom
    CUDA kernels: a 3xTF32 tensor-core implicit GEMM for the common case
    (near-fp32 accuracy, tensor-core throughput) plus a direct scalar
    fallback kernel. Falls back to the reference cuDNN implementation for
    unsupported configurations (e.g. bias or non-unit stride).
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int,
                 stride: int = 1, padding: int = 0, dilation: int = 1,
                 groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(
            in_channels, out_channels, (kernel_size, kernel_size, 1),
            stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        c = self.conv3d
        s, p, d = c.stride, c.padding, c.dilation
        if (x.is_cuda and x.dtype == torch.float32 and c.groups == 1
                and c.kernel_size[2] == 1
                and s[0] == s[1] and p[0] == p[1] and d[0] == d[1]
                and s[2] == 1 and p[2] == 0 and d[2] == 1):
            return _conv3d_mod.conv3d_fwd(x, c.weight, c.bias, s[0], p[0], d[0])
        return c(x)
```

Wait — I left a messy half-finished section in the SINGLE_K store path (the `while` loop with a bogus `break`). Since the actual test case (K=27) uses the multi-k-step path, let me clean that up properly and ship a correct, simpler SINGLE_K path that requires in-row tiles when K≤8 (guarded on the host: Q % 128 == 0), and otherwise falls back to the multi-k-step path with K≤8 (which handles boundaries generally). Here is the final, corrected version:

```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

# Custom CUDA kernels for Conv3d with kernel (K, K, 1).
#
# The operator is an independent 2D convolution over dims (2, 3) for every
# index of the innermost dim 4 ("depth"), i.e. a GEMM
#   out[(n,h,w,d), oc] = sum_{c,ky,kx} x[n,c,h+ky,w+kx,d] * W[oc,c,ky,kx]
# with M = N*HO*WO*D rows, N = OC columns, K = C*KS*KS.
#
# Fast path (sm_80+, K <= 32, OC % 8 == 0, M % 128 == 0): an implicit-GEMM
# tensor-core kernel using mma.m16n8k8.tf32 with the 3xTF32 decomposition
# (hi*hi + hi*lo + lo*hi), which delivers near-fp32 accuracy at tensor-core
# throughput:
#   * 128x64 block tile, 4 warps (2x2), warp tile 64x32,
#   * A staged through shared memory with im2col gather, coalesced global
#     reads, and an XOR swizzle that makes both staging and mma fragment
#     loads bank-conflict-free,
#   * B (weights, K x OC) is split into tf32 hi/lo parts once per launch and
#     kept in shared memory,
#   * fp32 accumulators (16 per thread) are drained through a padded shared
#     staging buffer so global stores are fully coalesced 128B transactions,
#   * when output tiles never cross row boundaries (Q % 128 == 0) the staged
#     output is a single contiguous span per column, so stores use one
#     conflict-free linear copy; otherwise a predicated per-element path with
#     boundary-safe clamping is used.
#
# A direct scalar fallback kernel handles every other configuration, and any
#