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

# ----------------------------------------------------------------------------
# Custom near-FP32-accurate tensor-core CUDA kernel for the dilated 2D conv
# (specialized: kernel=(5,9), stride=1, dilation=(2,3), any padding/channels).
#
# Method: implicit-GEMM style direct convolution on fp16 tensor cores using
# an error-compensated hi/lo split of both operands:
#     x = x_hi + x_lo      (each fp16; combined representation error ~2^-22)
#     a.b ~= a_hi*b_hi + a_lo*b_hi + a_hi*b_lo     (3 mmas, FP32 accumulate)
# This keeps accuracy in the same class as a plain FP32 kernel (error ~1e-7,
# far below TF32's ~1e-4), while running on the fp16 tensor cores.
#
# v6 improvements over v5 (which was L1-bandwidth bound on weight loads):
#   * Weights are pre-transposed into a fully-coalesced fragment layout
#     [tap][oc/8][channel-pair][8 oc] so every warp B-fragment load is a
#     single 128B cache-line access (v5 needed 4 sectors per load).
#   * Uses mma.m16n8k16 (sm_80+): half the instruction count and half the
#     accumulator-touches of m16n8k8 for the same math.
#   * Larger warp tile (m32 x n64): 2 m16 tiles x 8 n8 tiles per warp,
#     giving much better compute-per-byte on the fragment loads.
#   * The three error-compensation terms are issued as separate MMA sweeps
#     over all 16 accumulator quads, so dependent MMAs are ~16 instructions
#     apart (long dependency distance hides tensor-core latency).
#   * 4-warp blocks computing an 8x16 output-pixel tile for all 64 output
#     channels; the zero-padded dilated input tile (16x40 x 16 channels) is
#     staged in shared memory as bank-conflict-free hi/lo half2 plane pairs.
#
# Anything outside the supported configuration falls back to at::conv2d.
# ----------------------------------------------------------------------------

conv2d_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>

#define KH        5
#define KW        9
#define NTAP      (KH * KW)                 // 45
#define DH        2
#define DW        3
#define TILE_OH   8
#define TILE_OW   16
#define IN_H      (TILE_OH + (KH - 1) * DH) // 16
#define IN_W      (TILE_OW + (KW - 1) * DW) // 40
#define PIX       (IN_H * IN_W)             // 640
#define PLANE     648                       // padded: PLANE % 32 == 8
#define CHUNK     16                        // input channels staged at once
#define CPAIRS    (CHUNK / 2)               // 8 half2 planes (hi and lo)
#define NTHREADS  128                       // 4 warps

// m16n8k16 f16 -> f32 tensor core MMA (sm_80+)
#define MMA16816(C, A, B)                                                   \
    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]))

// ---------------------------------------------------------------------------
// Pre-kernel: transpose + split weights [OC][IC][45] into hi/lo half2
// channel-pair arrays laid out as [tap][oc8][pair][8] (32 contiguous words
// per (tap, oc8-tile, 4 pairs) so fragment loads are perfectly coalesced).
// Output index: w[((tap*NJ + j)*ICP + p)*8 + o]
// ---------------------------------------------------------------------------
__global__ void wsplit_kernel(const float* __restrict__ w,
                              unsigned* __restrict__ w_hi,
                              unsigned* __restrict__ w_lo,
                              const int IC, const int OC,
                              const int ICP, const int NJ) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int total = NTAP * NJ * ICP * 8;
    if (idx >= total) return;
    const int o   = idx & 7;                 // oc within the n8 tile
    const int t1  = idx >> 3;
    const int p   = t1 % ICP;                // input channel pair
    const int t2  = t1 / ICP;
    const int j   = t2 % NJ;                 // n8 tile
    const int tap = t2 / NJ;
    const int oc  = j * 8 + o;
    const int ic0 = 2 * p;
    const int ic1 = ic0 + 1;

    float v0 = 0.0f, v1 = 0.0f;
    if (oc < OC) {
        if (ic0 < IC) v0 = __ldg(w + ((long)oc * IC + ic0) * NTAP + tap);
        if (ic1 < IC) v1 = __ldg(w + ((long)oc * IC + ic1) * NTAP + tap);
    }
    const __half h0 = __float2half_rn(v0);
    const __half h1 = __float2half_rn(v1);
    const __half l0 = __float2half_rn(v0 - __half2float(h0));
    const __half l1 = __float2half_rn(v1 - __half2float(h1));
    __half2 hh = __halves2half2(h0, h1);
    __half2 ll = __halves2half2(l0, l1);
    w_hi[idx] = *reinterpret_cast<unsigned*>(&hh);
    w_lo[idx] = *reinterpret_cast<unsigned*>(&ll);
}

// ---------------------------------------------------------------------------
// Main kernel.
// grid: (ceil(WO/16), ceil(HO/8), N*ceil(OC/64));  block: 128 threads.
// Each warp computes a 2-row x 16-col x 64-oc tile (m32 x n64).
// ---------------------------------------------------------------------------
__global__ void __launch_bounds__(NTHREADS)
conv2d_f16x3_kernel(const float* __restrict__ x,
                    const unsigned* __restrict__ w_hi,  // [tap][oc8][pair][8]
                    const unsigned* __restrict__ w_lo,
                    const float* __restrict__ bias,     // may be null
                    float* __restrict__ out,
                    const int N, const int IC, const int H, const int W,
                    const int OC, const int HO, const int WO,
                    const int ICP, const int NJ,
                    const int pad_h, const int pad_w) {
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    __shared__ unsigned s_hi[CPAIRS * PLANE];
    __shared__ unsigned s_lo[CPAIRS * PLANE];

    const int tid  = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int lm0  = lane >> 2;   // 0..7  (pixel offset within m16 tile)
    const int lk0  = lane & 3;    // 0..3  (channel-pair group)

    const int oc_blocks = (OC + 63) >> 6;
    const int bz  = blockIdx.z;
    const int n   = bz / oc_blocks;
    const int oc0 = (bz % oc_blocks) << 6;
    const int oh0 = blockIdx.y * TILE_OH;
    const int ow0 = blockIdx.x * TILE_OW;

    // This warp's two output rows within the tile.
    const int r0 = 2 * warp;
    const int r1 = 2 * warp + 1;

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

    const float* x_base = x + (long)n * IC * H * W;
    const int in_r0 = oh0 - pad_h;
    const int in_c0 = ow0 - pad_w;

    for (int ic0 = 0; ic0 < IC; ic0 += CHUNK) {
        // ---- Stage zero-padded input tile, split into hi/lo half2 pairs ---
#pragma unroll 4
        for (int idx = tid; idx < CPAIRS * PIX; idx += NTHREADS) {
            const int p   = idx / PIX;
            const int rem = idx - p * PIX;
            const int sr  = rem / IN_W;
            const int sc  = rem - sr * IN_W;
            const int ih  = in_r0 + sr;
            const int iw  = in_c0 + sc;
            float v0 = 0.0f, v1 = 0.0f;
            const int icA = ic0 + 2 * p;
            if ((unsigned)ih < (unsigned)H && (unsigned)iw < (unsigned)W
                && icA < IC) {
                const long off = ((long)icA * H + ih) * W + iw;
                v0 = __ldg(x_base + off);
                if (icA + 1 < IC) v1 = __ldg(x_base + (long)H * W + off);
            }
            const __half h0 = __float2half_rn(v0);
            const __half h1 = __float2half_rn(v1);
            const __half l0 = __float2half_rn(v0 - __half2float(h0));
            const __half l1 = __float2half_rn(v1 - __half2float(h1));
            __half2 hh = __halves2half2(h0, h1);
            __half2 ll = __halves2half2(l0, l1);
            const int so = p * PLANE + sr * IN_W + sc;
            s_hi[so] = *reinterpret_cast<unsigned*>(&hh);
            s_lo[so] = *reinterpret_cast<unsigned*>(&ll);
        }
        __syncthreads();

        // ---- Accumulate over the 45 taps ----------------------------------
        const int icp0 = ic0 >> 1;
        const unsigned* wp_h = w_hi + ((long)icp0 + lk0) * 8 + lm0;
        const unsigned* wp_l = w_lo + ((long)icp0 + lk0) * 8 + lm0;
        const int bstride = 8 * ICP * NJ;   // words per tap

#pragma unroll 1
        for (int tap = 0; tap < NTAP; ++tap) {
            const int kg = tap / KW;
            const int kw = tap - kg * KW;
            const int sr  = r0 + DH * kg;
            const int sc0 = DW * kw;

            // A fragments (hi/lo) for the two m16 tiles.
            unsigned ah[2][4], al[2][4];
            {
                const unsigned* sh = s_hi + sr * IN_W + sc0 + lm0;
                const unsigned* sl = s_lo + sr * IN_W + sc0 + lm0;
                const int plo = lk0 * PLANE;
                const int phi = (lk0 + 4) * PLANE;
                ah[0][0] = sh[plo];
                ah[0][1] = sh[plo + 8];
                ah[0][2] = sh[phi];
                ah[0][3] = sh[phi + 8];
                al[0][0] = sl[plo];
                al[0][1] = sl[plo + 8];
                al[0][2] = sl[phi];
                al[0][3] = sl[phi + 8];
                ah[1][0] = sh[plo + IN_W];
                ah[1][1] = sh[plo + IN_W + 8];
                ah[1][2] = sh[phi + IN_W];
                ah[1][3] = sh[phi + IN_W + 8];
                al[1][0] = sl[plo + IN_W];
                al[1][1] = sl[plo + IN_W + 8];
                al[1][2] = sl[phi + IN_W];
                al[1][3] = sl[phi + IN_W + 8];
            }

            // B fragments (hi/lo) for the eight n8 tiles (coalesced).
            unsigned bh[8][2], bl[8][2];
            {
                const unsigned* ph = wp_h;
                const unsigned* pl = wp_l;
#pragma unroll
                for (int j = 0; j < 8; ++j) {
                    bh[j][0] = __ldg(ph);
                    bh[j][1] = __ldg(ph + 32);
                    bl[j][0] = __ldg(pl);
                    bl[j][1] = __ldg(pl + 32);
                    ph += 8 * ICP;
                    pl += 8 * ICP;
                }
                wp_h += bstride;
                wp_l += bstride;
            }

            // Term 1: hi * hi
#pragma unroll
            for (int j = 0; j < 8; ++j) {
#pragma unroll
                for (int m = 0; m < 2; ++m) {
                    MMA16816(c[m][j], ah[m], bh[j]);
                }
            }
            // Term 2: lo * hi
#pragma unroll
            for (int j = 0; j < 8; ++j) {
#pragma unroll
                for (int m = 0; m < 2; ++m) {
                    MMA16816(c[m][j], al[m], bh[j]);
                }
            }
            // Term 3: hi * lo
#pragma unroll
            for (int j = 0; j < 8; ++j) {
#pragma unroll
                for (int m = 0; m < 2; ++m) {
                    MMA16816(c[m][j], ah[m], bl[j]);
                }
            }
        }
        __syncthreads();
    }

    // ---- Epilogue: bias + store -------------------------------------------
#pragma unroll
    for (int m = 0; m < 2; ++m) {
        const int oh = oh0 + (m ? r1 : r0);
        if (oh >= HO) continue;
#pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int oc = oc0 + j * 8 + 2 * lk0;
            float c0 = c[m][j][0], c1 = c[m][j][1];
            float c2 = c[m][j][2], c3 = c[m][j][3];
            if (bias != nullptr) {
                if (oc < OC) {
                    const float b0 = __ldg(bias + oc);
                    c0 += b0; c2 += b0;
                }
                if (oc + 1 < OC) {
                    const float b1 = __ldg(bias + oc + 1);
                    c1 += b1; c3 += b1;
                }
            }
            const int ow_a = ow0 + lm0;
            const int ow_b = ow0 + lm0 + 8;
            const bool wa = (ow_a < WO);
            const bool wb = (ow_b < WO);
            float* p0 = out + ((long)(n * OC + oc) * HO + oh) * WO;
            if (oc < OC) {
                if (wa) p0[ow_a] = c0;
                if (wb) p0[ow_b] = c2;
            }
            if (oc + 1 < OC) {
                float* p1 = p0 + (long)HO * WO;
                if (wa) p1[ow_a] = c1;
                if (wb) p1[ow_b] = c3;
            }
        }
    }
#endif
}

// ---------------------------------------------------------------------------
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias_opt,
                             std::vector<int64_t> stride,
                             std::vector<int64_t> padding,
                             std::vector<int64_t> dilation) {
    const bool has_bias =
        bias_opt.has_value() && bias_opt->defined() && bias_opt->numel() > 0;

    bool ok = x.is_cuda() && w.is_cuda()
        && x.scalar_type() == torch::kFloat32
        && x.dim() == 4 && w.dim() == 4
        && stride[0] == 1 && stride[1] == 1
        && dilation[0] == DH && dilation[1] == DW
        && w.size(2) == KH && w.size(3) == KW;
    if (ok) {
        auto prop = at::cuda::getCurrentDeviceProperties();
        ok = (prop->major >= 8);
    }
    if (!ok) {
        return at::conv2d(x, w, bias_opt, stride, padding, dilation, 1);
    }

    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(x));

    auto x_c = x.contiguous();
    auto w_c = w.contiguous();

    const int N  = x_c.size(0);
    const int IC = x_c.size(1);
    const int H  = x_c.size(2);
    const int W  = x_c.size(3);
    const int OC = w_c.size(0);

    const int pad_h = (int)padding[0];
    const int pad_w = (int)padding[1];

    const int HO = H + 2 * pad_h - DH * (KH - 1);   // stride 1
    const int WO = W + 2 * pad_w - DW * (KW - 1);
    if (HO <= 0 || WO <= 0) {
        return at::conv2d(x, w, bias_opt, stride, padding, dilation, 1);
    }

    const int ICP = (IC + 1) / 2;
    const int NJ  = (OC + 63) / 8;            // padded n8 tiles (>= 8)

    auto out   = torch::empty({N, OC, HO, WO}, x_c.options());
    auto uopts = w_c.options().dtype(torch::kInt32);
    const long wsize = (long)NTAP * NJ * ICP * 8;
    auto w_hi  = torch::empty({wsize}, uopts);
    auto w_lo  = torch::empty({wsize}, uopts);

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

    {
        const int total = (int)wsize;
        const int threads = 256;
        const int blocks = (total + threads - 1) / threads;
        wsplit_kernel<<<blocks, threads, 0, stream>>>(
            w_c.data_ptr<float>(),
            reinterpret_cast<unsigned*>(w_hi.data_ptr<int>()),
            reinterpret_cast<unsigned*>(w_lo.data_ptr<int>()),
            IC, OC, ICP, NJ);
    }

    dim3 grid((WO + TILE_OW - 1) / TILE_OW,
              (HO + TILE_OH - 1) / TILE_OH,
              N * ((OC + 63) / 64));

    const float* bias_ptr = nullptr;
    torch::Tensor bias_c;
    if (has_bias) {
        bias_c = bias_opt->contiguous();
        bias_ptr = bias_c.data_ptr<float>();
    }

    conv2d_f16x3_kernel<<<grid, NTHREADS, 0, stream>>>(
        x_c.data_ptr<float>(),
        reinterpret_cast<const unsigned*>(w_hi.data_ptr<int>()),
        reinterpret_cast<const unsigned*>(w_lo.data_ptr<int>()),
        bias_ptr, out.data_ptr<float>(),
        N, IC, H, W, OC, HO, WO, ICP, NJ, pad_h, pad_w);

    return out;
}
"""

conv2d_cpp_source = r"""
#include <torch/extension.h>
#include <vector>
torch::Tensor conv2d_forward(torch::Tensor x, torch::Tensor w,
                             c10::optional<torch::Tensor> bias,
                             std::vector<int64_t> stride,
                             std::vector<int64_t> padding,
                             std::vector<int64_t> dilation);
"""

conv2d_ext = load_inline(
    name="conv2d_dilated_f16x3_v6",
    cpp_sources=[conv2d_cpp_source],
    cuda_sources=[conv2d_cuda_source],
    functions=["conv2d_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
    extra_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.  Keeps an nn.Conv2d submodule so that
    initialization and state_dict are identical to the reference model, but
    computes the forward pass with a custom tensor-core CUDA kernel that
    uses error-compensated fp16 hi/lo splitting (3x mma.m16n8k16 with FP32
    accumulation) to reach far above FP32 CUDA-core throughput while keeping
    FP32-class accuracy.  Falls back to the standard implementation for any
    configuration or device the kernel does not support.
    """

    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)

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

        s = _pair(stride)
        p = _pair(padding)
        d = _pair(dilation)
        self._stride = [s[0], s[1]]
        self._padding = [p[0], p[1]]
        self._dilation = [d[0], d[1]]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 4:
            return conv2d_ext.conv2d_forward(
                x, self.conv2d.weight, self.conv2d.bias,
                self._stride, self._padding, self._dilation)
        return self.conv2d(x)
```