```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 OH 8
#define OW 32
#define SDIM 64
#define NBUF 2

__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 cp4(uint32_t dst, const float* src, bool pred) {
    int sz = pred ? 4 : 0;
    asm volatile("cp.async.ca.shared.global [%0], [%1], 4, %2;\n"
                 :: "r"(dst), "l"(src), "r"(sz));
}
__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)); }

// ---------------------------------------------------------------------------
// Direct conv kernel: NCHW, stride 1, C=32, N=64, TF32 tensor cores.
// grid: (tiles_h * tiles_w, batch). CTA: 256 threads computes
// out[b, 0:64, H0:H0+8, W0:W0+32].
// Halo in smem: Hs[c][r][s], strides (Cs, S, 1), channel-pair interleaved.
// Weights: per-tap images, double buffered, cp.async prefetched.
// ---------------------------------------------------------------------------
__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, int Cs, int R, int S, int Wstride) {
    extern __shared__ float smem[];
    const int xs_floats = 32 * Cs;
    const int smem_floats = xs_floats + NBUF * Wstride + 64;
    float* Hs = smem;
    float* Wp = smem + xs_floats;
    float* Bs = smem + xs_floats + NBUF * Wstride;

    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) * OH;
    const int W0 = (tile % ntiles_w) * OW;
    const int b = blockIdx.y;

    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 into Hs (channel-pair interleaved), cp.async ----
    {
        const int hh0 = H0 - ph;
        const int ww0 = W0 - pw;
        const int S4 = S >> 2;
        const int npair = 16 * R * S4;
        const uint32_t hbase = smem_u32(Hs);
        for (int idx = tid; idx < npair; idx += NT) {
            int s4 = idx % S4;
            int t2 = idx / S4;
            int r = t2 % R;
            int p = t2 / R;
            int hh = hh0 + r;
            int ww = ww0 + s4 * 4;
            int wp2 = ww + ((pw & 1) ? -1 : 0);
            bool rok = ((unsigned)hh < (unsigned)Hin);
            const float* srcc = xb + (long)(2 * p) * inHW + (long)hh * Win;
            bool h0 = rok && ((unsigned)(wp2) < (unsigned)Win);
            bool h1 = rok && ((unsigned)(wp2 + 1) < (unsigned)Win);
            bool h2 = rok && ((unsigned)(wp2 + 2) < (unsigned)Win);
            bool h3 = rok && ((unsigned)(wp2 + 3) < (unsigned)Win);
            uint32_t dst = hbase + ((2 * p) * Cs + r * S + s4 * 4) * 4;
            cp4(dst + 0, srcc + wp2, h0);
            cp4(dst + 4, srcc + inHW + wp2, h1);
            cp4(dst + 8, srcc + wp2 + 2, h2);
            cp4(dst + 12, srcc + inHW + wp2 + 2, h3);
        }
        cp_commit();
    }

    // ---- prefetch tap 0 and tap 1 weight images ----
    {
        const float* src = wt;
        uint32_t dst = smem_u32(Wp);
        int n16 = Wstride >> 2;
        for (int idx = tid; idx < n16; idx += NT) cp16(dst + idx * 16, src + idx * 4);
        cp_commit();
        if (ntaps > 1) {
            const float* src1 = wt + (long)Wstride;
            uint32_t dst1 = smem_u32(Wp + Wstride);
            for (int idx = tid; idx < n16; idx += NT) 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 r0 = warp & 7;
    const int w0 = (warp >> 3) * 16;
    const int pxbase = g * S + w0;
    const int Bbase = g * 2;
    const int Brow = tq * (8 * 64);

    cp_wait<1>();   // halo + tap0 ready (tap1 may be 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) * Wstride;

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

        #pragma unroll
        for (int ks = 0; ks < 4; ++ks) {
            const int kbase = ks * 8 * Cs;
            const float* pA = Hs + kbase + tq * 2 * Cs + aoff + pxbase;
            const float* pB = Hs + kbase + (tq * 2 + 1) * Cs + aoff + pxbase;
            uint32_t a_0[4], a_1[4];
            float2 fa, fb;
            fa = *reinterpret_cast<const float2*>(pA);
            fb = *reinterpret_cast<const float2*>(pB);
            a_0[0] = __float_as_uint(fa.x); a_0[1] = __float_as_uint(fb.x);
            a_1[0] = __float_as_uint(fa.y); a_1[1] = __float_as_uint(fb.y);
            fa = *reinterpret_cast<const float2*>(pA + 4 * Cs);
            fb = *reinterpret_cast<const float2*>(pB + 4 * Cs);
            a_0[2] = __float_as_uint(fa.x); a_0[3] = __float_as_uint(fb.x);
            a_1[2] = __float_as_uint(fa.y); a_1[3] = __float_as_uint(fb.y);

            const float* bp = wimg + Brow + 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) * Wstride;
            uint32_t dst = smem_u32(Wp + (tap & 1) * Wstride);
            int n16 = Wstride >> 2;
            for (int idx = tid; idx < n16; idx += NT) cp16(dst + idx * 16, src + idx * 4);
        }
        cp_commit();
        __syncthreads();
    }

    // ---- epilogue: stage via smem (reuse halo), coalesced stores ----
    float* SW = Hs;   // reuse: 64 * 272 floats needed
    const int SW_STRIDE = 272;
    if (bias && tid < 64) Bs[tid] = bias[tid];
    __syncthreads();

    const float* biasp = bias ? Bs : nullptr;
    #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) {
                int rr = g + ((v >> 1) ? 8 : 0);
                int n = nt * 8 + 2 * tq + (v & 1);
                int hh = rr & 7;
                int ww = w0 + 2 * mt + (rr >> 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(OH, Ho - H0);
    const int w_lim = min(OW, 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 + tf32 rounding:
// w[N=64][C=32][kh][kw] -> wt[tap][ks][csub(8)][n(64)] (permuted taps first).
// ---------------------------------------------------------------------------
__global__ void repack_w_kernel(const float* __restrict__ w, float* __restrict__ wt,
                                int kh, int kw, int Wstride, int ntaps) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = ntaps * 2048;
    if (idx >= total) return;
    const int q = idx & 2047;
    const int tapp = idx >> 11;
    const int i = tapp / kw;
    const int jp = tapp - i * kw;
    const int j = (jp == kw - 1) ? jp : (jp ^ 1);
    const int tap = i * kw + j;
    const int n = q & 63;
    const int c = q >> 6;              // 0..31
    const float v = w[((long)n * 32 + c) * ntaps + tap];
    wt[tapp * Wstride + c * 64 + n] = rnd_tf32(v);
}

torch::Tensor repack_w(torch::Tensor w, int64_t Wstride) {
    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 * Wstride}, 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, (int)Wstride, ntaps);
    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 R = OH + ((int)kh - 1) * (int)dh;
    const int S = OW + ((int)kw - 1) * (int)dw;
    TORCH_CHECK(R <= 64 && S <= SDIM && (S % 4) == 0, "halo too large");
    const int Cs = R * S;
    const int xs_floats = 32 * Cs;
    const int Wstride = 2048;
    const int smem = (xs_floats + NBUF * Wstride + 64) * 4;
    TORCH_CHECK(smem <= 160 * 1024, "smem too large");

    const int ntiles_h = (Ho + OH - 1) / OH;
    const int ntiles_w = (Wo + OW - 1) / OW;
    dim3 grid(ntiles_h * ntiles_w, Bc, 1);

    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, Cs, R, S, Wstride);
    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, int64_t Wstride);
"""

_conv_ext = load_inline(
    name="conv2d_tc_v6",
    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
                kh, kw = int(w.shape[2]), int(w.shape[3])
                dh, dw = self._dilation
                R = 8 + (kh - 1) * dh
                S = 32 + (kw - 1) * dw
                smem = (32 * R * S + 2 * 2048 + 64) * 4
                ok = (cap[0] >= 8 and int(w.shape[0]) == 64 and int(w.shape[1]) == 32
                      and self._stride == (1, 1)
                      and dh >= 1 and dw >= 1
                      and R <= 64 and S <= 64 and (S % 4) == 0
                      and smem <= 160 * 1024)
            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(), 2048)
                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)
```