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


_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor implicit_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t stride_h, int64_t stride_w,
    int64_t pad_h, int64_t pad_w,
    int64_t dilation_h, int64_t dilation_w);
"""


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

namespace {

__device__ __forceinline__ float round_tf32(float value) {
    unsigned bits = __float_as_uint(value);
    if ((bits & 0x7f800000u) == 0x7f800000u)
        return value;
    bits = (bits + 0xfffu + ((bits >> 13) & 1u)) & 0xffffe000u;
    return __uint_as_float(bits);
}

template <bool Specialized>
__global__ void implicit_conv_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int channels, int height, int width,
    int output_channels, int output_h, int output_w,
    int kernel_h_arg, int kernel_w_arg,
    int stride_h_arg, int stride_w_arg,
    int pad_h_arg, int pad_w_arg,
    int dilation_h_arg, int dilation_w_arg)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    using namespace nvcuda;

    constexpr int BM = 64;
    constexpr int BN = 64;
    constexpr int BK = 32;

    const int kh = Specialized ? 5 : kernel_h_arg;
    const int kw = Specialized ? 9 : kernel_w_arg;
    const int sh = Specialized ? 1 : stride_h_arg;
    const int sw = Specialized ? 1 : stride_w_arg;
    const int ph = Specialized ? 2 : pad_h_arg;
    const int pw = Specialized ? 4 : pad_w_arg;
    const int dh = Specialized ? 2 : dilation_h_arg;
    const int dw = Specialized ? 3 : dilation_w_arg;

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int warp_row = (warp / 2) * 16;
    const int warp_col = (warp % 2) * 32;

    const int spatial_size = output_h * output_w;
    const int reduction_size = channels * kh * kw;
    const int channel_base = blockIdx.y * BM;
    const int spatial_base = blockIdx.x * BN;
    const int batch = blockIdx.z;

    __shared__ __align__(32) float shared[8192];
    float* a_hi = shared;
    float* a_lo = shared + 2048;
    float* b_hi = shared + 4096;
    float* b_lo = shared + 6144;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2];
    wmma::fill_fragment(acc[0], 0.0f);
    wmma::fill_fragment(acc[1], 0.0f);

    for (int reduction_base = 0;
         reduction_base < reduction_size;
         reduction_base += BK) {

        for (int index = tid; index < BM * BK; index += 256) {
            const int row = index / BK;
            const int k = reduction_base + index % BK;
            const int oc = channel_base + row;

            float value = 0.0f;
            if (oc < output_channels && k < reduction_size) {
                value = weight[
                    static_cast<int64_t>(oc) * reduction_size + k];
            }

            const float high = round_tf32(value);
            a_hi[index] = high;
            a_lo[index] = round_tf32(value - high);
        }

        for (int index = tid; index < BK * BN; index += 256) {
            const int k = reduction_base + index / BN;
            const int position = spatial_base + index % BN;

            float value = 0.0f;
            if (k < reduction_size && position < spatial_size) {
                const int ic = k / (kh * kw);
                const int remainder = k % (kh * kw);
                const int ky = remainder / kw;
                const int kx = remainder % kw;

                const int oy = position / output_w;
                const int ox = position % output_w;
                const int iy = oy * sh - ph + ky * dh;
                const int ix = ox * sw - pw + kx * dw;

                if (static_cast<unsigned>(iy) <
                        static_cast<unsigned>(height) &&
                    static_cast<unsigned>(ix) <
                        static_cast<unsigned>(width)) {
                    const int64_t offset =
                        ((static_cast<int64_t>(batch) * channels + ic)
                         * height + iy) * width + ix;
                    value = x[offset];
                }
            }

            const float high = round_tf32(value);
            b_hi[index] = high;
            b_lo[index] = round_tf32(value - high);
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; k += 8) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major> ah, al;

            wmma::load_matrix_sync(
                ah, a_hi + warp_row * BK + k, BK);
            wmma::load_matrix_sync(
                al, a_lo + warp_row * BK + k, BK);

            #pragma unroll
            for (int tile = 0; tile < 2; ++tile) {
                wmma::fragment<
                    wmma::matrix_b, 16, 16, 8,
                    wmma::precision::tf32, wmma::row_major> bh, bl;

                const int offset = k * BN + warp_col + tile * 16;
                wmma::load_matrix_sync(bh, b_hi + offset, BN);
                wmma::load_matrix_sync(bl, b_lo + offset, BN);

                // Recover FP32-like accuracy with TF32 residual products.
                wmma::mma_sync(acc[tile], al, bh, acc[tile]);
                wmma::mma_sync(acc[tile], ah, bl, acc[tile]);
                wmma::mma_sync(acc[tile], ah, bh, acc[tile]);
            }
        }

        __syncthreads();
    }

    // Reuse the operand staging buffer for the output tile.
    #pragma unroll
    for (int tile = 0; tile < 2; ++tile) {
        wmma::store_matrix_sync(
            shared + warp_row * BN + warp_col + tile * 16,
            acc[tile], BN, wmma::mem_row_major);
    }

    __syncthreads();

    for (int index = tid; index < BM * BN; index += 256) {
        const int oc = channel_base + index / BN;
        const int position = spatial_base + index % BN;
        if (oc < output_channels && position < spatial_size) {
            float value = shared[index];
            if (bias != nullptr)
                value += bias[oc];

            output[
                (static_cast<int64_t>(batch) * output_channels + oc)
                * spatial_size + position] = value;
        }
    }
#endif
}

} // namespace

torch::Tensor implicit_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    int64_t stride_h, int64_t stride_w,
    int64_t pad_h, int64_t pad_w,
    int64_t dilation_h, int64_t dilation_w)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(),
                "Input and weight must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat &&
                weight.scalar_type() == at::kFloat,
                "Input and weight must be float32");
    TORCH_CHECK(x.dim() == 4 && weight.dim() == 4,
                "Expected four-dimensional tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(),
                "Expected contiguous NCHW tensors");
    TORCH_CHECK(x.device() == weight.device(),
                "Input and weight must be on the same device");
    TORCH_CHECK(x.size(1) == weight.size(1),
                "Input channel count does not match weight");

    const c10::cuda::CUDAGuard device_guard(x.device());
    TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8,
                "This kernel requires compute capability 8.0 or newer");

    const int batch = x.size(0);
    const int channels = x.size(1);
    const int height = x.size(2);
    const int width = x.size(3);
    const int output_channels = weight.size(0);
    const int kernel_h = weight.size(2);
    const int kernel_w = weight.size(3);

    const int64_t effective_h = dilation_h * (kernel_h - 1) + 1;
    const int64_t effective_w = dilation_w * (kernel_w - 1) + 1;
    TORCH_CHECK(height + 2 * pad_h >= effective_h &&
                width + 2 * pad_w >= effective_w,
                "Kernel exceeds padded input dimensions");

    const int output_h =
        (height + 2 * pad_h - effective_h) / stride_h + 1;
    const int output_w =
        (width + 2 * pad_w - effective_w) / stride_w + 1;

    auto output = torch::empty(
        {batch, output_channels, output_h, output_w}, x.options());

    if (output.numel() == 0)
        return output;

    const float* bias_ptr = nullptr;
    if (bias.numel() != 0) {
        TORCH_CHECK(
            bias.is_cuda() && bias.device() == x.device() &&
            bias.scalar_type() == at::kFloat && bias.is_contiguous() &&
            bias.numel() == output_channels,
            "Invalid bias tensor");
        bias_ptr = bias.data_ptr<float>();
    }

    const dim3 block(256);
    const dim3 grid(
        (output_h * output_w + 63) / 64,
        (output_channels + 63) / 64,
        batch);
    const cudaStream_t stream =
        at::cuda::getCurrentCUDAStream(x.get_device());

    const bool specialized =
        kernel_h == 5 && kernel_w == 9 &&
        stride_h == 1 && stride_w == 1 &&
        pad_h == 2 && pad_w == 4 &&
        dilation_h == 2 && dilation_w == 3;

    if (specialized) {
        implicit_conv_kernel<true><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            bias_ptr, output.data_ptr<float>(),
            channels, height, width, output_channels, output_h, output_w,
            kernel_h, kernel_w, stride_h, stride_w,
            pad_h, pad_w, dilation_h, dilation_w);
    } else {
        implicit_conv_kernel<false><<<grid, block, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(),
            bias_ptr, output.data_ptr<float>(),
            channels, height, width, output_channels, output_h, output_w,
            kernel_h, kernel_w, stride_h, stride_w,
            pad_h, pad_w, dilation_h, dilation_w);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""


_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="asymmetric_dilated_implicit_conv_tf32x3",
            cpp_sources=_cpp_source,
            cuda_sources=_cuda_source,
            functions=["implicit_conv_cuda"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            with_cuda=True,
            verbose=False,
        )
    return _extension


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().__init__()
        self.conv2d = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=bias,
        )
        self.register_buffer(
            "_empty_bias", torch.empty(0), persistent=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv2d

        needs_autograd = torch.is_grad_enabled() and (
            x.requires_grad
            or conv.weight.requires_grad
            or (conv.bias is not None and conv.bias.requires_grad)
        )

        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or conv.weight.dtype != torch.float32
            or x.ndim != 4
            or needs_autograd
            or torch.is_autocast_enabled()
            or isinstance(conv.padding, str)
            or torch.cuda.get_device_capability(x.device)[0] < 8
            or x.shape[0] > 65535
        ):
            return conv(x)

        bias = conv.bias if conv.bias is not None else self._empty_bias

        return _get_extension().implicit_conv_cuda(
            x.contiguous(),
            conv.weight.contiguous(),
            bias,
            conv.stride[0],
            conv.stride[1],
            conv.padding[0],
            conv.padding[1],
            conv.dilation[0],
            conv.dilation[1],
        )
```