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


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

torch::Tensor asymmetric_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    bool has_bias);
"""

_CUDA = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cstdint>

namespace {

constexpr int IC = 32;
constexpr int OC = 64;
constexpr int IH = 512;
constexpr int IW = 512;
constexpr int KH = 5;
constexpr int KW = 9;
constexpr int OH = IH - KH + 1;
constexpr int OW = IW - KW + 1;
constexpr int SPATIAL = OH * OW;

constexpr int TN = 128;
constexpr int PATCH = TN + KW - 1;
constexpr int LD = 40;
constexpr int THREADS = 256;

__device__ __forceinline__ half2 split_float(float value) {
    const half high = __float2half_rn(value);
    const half low = __float2half_rn(
        (value - __half2float(high)) * 4096.0f);
    return __halves2half2(high, low);
}

__global__ void pack_input_kernel(
    const float* __restrict__ input,
    half2* __restrict__ packed)
{
    __shared__ float tile[IC][129];

    const int tid = threadIdx.x;
    const int start = blockIdx.x * 128;
    const int64_t batch_offset =
        static_cast<int64_t>(blockIdx.y) * IC * IH * IW;

    #pragma unroll
    for (int i = tid; i < IC * 128; i += THREADS) {
        const int c = i / 128;
        const int p = i % 128;
        tile[c][p] = input[
            batch_offset + c * IH * IW + start + p];
    }

    __syncthreads();

    #pragma unroll
    for (int i = tid; i < IC * 128; i += THREADS) {
        const int p = i / IC;
        const int c = i % IC;
        packed[batch_offset + (start + p) * IC + c] =
            split_float(tile[c][p]);
    }
}

__global__ void pack_weight_kernel(
    const float* __restrict__ weight,
    half2* __restrict__ packed)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < KH * KW * OC * IC) {
        const int c = i % IC;
        const int o = (i / IC) % OC;
        const int k = i / (IC * OC);
        packed[i] = split_float(
            weight[(o * IC + c) * KH * KW + k]);
    }
}

struct __align__(32) SharedStorage {
    half ah[OC * LD];
    half al[OC * LD];
    half bh[PATCH * LD];
    half bl[PATCH * LD];
};

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800

__device__ __forceinline__ void load_a(
    uint32_t (&fragment)[4],
    const half* pointer)
{
    const uint32_t address =
        static_cast<uint32_t>(__cvta_generic_to_shared(pointer));

    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 "
        "{%0, %1, %2, %3}, [%4];\n"
        : "=r"(fragment[0]), "=r"(fragment[1]),
          "=r"(fragment[2]), "=r"(fragment[3])
        : "r"(address));
}

__device__ __forceinline__ void load_b(
    uint32_t (&fragment)[2],
    const half* pointer)
{
    const uint32_t address =
        static_cast<uint32_t>(__cvta_generic_to_shared(pointer));

    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x2.shared.b16 "
        "{%0, %1}, [%2];\n"
        : "=r"(fragment[0]), "=r"(fragment[1])
        : "r"(address));
}

__device__ __forceinline__ void mma(
    float (&accumulator)[4],
    const uint32_t (&a)[4],
    const uint32_t (&b)[2])
{
    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"(accumulator[0]), "+f"(accumulator[1]),
          "+f"(accumulator[2]), "+f"(accumulator[3])
        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
          "r"(b[0]), "r"(b[1]));
}

#endif

__global__ __launch_bounds__(THREADS, 2)
void convolution_kernel(
    const uint2* __restrict__ input,
    const uint2* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    bool has_bias)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    __shared__ SharedStorage shared;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;

    const int row_start = (warp >> 2) * 32;
    const int col_start = (warp & 3) * 32;

    const int output_x = blockIdx.x * TN;
    const int output_y = blockIdx.y;
    const int batch = blockIdx.z;

    float main_acc[2][4][4];
    float correction_acc[2][4][4];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            #pragma unroll
            for (int e = 0; e < 4; ++e) {
                main_acc[i][j][e] = 0.0f;
                correction_acc[i][j][e] = 0.0f;
            }
        }
    }

    uint32_t* shared_ah = reinterpret_cast<uint32_t*>(shared.ah);
    uint32_t* shared_al = reinterpret_cast<uint32_t*>(shared.al);
    uint32_t* shared_bh = reinterpret_cast<uint32_t*>(shared.bh);
    uint32_t* shared_bl = reinterpret_cast<uint32_t*>(shared.bl);

    #pragma unroll 1
    for (int ky = 0; ky < KH; ++ky) {
        const uint2* input_row = input +
            ((static_cast<int64_t>(batch) * IH + output_y + ky) * IW
             + output_x) * (IC / 2);

        // Stage an entire horizontal patch once, reusing it for all
        // nine kernel columns. Each load handles two input channels.
        #pragma unroll
        for (int index = tid; index < PATCH * (IC / 2);
             index += THREADS) {
            const int column = index / (IC / 2);
            const int channel_pair = index % (IC / 2);

            uint2 value = make_uint2(0u, 0u);
            if (output_x + column < IW) {
                value = input_row[index];
            }

            const int destination = column * (LD / 2) + channel_pair;
            shared_bh[destination] =
                (value.x & 0xffffu) | (value.y << 16);
            shared_bl[destination] =
                (value.x >> 16) | (value.y & 0xffff0000u);
        }

        #pragma unroll 1
        for (int kx = 0; kx < KW; ++kx) {
            const uint2* kernel_weight =
                weight + (ky * KW + kx) * OC * (IC / 2);

            #pragma unroll
            for (int index = tid; index < OC * (IC / 2);
                 index += THREADS) {
                const int row = index / (IC / 2);
                const int channel_pair = index % (IC / 2);
                const uint2 value = kernel_weight[index];

                const int destination = row * (LD / 2) + channel_pair;
                shared_ah[destination] =
                    (value.x & 0xffffu) | (value.y << 16);
                shared_al[destination] =
                    (value.x >> 16) | (value.y & 0xffff0000u);
            }

            __syncthreads();

            #pragma unroll
            for (int reduction = 0; reduction < IC; reduction += 16) {
                uint32_t ah[2][4];
                uint32_t al[2][4];
                uint32_t bh[4][2];
                uint32_t bl[4][2];

                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    const int offset =
                        (row_start + i * 16 + (lane & 15)) * LD
                        + reduction + (lane >> 4) * 8;
                    load_a(ah[i], shared.ah + offset);
                    load_a(al[i], shared.al + offset);
                }

                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    const int offset =
                        (col_start + j * 8 + (lane & 7) + kx) * LD
                        + reduction + ((lane >> 3) & 1) * 8;
                    load_b(bh[j], shared.bh + offset);
                    load_b(bl[j], shared.bl + offset);
                }

                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    #pragma unroll
                    for (int j = 0; j < 4; ++j) {
                        mma(main_acc[i][j], ah[i], bh[j]);
                        mma(correction_acc[i][j], ah[i], bl[j]);
                        mma(correction_acc[i][j], al[i], bh[j]);
                    }
                }
            }

            __syncthreads();
        }
    }

    // MMA accumulator lanes already hold adjacent pairs of output
    // positions, allowing direct vectorized NCHW stores.
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int channel0 = row_start + i * 16 + (lane >> 2);
        const int channel1 = channel0 + 8;

        const float bias0 = has_bias ? bias[channel0] : 0.0f;
        const float bias1 = has_bias ? bias[channel1] : 0.0f;

        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int column =
                output_x + col_start + j * 8 + (lane & 3) * 2;

            if (column < OW) {
                float2 value0;
                float2 value1;

                value0.x = fmaf(
                    correction_acc[i][j][0], 1.0f / 4096.0f,
                    main_acc[i][j][0]) + bias0;
                value0.y = fmaf(
                    correction_acc[i][j][1], 1.0f / 4096.0f,
                    main_acc[i][j][1]) + bias0;
                value1.x = fmaf(
                    correction_acc[i][j][2], 1.0f / 4096.0f,
                    main_acc[i][j][2]) + bias1;
                value1.y = fmaf(
                    correction_acc[i][j][3], 1.0f / 4096.0f,
                    main_acc[i][j][3]) + bias1;

                const int64_t offset0 =
                    (static_cast<int64_t>(batch) * OC + channel0)
                    * SPATIAL + output_y * OW + column;
                const int64_t offset1 = offset0 + 8 * SPATIAL;

                *reinterpret_cast<float2*>(output + offset0) = value0;
                *reinterpret_cast<float2*>(output + offset1) = value1;
            }
        }
    }
#endif
}

} // namespace

torch::Tensor asymmetric_conv_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    bool has_bias)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(),
                "Expected CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                weight.scalar_type() == torch::kFloat32,
                "Expected float32 tensors");
    TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(),
                "Expected contiguous tensors");
    TORCH_CHECK(x.device() == weight.device(),
                "Input and weight must share a device");
    TORCH_CHECK(x.dim() == 4 && x.size(1) == IC &&
                x.size(2) == IH && x.size(3) == IW &&
                x.size(0) <= 65535,
                "Unsupported input shape");
    TORCH_CHECK(weight.dim() == 4 && weight.size(0) == OC &&
                weight.size(1) == IC && weight.size(2) == KH &&
                weight.size(3) == KW,
                "Unsupported weight shape");

    if (has_bias) {
        TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() &&
                    bias.scalar_type() == torch::kFloat32 &&
                    bias.is_contiguous() && bias.numel() == OC,
                    "Invalid bias");
    }

    const c10::cuda::CUDAGuard guard(x.device());

    auto output = torch::empty({x.size(0), OC, OH, OW}, x.options());
    if (x.size(0) == 0) {
        return output;
    }

    auto packed_input = torch::empty(
        {x.size(0), IH, IW, IC}, x.options().dtype(torch::kInt32));
    auto packed_weight = torch::empty(
        {KH * KW, OC, IC}, x.options().dtype(torch::kInt32));

    const auto stream = c10::cuda::getCurrentCUDAStream(x.get_device());

    auto* input_ptr =
        reinterpret_cast<half2*>(packed_input.data_ptr<int32_t>());
    auto* weight_ptr =
        reinterpret_cast<half2*>(packed_weight.data_ptr<int32_t>());

    pack_input_kernel<<<
        dim3(IH * IW / 128, static_cast<unsigned int>(x.size(0))),
        THREADS, 0, stream.stream()>>>(
            x.data_ptr<float>(), input_ptr);

    pack_weight_kernel<<<
        (KH * KW * OC * IC + THREADS - 1) / THREADS,
        THREADS, 0, stream.stream()>>>(
            weight.data_ptr<float>(), weight_ptr);

    convolution_kernel<<<
        dim3((OW + TN - 1) / TN, OH,
             static_cast<unsigned int>(x.size(0))),
        THREADS, 0, stream.stream()>>>(
            reinterpret_cast<const uint2*>(input_ptr),
            reinterpret_cast<const uint2*>(weight_ptr),
            has_bias ? bias.data_ptr<float>() : nullptr,
            output.data_ptr<float>(),
            has_bias);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="asymmetric_conv_patch_mma_v6",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=["asymmetric_conv_cuda"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            verbose=False,
        )
    return _extension


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple,
        stride: int = 1,
        padding: int = 0,
        dilation: int = 1,
        groups: int = 1,
        bias: bool = False,
    ):
        super().__init__()
        self.conv2d = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )

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

        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and weight.dtype == torch.float32
            and weight.device == x.device
            and x.ndim == 4
            and tuple(x.shape[1:]) == (32, 512, 512)
            and x.shape[0] <= 65535
            and conv.in_channels == 32
            and conv.out_channels == 64
            and conv.kernel_size == (5, 9)
            and conv.stride == (1, 1)
            and conv.padding == (0, 0)
            and conv.dilation == (1, 1)
            and conv.groups == 1
            and not torch.is_autocast_enabled()
        )

        if not supported:
            return conv(x)

        if torch.cuda.get_device_capability(x.device)[0] < 8:
            return conv(x)

        bias = conv.bias
        if bias is not None and (
            bias.dtype != torch.float32 or bias.device != x.device
        ):
            return conv(x)

        with torch.cuda.device(x.device):
            return _get_extension().asymmetric_conv_cuda(
                x.contiguous(),
                weight.contiguous(),
                bias.contiguous() if bias is not None else weight,
                bias is not None,
            )
```