```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 <mma.h>

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 = 508;
constexpr int OW = 504;
constexpr int SPATIAL = OH * OW;
constexpr int TN = 128;
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);
}

// Transpose NCHW to NHWC while computing both components once.
__global__ void pack_input_kernel(
    const float* __restrict__ input,
    half2* __restrict__ packed)
{
    __shared__ float tile[IC][129];

    const int tid = threadIdx.x;
    const int spatial_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 + spatial_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 + (spatial_start + p) * IC + c] =
            split_float(tile[c][p]);
    }
}

// Reduction order is (kernel row, kernel column, input channel).
__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]);
    }
}

union __align__(32) SharedStorage {
    struct {
        half ah[OC * LD];
        half al[OC * LD];
        half bh[TN * LD];
        half bl[TN * LD];
    } operands;
    float result[OC * TN];
};

__global__ __launch_bounds__(THREADS, 2)
void convolution_kernel(
    const half2* __restrict__ input,
    const half2* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    bool has_bias)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
    using namespace nvcuda;

    __shared__ SharedStorage shared;

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int lane = tid % 32;
    const int row_start = (warp / 4) * 32;
    const int col_start = (warp % 4) * 32;
    const int position_start = blockIdx.x * TN;

    const half2* batch_input =
        input + static_cast<int64_t>(blockIdx.y) * IC * IH * IW;

    // Each warp loads all 32 channels of one spatial position.
    int input_offset[TN / 8];
    #pragma unroll
    for (int i = 0; i < TN / 8; ++i) {
        const int p = position_start + warp + i * 8;
        input_offset[i] = (p + (IW - OW) * (p / OW)) * IC + lane;
    }

    wmma::fragment<wmma::accumulator, 16, 16, 16, float>
        main_acc[2][2], correction_acc[2][2];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::fill_fragment(main_acc[i][j], 0.0f);
            wmma::fill_fragment(correction_acc[i][j], 0.0f);
        }
    }

    #pragma unroll 1
    for (int k = 0; k < KH * KW; ++k) {
        const half2* kernel_weight = weight + k * OC * IC;
        const int kernel_offset = ((k / KW) * IW + k % KW) * IC;

        #pragma unroll
        for (int i = 0; i < OC / 8; ++i) {
            const int row = warp + i * 8;
            const half2 value = kernel_weight[row * IC + lane];
            shared.operands.ah[row * LD + lane] = __low2half(value);
            shared.operands.al[row * LD + lane] = __high2half(value);
        }

        #pragma unroll
        for (int i = 0; i < TN / 8; ++i) {
            const int column = warp + i * 8;
            half2 value = __float2half2_rn(0.0f);
            if (position_start + column < SPATIAL) {
                value = batch_input[input_offset[i] + kernel_offset];
            }
            shared.operands.bh[column * LD + lane] = __low2half(value);
            shared.operands.bl[column * LD + lane] = __high2half(value);
        }

        __syncthreads();

        #pragma unroll
        for (int reduction_tile = 0; reduction_tile < IC;
             reduction_tile += 16) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 16, half, wmma::row_major>
                ah[2], al[2];

            wmma::fragment<
                wmma::matrix_b, 16, 16, 16, half, wmma::col_major>
                bh[2], bl[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset =
                    (row_start + i * 16) * LD + reduction_tile;
                wmma::load_matrix_sync(
                    ah[i], shared.operands.ah + offset, LD);
                wmma::load_matrix_sync(
                    al[i], shared.operands.al + offset, LD);
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int offset =
                    (col_start + j * 16) * LD + reduction_tile;
                wmma::load_matrix_sync(
                    bh[j], shared.operands.bh + offset, LD);
                wmma::load_matrix_sync(
                    bl[j], shared.operands.bl + offset, LD);
            }

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

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int e = 0; e < main_acc[i][j].num_elements; ++e) {
                main_acc[i][j].x[e] = fmaf(
                    correction_acc[i][j].x[e],
                    1.0f / 4096.0f,
                    main_acc[i][j].x[e]);
            }
            wmma::store_matrix_sync(
                shared.result
                    + (row_start + i * 16) * TN + col_start + j * 16,
                main_acc[i][j], TN, wmma::mem_row_major);
        }
    }

    __syncthreads();

    #pragma unroll
    for (int i = tid; i < OC * TN; i += THREADS) {
        const int channel = i / TN;
        const int position = position_start + i % TN;
        if (position < SPATIAL) {
            float value = shared.result[i];
            if (has_bias) {
                value += bias[channel];
            }
            output[
                (static_cast<int64_t>(blockIdx.y) * OC + channel)
                * SPATIAL + position] = value;
        }
    }
#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((SPATIAL + TN - 1) / TN,
             static_cast<unsigned int>(x.size(0))),
        THREADS, 0, stream.stream()>>>(
            input_ptr, 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_packed_wmma_v5",
            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] < 7:
            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,
            )
```