```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_mma_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 = 128;
constexpr int A_STRIDE = 40;
constexpr int B_STRIDE = 136;
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)[2],
    const half* pointer)
{
    const unsigned int address =
        static_cast<unsigned int>(__cvta_generic_to_shared(pointer));

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

__device__ __forceinline__ void mma(
    float (&accumulator)[4],
    const unsigned int (&a)[4],
    const unsigned int (&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__(256, 2)
void conv3d_mma_kernel(
    const unsigned int* __restrict__ x,
    const unsigned int* __restrict__ 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;

    // Each warp computes 32 output channels by 32 spatial positions.
    const int channel_base = (warp & 1) * 32;
    const int column_base = (warp >> 1) * 32;
    const int spatial_base = blockIdx.x * TILE_N;
    const int batch = blockIdx.y;

    float 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 r = 0; r < 4; ++r)
                acc[i][j][r] = 0.0f;
        }
    }

    const int load_position = spatial_base + (tid & 127);
    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 < 81; k_base += 32) {
        #pragma unroll
        for (int t = 0; t < 8; ++t) {
            const int index = tid + t * 256;
            const int row = index >> 5;
            const int col = index & 31;
            const int k = k_base + col;

            const unsigned int bits =
                k < 81 ? weight[row * 81 + k] : 0u;

            a_hi[row * A_STRIDE + col] =
                __ushort_as_half(static_cast<unsigned short>(bits));
            a_lo[row * A_STRIDE + col] =
                __ushort_as_half(static_cast<unsigned short>(bits >> 16));
        }

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

            unsigned int bits = 0u;
            if (k < 81 && 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;

                bits = x[
                    input_base + ic * INPUT_SPATIAL +
                    kz * (64 * 64) + ky * 64 + kx];
            }

            b_hi[row * B_STRIDE + col] =
                __ushort_as_half(static_cast<unsigned short>(bits));
            b_lo[row * B_STRIDE + col] =
                __ushort_as_half(static_cast<unsigned short>(bits >> 16));
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < 32; kk += 16) {
            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[2];
                unsigned int bl[2];

                const int offset =
                    (kk + (lane & 15)) * B_STRIDE +
                    column_base + j * 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], ah[i], bh);
                    mma(acc[i][j], ah[i], bl);
                    mma(acc[i][j], al[i], bh);
                }
            }
        }

        __syncthreads();
    }

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

        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int position =
                spatial_base + column_base + j * 8 + (lane & 3) * 2;

            if (position < SPATIAL) {
                float* destination =
                    output +
                    (static_cast<int64_t>(batch) * 64 + channel) *
                    SPATIAL + position;

                *reinterpret_cast<float2*>(destination) = make_float2(
                    acc[i][j][0] + bias0,
                    acc[i][j][1] + bias0);

                *reinterpret_cast<float2*>(destination + 8 * SPATIAL) =
                    make_float2(
                        acc[i][j][2] + bias1,
                        acc[i][j][3] + bias1);
            }
        }
    }
#endif
}

torch::Tensor conv3d_mma_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");

    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_mma_kernel<<<grid, 256, 0, stream>>>(
        packed_ptr,
        packed_ptr + input_elements,
        bias_ptr,
        output.data_ptr<float>());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="conv3d_compensated_direct_mma_v4",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["conv3d_mma_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_mma_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 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)
```