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

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

#define NT 256
#define CS 680
#define XS_FLOATS (32 * CS)
#define TAPIMAGE 2176
#define SW_STRIDE 260
#define SW_FLOATS (64 * SW_STRIDE)

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

__device__ __forceinline__ void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                                         uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
                                         uint32_t b0, uint32_t 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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

__device__ __forceinline__ uint32_t smem_u32(const void* p) {
    return (uint32_t)__cvta_generic_to_shared(p);
}

__device__ __forceinline__ void cp16(uint32_t dst, const float* src) {
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n" :: "r"(dst), "l"(src));
}
__device__ __forceinline__ void cp_commit() { asm volatile("cp.async.commit_group;\n"); }
template <int N>
__device__ __forceinline__ void cp_wait() { asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); }

// ---------------------------------------------------------------------------
// Main conv kernel: NCHW, stride 1, C=32, N=64, TF32 tensor cores.
// grid: (tiles, batch). CTA: 256 threads computes out[b, 0:64, H0:H0+8, W0:W0+32].
// smem: Xs = x halo [c][r][s] strides (680, 41, 1), tf32-rounded;
//       Wp = double-buffered per-tap weight images (pre-packed, tf32).
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NT)
conv2d_tc_kernel(const float* __restrict__ x, const float* __restrict__ wt,
                 const float* __restrict__ bias, float* __restrict__ out,
                 int Bc, int Hin, int Win, int Ho, int Wo,
                 int kh, int kw, int dh, int dw, int ph, int pw,
                 int ntiles_w) {
    extern __shared__ float smem[];
    float* Xs = smem;
    float* Wp = smem + XS_FLOATS;
    float* SW = Xs;

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g = lane >> 2;
    const int tq = lane & 3;

    const int tile = blockIdx.x;
    const int H0 = (tile / ntiles_w) * 8;
    const int W0 = (tile % ntiles_w) * 32;
    const int b = blockIdx.y;

    const int R = 8 + (kh - 1) * dh;
    const int S = 32 + (kw - 1) * dw;
    const int ntaps = kh * kw;
    const long inHW = (long)Hin * Win;
    const long outHW = (long)Ho * Wo;
    const float* xb = x + (long)b * 32 * inHW;

    // ---- load x halo, tf32-rounded, into Xs ----
    {
        const int hh0 = H0 - ph;
        const int ww0 = W0 - pw;
        const bool aligned = (((uintptr_t)xb & 15) == 0) && ((Win & 3) == 0)
                             && ((ww0 & 3) == 0);
        const int total = 32 * R * S;
        if (aligned && (S % 4) == 0) {
            const int n4 = total >> 2;
            for (int idx = tid; idx < n4; idx += NT) {
                int s4 = idx % (S >> 2);
                int t2 = idx / (S >> 2);
                int r = t2 & (R - 1);   // R is a power of two only when dh*... guard below
                int c = t2 / R;
                // generic (R not necessarily pow2): recompute
                r = t2 % R;
                c = t2 / R;
                int hh = hh0 + r;
                int ww = ww0 + s4 * 4;
                float4 v = make_float4(0.f, 0.f, 0.f, 0.f);
                if ((unsigned)hh < (unsigned)Hin) {
                    if (ww >= 0 && ww + 3 < Win) {
                        v = *reinterpret_cast<const float4*>(
                            xb + (long)c * inHW + (long)hh * Win + ww);
                    } else {
                        float tmp[4] = {0.f, 0.f, 0.f, 0.f};
                        #pragma unroll
                        for (int q = 0; q < 4; ++q) {
                            int wq = ww + q;
                            if ((unsigned)wq < (unsigned)Win)
                                tmp[q] = __ldg(xb + (long)c * inHW + (long)hh * Win + wq);
                        }
                        v = make_float4(tmp[0], tmp[1], tmp[2], tmp[3]);
                    }
                }
                float4 d;
                d.x = rnd_tf32(v.x); d.y = rnd_tf32(v.y);
                d.z = rnd_tf32(v.z); d.w = rnd_tf32(v.w);
                *reinterpret_cast<float4*>(Xs + c * CS + r * 41 + s4 * 4) = d;
            }
        } else {
            for (int idx = tid; idx < total; idx += NT) {
                int s = idx % S;
                int t2 = idx / S;
                int r = t2 % R;
                int c = t2 / R;
                int hh = hh0 + r;
                int ww = ww0 + s;
                float v = 0.f;
                if ((unsigned)hh < (unsigned)Hin && (unsigned)ww < (unsigned)Win)
                    v = __ldg(xb + (long)c * inHW + (long)hh * Win + ww);
                Xs[c * CS + r * 41 + s] = rnd_tf32(v);
            }
        }
    }

    // ---- prefetch tap 0 and tap 1 weight images ----
    {
        const float* src = wt;
        uint32_t dst = smem_u32(Wp);
        #pragma unroll
        for (int q = 0; q < 2; ++q) {
            int idx = q * NT + tid;
            if (idx < 544) cp16(dst + idx * 16, src + idx * 4);
        }
        cp_commit();
        if (ntaps > 1) {
            const float* src1 = wt + (long)TAPIMAGE;
            uint32_t dst1 = smem_u32(Wp + TAPIMAGE);
            #pragma unroll
            for (int q = 0; q < 2; ++q) {
                int idx = q * NT + tid;
                if (idx < 544) cp16(dst1 + idx * 16, src1 + idx * 4);
            }
        }
        cp_commit();
    }

    // ---- accumulators ----
    float acc[2][8][4];
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt)
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt)
            #pragma unroll
            for (int v = 0; v < 4; ++v) acc[mt][nt][v] = 0.f;

    const int mt0 = warp * 2;
    const int Abase0 = g * 41 + 2 * mt0;
    const int Abase1 = g * 41 + 2 * (mt0 + 1);
    const int Bbase = g * 2;

    cp_wait<1>();   // X ready, tap0 ready, tap1 maybe in flight
    __syncthreads();

    for (int tap = 0; tap < ntaps; ++tap) {
        if (tap + 1 < ntaps) cp_wait<0>();  // wait tap+1 prefetch
        __syncthreads();
        const float* wimg = Wp + (tap & 1) * TAPIMAGE;

        const int i = tap / kw;
        const int j = tap - i * kw;
        const int aoff = i * dh * 41 + j * dw;

        #pragma unroll
        for (int ks = 0; ks < 4; ++ks) {
            const int kbase = ks * 8 * CS;
            const float* a0p = Xs + kbase + tq * CS + aoff + Abase0;
            const float* a1p = Xs + kbase + tq * CS + aoff + Abase1;
            const float* a2p = a0p + 4 * CS;
            const float* a3p = a1p + 4 * CS;
            float2 f0 = *reinterpret_cast<const float2*>(a0p);
            float2 f1 = *reinterpret_cast<const float2*>(a1p);
            float2 f2 = *reinterpret_cast<const float2*>(a2p);
            float2 f3 = *reinterpret_cast<const float2*>(a3p);
            uint32_t a_0[4] = {__float_as_uint(f0.x), __float_as_uint(f0.y),
                               __float_as_uint(f2.x), __float_as_uint(f2.y)};
            uint32_t a_1[4] = {__float_as_uint(f1.x), __float_as_uint(f1.y),
                               __float_as_uint(f3.x), __float_as_uint(f3.y)};
            const float* bp = wimg + ks * 544 + tq * 136 + Bbase;
            #pragma unroll
            for (int nt = 0; nt < 8; ++nt) {
                float2 bb = *reinterpret_cast<const float2*>(bp + nt * 16);
                uint32_t b0 = __float_as_uint(bb.x);
                uint32_t b1 = __float_as_uint(bb.y);
                mma_tf32(acc[0][nt][0], acc[0][nt][1], acc[0][nt][2], acc[0][nt][3],
                         a_0[0], a_0[1], a_0[2], a_0[3], b0, b1);
                mma_tf32(acc[1][nt][0], acc[1][nt][1], acc[1][nt][2], acc[1][nt][3],
                         a_1[0], a_1[1], a_1[2], a_1[3], b0, b1);
            }
        }
        // prefetch tap+2 into the buffer just freed
        if (tap + 2 < ntaps) {
            const float* src = wt + (long)(tap + 2) * TAPIMAGE;
            uint32_t dst = smem_u32(Wp + (tap & 1) * TAPIMAGE);
            #pragma unroll
            for (int q = 0; q < 2; ++q) {
                int idx = q * NT + tid;
                if (idx < 544) cp16(dst + idx * 16, src + idx * 4);
            }
        }
        cp_commit();
        __syncthreads();
    }

    // ---- epilogue: stage via smem, coalesced stores ----
    if (bias && tid < 64) SW[tid] = bias[tid];
    __syncthreads();

    const float* biasp = bias ? SW : nullptr;
    #pragma unroll
    for (int mt = 0; mt < 2; ++mt) {
        const int mtile = mt0 + mt;
        #pragma unroll
        for (int nt = 0; nt < 8; ++nt) {
            #pragma unroll
            for (int v = 0; v < 4; ++v) {
                int r = g + ((v >> 1) ? 8 : 0);
                int n = nt * 8 + 2 * tq + (v & 1);
                int hh = r & 7;
                int ww = 2 * mtile + (r >> 3);
                float val = acc[mt][nt][v] + (biasp ? biasp[n] : 0.f);
                SW[n * SW_STRIDE + ww * 8 + hh] = val;
            }
        }
    }
    __syncthreads();

    const int h_lim = min(8, Ho - H0);
    const int w_lim = min(32, Wo - W0);
    float* outb = out + (long)b * 64 * outHW;
    const int m4 = tid & 63;
    const int n0t = tid >> 6;
    const int hh = m4 >> 3;
    const int ww = (m4 & 7) * 4;
    if (hh < h_lim) {
        const int wleft = w_lim - ww;
        float* ob = outb + (long)(H0 + hh) * Wo + (W0 + ww);
        if (wleft >= 4 && ((Wo & 3) == 0)) {
            #pragma unroll
            for (int nn = 0; nn < 16; ++nn) {
                const int n = n0t + nn * 4;
                const float4 v = *reinterpret_cast<const float4*>(
                    SW + n * SW_STRIDE + m4 * 4);
                *reinterpret_cast<float4*>(ob + (long)n * outHW) = v;
            }
        } else if (wleft > 0) {
            #pragma unroll
            for (int nn = 0; nn < 16; ++nn) {
                const int n = n0t + nn * 4;
                const float* src = SW + n * SW_STRIDE + m4 * 4;
                float* dst = ob + (long)n * outHW;
                for (int q = 0; q < wleft; ++q) dst[q] = src[q];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Weight repack: w[N=64][C=32][kh][kw] -> wt[tap][ks][t][n][2] (tf32 rounded),
// image stride per tap = 2176 floats: ((ks*4 + t)*136 + n*2 + {0,1}).
// ---------------------------------------------------------------------------
__global__ void repack_w_kernel(const float* __restrict__ w, float* __restrict__ wt,
                                int kh, int kw) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int ntaps = kh * kw;
    const int total = ntaps * 2048;
    if (idx >= total) return;
    const int q = idx & 2047;
    const int tap = idx >> 11;
    const int half = q & 1;
    const int n = (q >> 1) & 63;
    const int t = (q >> 7) & 3;
    const int ks = q >> 9;
    const int c = ks * 8 + t + half * 4;
    const float v = w[((long)n * 32 + c) * ntaps + tap];
    const int dst = tap * TAPIMAGE + ((ks * 4 + t) * 136 + n * 2 + half);
    wt[dst] = rnd_tf32(v);
}

torch::Tensor repack_w(torch::Tensor w) {
    TORCH_CHECK(w.is_cuda() && w.dtype() == torch::kFloat && w.dim() == 4, "bad w");
    auto wc = w.contiguous();
    const int N = (int)wc.size(0);
    const int C = (int)wc.size(1);
    const int kh = (int)wc.size(2);
    const int kw = (int)wc.size(3);
    TORCH_CHECK(N == 64 && C == 32, "need N=64, C=32");
    const int ntaps = kh * kw;
    auto wt = torch::empty({(long)ntaps * TAPIMAGE}, wc.options());
    const int total = ntaps * 2048;
    const int threads = 256;
    const int blocks = (total + threads - 1) / threads;
    auto stream = at::cuda::getCurrentCUDAStream();
    repack_w_kernel<<<blocks, threads, 0, stream>>>(
        wc.data_ptr<float>(), wt.data_ptr<float>(), kh, kw);
    return wt;
}

torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor wt, torch::Tensor bias,
                             int64_t ph, int64_t pw, int64_t dh, int64_t dw,
                             int64_t kh, int64_t kw) {
    auto xc = x.contiguous();
    auto wtc = wt.contiguous();
    const int Bc = (int)xc.size(0);
    const int Hin = (int)xc.size(2);
    const int Win = (int)xc.size(3);
    const int Ho = Hin + 2 * (int)ph - (int)dh * ((int)kh - 1);
    const int Wo = Win + 2 * (int)pw - (int)dw * ((int)kw - 1);
    TORCH_CHECK(Ho > 0 && Wo > 0, "bad output size");
    auto out = torch::empty({Bc, 64, Ho, Wo}, xc.options());

    const int ntiles_h = (Ho + 7) / 8;
    const int ntiles_w = (Wo + 31) / 32;
    dim3 grid(ntiles_h * ntiles_w, Bc, 1);

    const int smem = (XS_FLOATS + 2 * TAPIMAGE) * 4;
    static int configured = 0;
    if (!configured) {
        cudaFuncSetAttribute(conv2d_tc_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             160 * 1024);
        configured = 1;
    }
    const float* bias_ptr = (bias.defined() && bias.numel() > 0)
                                ? bias.data_ptr<float>() : nullptr;
    auto stream = at::cuda::getCurrentCUDAStream();
    conv2d_tc_kernel<<<grid, NT, smem, stream>>>(
        xc.data_ptr<float>(), wtc.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(), Bc, Hin, Win, Ho, Wo,
        (int)kh, (int)kw, (int)dh, (int)dw, (int)ph, (int)pw, ntiles_w);
    return out;
}
"""

_cpp_src = """
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor wt, torch::Tensor bias,
                             int64_t ph, int64_t pw, int64_t dh, int64_t dw,
                             int64_t kh, int64_t kw);
torch::Tensor repack_w(torch::Tensor w);
"""

_conv_ext = load_inline(
    name="conv2d_tc_v5",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv2d_forward", "repack_w"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


def _pair(v):
    if isinstance(v, int):
        return (v, v)
    return (int(v[0]), int(v[1]))


class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: tuple,
                 stride: int = 1, padding: tuple = (0, 0), dilation: tuple = (1, 1),
                 bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size,
                                stride=stride, padding=padding,
                                dilation=dilation, bias=bias)
        self._stride = _pair(stride)
        self._padding = _pair(padding)
        self._dilation = _pair(dilation)
        self._cached_wt = None
        self._cached_version = -1
        self._fast = None
        self._empty_bias = None

    def _fast_ok(self, x):
        if self._fast is None:
            ok = False
            try:
                cap = torch.cuda.get_device_capability(x.device)
                w = self.conv2d.weight
                ok = (cap[0] >= 8 and w.shape[0] == 64 and w.shape[1] == 32
                      and self._stride == (1, 1)
                      and 8 + (w.shape[2] - 1) * self._dilation[0] <= 16
                      and 32 + (w.shape[3] - 1) * self._dilation[1] <= 40
                      and self._dilation[0] >= 1 and self._dilation[1] >= 1)
            except Exception:
                ok = False
            self._fast = ok
        return self._fast

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.conv2d.weight
        if (x.is_cuda and x.dtype == torch.float32 and w.dtype == torch.float32
                and x.dim() == 4 and self._fast_ok(x)):
            if (self._cached_wt is None or self._cached_version != w._version
                    or self._cached_wt.device != w.device):
                self._cached_wt = _conv_ext.repack_w(w.detach())
                self._cached_version = w._version
            b = self.conv2d.bias
            if b is None:
                if self._empty_bias is None:
                    self._empty_bias = w.new_empty(0)
                b = self._empty_bias
            return _conv_ext.conv2d_forward(
                x, self._cached_wt, b,
                self._padding[0], self._padding[1],
                self._dilation[0], self._dilation[1],
                w.shape[2], w.shape[3])
        return self.conv2d(x)
```