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


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

torch::Tensor conv3d_split_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_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>

namespace {

constexpr int TILE_M = 32;
constexpr int TILE_N = 64;
constexpr int PAD_K = 96;

__device__ __forceinline__ void split_float(
    float value, half& high, half& low) {
    high = __float2half_rn(value);
    low = __float2half_rn(
        (value - __half2float(high)) * 2048.0f);
}

__global__ void conv3d_split_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int channels_out,
    int depth,
    int height,
    int width,
    int output_height,
    int output_width,
    int spatial_size) {

    const int tid = threadIdx.x;
    const int channel_start = blockIdx.y * TILE_M;
    const int spatial_start = blockIdx.x * TILE_N;
    const int batch = blockIdx.z;

    const int64_t input_volume =
        static_cast<int64_t>(depth) * height * width;
    const float* input_batch = x + batch * 3 * input_volume;

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
    using namespace nvcuda;

    __shared__ __align__(32) half weight_high[TILE_M * PAD_K];
    __shared__ __align__(32) half weight_low[TILE_M * PAD_K];

    union SharedStorage {
        struct {
            half high[PAD_K * TILE_N];
            half low[PAD_K * TILE_N];
        } input;
        float output[TILE_M * TILE_N];
    };
    __shared__ __align__(32) SharedStorage storage;

    for (int index = tid; index < TILE_M * PAD_K; index += 256) {
        const int row = index / PAD_K;
        const int k = index % PAD_K;
        const int channel = channel_start + row;

        float value = 0.0f;
        if (channel < channels_out && k < 81) {
            value = weight[channel * 81 + k];
        }
        split_float(value, weight_high[index], weight_low[index]);
    }

    const int column = tid % TILE_N;
    const int position = spatial_start + column;
    const int plane = output_height * output_width;
    const int od = position / plane;
    const int remainder = position % plane;
    const int oh = remainder / output_width;
    const int ow = remainder % output_width;

    const int64_t input_base =
        (static_cast<int64_t>(od) * height + oh) * width + ow;

    for (int k = tid / TILE_N; k < PAD_K; k += 4) {
        float value = 0.0f;
        if (position < spatial_size && k < 81) {
            const int ic = k / 27;
            const int kernel_index = k % 27;
            const int kd = kernel_index / 9;
            const int kh = (kernel_index % 9) / 3;
            const int kw = kernel_index % 3;

            const int64_t offset =
                ic * input_volume + input_base
                + static_cast<int64_t>(kd) * height * width
                + kh * width + kw;
            value = input_batch[offset];
        }

        const int index = k * TILE_N + column;
        split_float(
            value, storage.input.high[index], storage.input.low[index]);
    }

    __syncthreads();

    const int warp = tid / 32;
    const int row_start = (warp / 4) * 16;
    const int col_start = (warp % 4) * 16;

    wmma::fragment<
        wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_high, a_low;
    wmma::fragment<
        wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_high, b_low;
    wmma::fragment<
        wmma::accumulator, 16, 16, 16, float> main_sum, cross_sum, low_sum;

    wmma::fill_fragment(main_sum, 0.0f);
    wmma::fill_fragment(cross_sum, 0.0f);
    wmma::fill_fragment(low_sum, 0.0f);

    #pragma unroll
    for (int k = 0; k < PAD_K; k += 16) {
        wmma::load_matrix_sync(
            a_high, weight_high + row_start * PAD_K + k, PAD_K);
        wmma::load_matrix_sync(
            a_low, weight_low + row_start * PAD_K + k, PAD_K);
        wmma::load_matrix_sync(
            b_high, storage.input.high + k * TILE_N + col_start, TILE_N);
        wmma::load_matrix_sync(
            b_low, storage.input.low + k * TILE_N + col_start, TILE_N);

        wmma::mma_sync(main_sum, a_high, b_high, main_sum);
        wmma::mma_sync(cross_sum, a_high, b_low, cross_sum);
        wmma::mma_sync(cross_sum, a_low, b_high, cross_sum);
        wmma::mma_sync(low_sum, a_low, b_low, low_sum);
    }

    #pragma unroll
    for (int i = 0; i < main_sum.num_elements; ++i) {
        main_sum.x[i] +=
            cross_sum.x[i] * (1.0f / 2048.0f)
            + low_sum.x[i] * (1.0f / 4194304.0f);
    }

    // All warps must finish reading before reusing shared input storage.
    __syncthreads();

    wmma::store_matrix_sync(
        storage.output + row_start * TILE_N + col_start,
        main_sum,
        TILE_N,
        wmma::mem_row_major);

    __syncthreads();

    for (int index = tid; index < TILE_M * TILE_N; index += 256) {
        const int channel = channel_start + index / TILE_N;
        const int position_out = spatial_start + index % TILE_N;

        if (channel < channels_out && position_out < spatial_size) {
            float value = storage.output[index];
            if (bias != nullptr) {
                value += bias[channel];
            }
            const int64_t output_index =
                (static_cast<int64_t>(batch) * channels_out + channel)
                * spatial_size + position_out;
            output[output_index] = value;
        }
    }
#else
    for (int index = tid; index < TILE_M * TILE_N; index += 256) {
        const int channel = channel_start + index / TILE_N;
        const int position = spatial_start + index % TILE_N;
        if (channel >= channels_out || position >= spatial_size) {
            continue;
        }

        const int od = position / (output_height * output_width);
        const int remainder = position % (output_height * output_width);
        const int oh = remainder / output_width;
        const int ow = remainder % output_width;

        float sum = 0.0f;
        #pragma unroll
        for (int ic = 0; ic < 3; ++ic) {
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd) {
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    #pragma unroll
                    for (int kw = 0; kw < 3; ++kw) {
                        const int64_t input_index =
                            ic * input_volume
                            + (static_cast<int64_t>(od + kd) * height
                               + oh + kh) * width + ow + kw;
                        const int weight_index =
                            channel * 81 + ic * 27 + kd * 9 + kh * 3 + kw;
                        sum = fmaf(
                            input_batch[input_index],
                            weight[weight_index],
                            sum);
                    }
                }
            }
        }

        if (bias != nullptr) {
            sum += bias[channel];
        }
        output[
            (static_cast<int64_t>(batch) * channels_out + channel)
            * spatial_size + position] = sum;
    }
#endif
}

} // namespace

torch::Tensor conv3d_split_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias) {

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

    const int batch = x.size(0);
    const int channels_out = weight.size(0);
    const int depth = x.size(2);
    const int height = x.size(3);
    const int width = x.size(4);
    const int output_depth = depth - 2;
    const int output_height = height - 2;
    const int output_width = width - 2;
    const int spatial_size =
        output_depth * output_height * output_width;

    auto output = torch::empty(
        {batch, channels_out, output_depth, output_height, output_width},
        x.options());

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

    const float* bias_pointer = nullptr;
    if (bias.has_value()) {
        bias_pointer = bias.value().data_ptr<float>();
    }

    const dim3 grid(
        (spatial_size + TILE_N - 1) / TILE_N,
        (channels_out + TILE_M - 1) / TILE_M,
        batch);

    conv3d_split_kernel<<<
        grid, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias_pointer,
        output.data_ptr<float>(),
        channels_out,
        depth,
        height,
        width,
        output_height,
        output_width,
        spatial_size);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_conv3d_extension = load_inline(
    name="conv3d_split_m32n64k96",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["conv3d_split_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)
        return _conv3d_extension.conv3d_split_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.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
        weight = conv.weight
        bias = conv.bias

        use_custom_kernel = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 5
            and x.is_contiguous()
            and x.shape[1] == 3
            and min(x.shape[2:]) >= 3
            and weight.device == x.device
            and weight.dtype == torch.float32
            and weight.is_contiguous()
            and conv.in_channels == 3
            and conv.kernel_size == (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 (
                bias is None
                or (
                    bias.device == x.device
                    and bias.dtype == torch.float32
                    and bias.is_contiguous()
                )
            )
            and not torch.is_autocast_enabled()
        )

        if use_custom_kernel:
            return _Conv3dFunction.apply(x, weight, bias)

        return conv(x)
```