```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 conv3d_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias);
"""

_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_fp16.h>
#include <cuda_runtime.h>

constexpr int SPATIAL = 62 * 62 * 62;
constexpr int INPUT_SPATIAL = 64 * 64 * 64;
constexpr int TILE_N = 256;
constexpr int A_STRIDE = 40;
constexpr int B_STRIDE = 264;
constexpr int SHARED_HALVES =
    2 * 64 * A_STRIDE + 2 * 32 * B_STRIDE;

__device__ __forceinline__ unsigned int pack_float(float value) {
    const half hi = __float2half_rn(value);
    const half lo = __float2half_rn(value - __half2float(hi));
    return static_cast<unsigned int>(__half_as_ushort(hi)) |
           (static_cast<unsigned int>(__half_as_ushort(lo)) << 16);
}

__global__ void pack_operands_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    unsigned int* __restrict__ packed,
    int64_t input_vectors,
    int64_t total_vectors)
{
    const int64_t index =
        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (index >= total_vectors)
        return;

    const float4 values = index < input_vectors
        ? reinterpret_cast<const float4*>(x)[index]
        : reinterpret_cast<const float4*>(weight)[index - input_vectors];

    reinterpret_cast<uint4*>(packed)[index] = make_uint4(
        pack_float(values.x), pack_float(values.y),
        pack_float(values.z), pack_float(values.w));
}

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

__device__ __forceinline__ void load_a(
    unsigned int (&fragment)[4], const half* pointer)
{
    const unsigned int address =
        static_cast<unsigned int>(__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(
    unsigned int (&fragment)[4], const half* pointer)
{
    const unsigned int address =
        static_cast<unsigned int>(__cvta_generic_to_shared(pointer));
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.trans.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 mma(
    float (&accumulator)[4],
    const unsigned int (&a)[4],
    const unsigned int* 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"(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]));
}

__device__ __forceinline__ void store_half_pair(
    half* hi, half* lo, unsigned int first, unsigned int second)
{
    *reinterpret_cast<unsigned int*>(hi) =
        __byte_perm(first, second, 0x5410);
    *reinterpret_cast<unsigned int*>(lo) =
        __byte_perm(first, second, 0x7632);
}

__device__ __forceinline__ void store_streaming_pair(
    float* destination, float first, float second)
{
    asm volatile(
        "st.global.cs.v2.f32 [%0], {%1, %2};\n"
        :
        : "l"(destination), "f"(first), "f"(second)
        : "memory");
}

#endif

__global__ __launch_bounds__(256, 2)
void conv3d_kernel(
    const unsigned int* __restrict__ x,
    const unsigned int* __restrict__ weight,
    const float* __restrict__ original_x,
    const float* __restrict__ original_weight,
    const float* __restrict__ bias,
    float* __restrict__ output)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    __shared__ __align__(32) half storage[SHARED_HALVES];

    half* a_hi = storage;
    half* a_lo = a_hi + 64 * A_STRIDE;
    half* b_hi = a_lo + 64 * A_STRIDE;
    half* b_lo = b_hi + 32 * B_STRIDE;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int channel_base = (warp & 1) * 32;
    const int column_base = (warp >> 1) * 64;
    const int spatial_base = blockIdx.x * TILE_N;
    const int batch = blockIdx.y;

    float acc[2][8][4];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            #pragma unroll
            for (int r = 0; r < 4; ++r)
                acc[i][j][r] = 0.0f;
        }
    }

    const int load_position = spatial_base + (tid & 127) * 2;
    const int ox = load_position % 62;
    const int oy = (load_position / 62) % 62;
    const int oz = load_position / (62 * 62);
    const int64_t input_base =
        static_cast<int64_t>(batch) * (3 * INPUT_SPATIAL) +
        oz * (64 * 64) + oy * 64 + ox;

    #pragma unroll
    for (int k_base = 0; k_base < 80; k_base += 32) {
        #pragma unroll
        for (int t = 0; t < 4; ++t) {
            const int index = tid + t * 256;
            const int row = index >> 4;
            const int col = (index & 15) * 2;
            const int k = k_base + col;

            if (k < 80) {
                store_half_pair(
                    a_hi + row * A_STRIDE + col,
                    a_lo + row * A_STRIDE + col,
                    weight[row * 81 + k],
                    weight[row * 81 + k + 1]);
            }
        }

        #pragma unroll
        for (int t = 0; t < 16; ++t) {
            const int row = (tid >> 7) + t * 2;
            const int col = (tid & 127) * 2;
            const int k = k_base + row;

            if (k < 80) {
                unsigned int first = 0u;
                unsigned int second = 0u;
                if (load_position < SPATIAL) {
                    const int ic = k / 27;
                    const int kz = (k / 9) % 3;
                    const int ky = (k / 3) % 3;
                    const int kx = k % 3;
                    const int64_t offset =
                        input_base + ic * INPUT_SPATIAL +
                        kz * (64 * 64) + ky * 64 + kx;
                    first = x[offset];
                    second = x[offset + 1];
                }

                store_half_pair(
                    b_hi + row * B_STRIDE + col,
                    b_lo + row * B_STRIDE + col,
                    first, second);
            }
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < 32; kk += 16) {
            if (k_base + kk < 80) {
                unsigned int ah[2][4];
                unsigned int al[2][4];

                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    const int offset =
                        (channel_base + i * 16 + (lane & 15)) *
                        A_STRIDE + kk + (lane >> 4) * 8;
                    load_a(ah[i], a_hi + offset);
                    load_a(al[i], a_lo + offset);
                }

                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    unsigned int bh[4];
                    unsigned int bl[4];
                    const int offset =
                        (kk + (lane & 15)) * B_STRIDE +
                        column_base + j * 16 + (lane >> 4) * 8;

                    load_b(bh, b_hi + offset);
                    load_b(bl, b_lo + offset);

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

        if (k_base < 64)
            __syncthreads();
    }

    float w0[2], w1[2], bias0[2], bias1[2];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int channel = channel_base + i * 16 + (lane >> 2);
        w0[i] = original_weight[channel * 81 + 80];
        w1[i] = original_weight[(channel + 8) * 81 + 80];
        bias0[i] = bias != nullptr ? bias[channel] : 0.0f;
        bias1[i] = bias != nullptr ? bias[channel + 8] : 0.0f;
    }

    const int first_position =
        spatial_base + column_base + (lane & 3) * 2;
    const int px = first_position % 62;
    const int py = (first_position / 62) % 62;
    const int pz = first_position / (62 * 62);
    const int64_t tail_base =
        static_cast<int64_t>(batch) * (3 * INPUT_SPATIAL) +
        2 * INPUT_SPATIAL +
        (pz + 2) * (64 * 64) + (py + 2) * 64 + px + 2;
    const int row_correction = 2 + (py == 61 ? 128 : 0);

    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        const int position = first_position + j * 8;
        if (position < SPATIAL) {
            const int64_t offset =
                tail_base + j * 8 +
                (px + j * 8 >= 62 ? row_correction : 0);
            const float2 tail =
                *reinterpret_cast<const float2*>(original_x + offset);

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int channel =
                    channel_base + i * 16 + (lane >> 2);
                float* destination =
                    output +
                    (static_cast<int64_t>(batch) * 64 + channel) *
                    SPATIAL + position;

                store_streaming_pair(
                    destination,
                    fmaf(w0[i], tail.x, acc[i][j][0]) + bias0[i],
                    fmaf(w0[i], tail.y, acc[i][j][1]) + bias0[i]);

                store_streaming_pair(
                    destination + 8 * SPATIAL,
                    fmaf(w1[i], tail.x, acc[i][j][2]) + bias1[i],
                    fmaf(w1[i], tail.y, acc[i][j][3]) + bias1[i]);
            }
        }
    }
#endif
}

torch::Tensor conv3d_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "Expected CUDA tensors");
    const c10::cuda::CUDAGuard guard(x.device());

    TORCH_CHECK(
        x.scalar_type() == at::kFloat &&
        weight.scalar_type() == at::kFloat,
        "Expected float32 tensors");
    TORCH_CHECK(
        x.is_contiguous() && weight.is_contiguous(),
        "Expected contiguous tensors");
    TORCH_CHECK(x.device() == weight.device(), "Device mismatch");
    TORCH_CHECK(
        x.dim() == 5 && x.size(1) == 3 &&
        x.size(2) == 64 && x.size(3) == 64 && x.size(4) == 64 &&
        x.size(0) <= 65535,
        "Unsupported input shape");
    TORCH_CHECK(
        weight.dim() == 5 &&
        weight.size(0) == 64 && weight.size(1) == 3 &&
        weight.size(2) == 3 && weight.size(3) == 3 &&
        weight.size(4) == 3,
        "Unsupported weight shape");
    TORCH_CHECK(
        at::cuda::getCurrentDeviceProperties()->major >= 8,
        "Requires compute capability 8.0 or newer");

    auto output = torch::empty({x.size(0), 64, 62, 62, 62}, x.options());
    if (x.size(0) == 0)
        return output;

    const float* bias_ptr = nullptr;
    if (bias.has_value()) {
        TORCH_CHECK(
            bias->device() == x.device() &&
            bias->scalar_type() == at::kFloat &&
            bias->is_contiguous() && bias->numel() == 64,
            "Invalid bias");
        bias_ptr = bias->data_ptr<float>();
    }

    const int64_t input_elements = x.numel();
    const int64_t total_elements = input_elements + weight.numel();
    auto packed = torch::empty(
        {total_elements}, x.options().dtype(at::kInt));
    auto* packed_ptr =
        reinterpret_cast<unsigned int*>(packed.data_ptr<int>());

    const auto stream = at::cuda::getCurrentCUDAStream();
    const int64_t total_vectors = total_elements / 4;

    pack_operands_kernel<<<(total_vectors + 255) / 256, 256, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(),
        packed_ptr, input_elements / 4, total_vectors);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const dim3 grid(
        (SPATIAL + TILE_N - 1) / TILE_N,
        static_cast<unsigned int>(x.size(0)));

    conv3d_kernel<<<grid, 256, 0, stream>>>(
        packed_ptr, packed_ptr + input_elements,
        x.data_ptr<float>(), weight.data_ptr<float>(),
        bias_ptr, output.data_ptr<float>());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="conv3d_compensated_mma_n256_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["conv3d_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _Conv3dFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias):
        ctx.save_for_backward(x, weight)
        ctx.has_bias = bias is not None
        return _extension.conv3d_cuda(x, weight, bias)

    @staticmethod
    def backward(ctx, grad_output):
        x, weight = ctx.saved_tensors
        grad_output = grad_output.contiguous()
        grad_x = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_x = torch.nn.grad.conv3d_input(
                x.shape, weight, grad_output
            )
        if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv3d_weight(
                x, weight.shape, grad_output
            )
        if ctx.has_bias and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(dim=(0, 2, 3, 4))

        return grad_x, grad_weight, grad_bias


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv3d
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            and x.dim() == 5
            and tuple(x.shape[1:]) == (3, 64, 64, 64)
            and x.shape[0] <= 65535
            and x.data_ptr() % 16 == 0
            and conv.weight.dtype == torch.float32
            and conv.weight.device == x.device
            and conv.weight.is_contiguous()
            and conv.weight.data_ptr() % 16 == 0
            and tuple(conv.weight.shape) == (64, 3, 3, 3, 3)
            and conv.stride == (1, 1, 1)
            and conv.padding == (0, 0, 0)
            and conv.dilation == (1, 1, 1)
            and conv.groups == 1
            and (
                conv.bias is None
                or (
                    conv.bias.dtype == torch.float32
                    and conv.bias.device == x.device
                    and conv.bias.is_contiguous()
                )
            )
            and not torch.is_autocast_enabled()
            and torch.cuda.get_device_capability(x.device)[0] >= 8
        ):
            return _Conv3dFunction.apply(x, conv.weight, conv.bias)
        return conv(x)
```