```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_k331_cuda(
    torch::Tensor input,
    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_runtime.h>

__global__ __launch_bounds__(256)
void conv3d_k331_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int channels_out,
    int height,
    int width,
    int depth,
    int height_out,
    int width_out,
    int64_t spatial_out)
{
    constexpr int TILE_C = 8;
    __shared__ float weights[27][TILE_C];
    __shared__ float biases[TILE_C];

    const int channel_base = blockIdx.y * TILE_C;

    for (int i = threadIdx.x; i < 27 * TILE_C; i += blockDim.x) {
        const int k = i / TILE_C;
        const int j = i % TILE_C;
        const int oc = channel_base + j;
        weights[k][j] = oc < channels_out ? weight[oc * 27 + k] : 0.0f;
    }

    if (threadIdx.x < TILE_C) {
        const int oc = channel_base + threadIdx.x;
        biases[threadIdx.x] =
            bias != nullptr && oc < channels_out ? bias[oc] : 0.0f;
    }
    __syncthreads();

    const int64_t p =
        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (p >= spatial_out)
        return;

    const int od = p % depth;
    const int ow = (p / depth) % width_out;
    const int oh = p / (static_cast<int64_t>(depth) * width_out);
    const int batch = blockIdx.z;

    const int64_t plane = static_cast<int64_t>(width) * depth;
    const int64_t channel_stride = static_cast<int64_t>(height) * plane;
    const int64_t input_base =
        static_cast<int64_t>(batch) * 3 * channel_stride +
        static_cast<int64_t>(oh) * plane +
        static_cast<int64_t>(ow) * depth + od;

    float accum[TILE_C] = {};

    #pragma unroll
    for (int ic = 0; ic < 3; ++ic) {
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                const int k = ic * 9 + kh * 3 + kw;
                const float value = input[
                    input_base + ic * channel_stride +
                    kh * plane + kw * depth];

                #pragma unroll
                for (int j = 0; j < TILE_C; ++j) {
                    accum[j] = fmaf(value, weights[k][j], accum[j]);
                }
            }
        }
    }

    #pragma unroll
    for (int j = 0; j < TILE_C; ++j) {
        const int oc = channel_base + j;
        if (oc < channels_out) {
            const int64_t index =
                (static_cast<int64_t>(batch) * channels_out + oc) *
                spatial_out + p;
            output[index] = accum[j] + biases[j];
        }
    }
}

torch::Tensor conv3d_k331_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias)
{
    const c10::cuda::CUDAGuard device_guard(input.device());

    const int batch = input.size(0);
    const int height = input.size(2);
    const int width = input.size(3);
    const int depth = input.size(4);
    const int channels_out = weight.size(0);
    const int height_out = height - 2;
    const int width_out = width - 2;

    auto output = torch::empty(
        {batch, channels_out, height_out, width_out, depth},
        input.options());

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

    const int64_t spatial_out =
        static_cast<int64_t>(height_out) * width_out * depth;

    const dim3 block(256);
    const dim3 grid(
        (spatial_out + 255) / 256,
        (channels_out + 7) / 8,
        batch);

    const float* bias_ptr =
        bias.has_value() ? bias.value().data_ptr<float>() : nullptr;

    conv3d_k331_kernel<<<
        grid, block, 0, at::cuda::getCurrentCUDAStream()
    >>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias_ptr,
        output.data_ptr<float>(),
        channels_out,
        height,
        width,
        depth,
        height_out,
        width_out,
        spatial_out);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_conv3d_extension = load_inline(
    name="custom_conv3d_k331_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["conv3d_k331_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


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, 1),
            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.ndim == 5
            and x.shape[1] == 3
            and x.shape[2] >= 3
            and x.shape[3] >= 3
            and x.shape[4] > 0
            and x.shape[0] <= 65535
            and conv.in_channels == 3
            and conv.out_channels <= 8 * 65535
            and conv.kernel_size == (3, 3, 1)
            and conv.stride == (1, 1, 1)
            and conv.padding == (0, 0, 0)
            and conv.dilation == (1, 1, 1)
            and conv.groups == 1
            and conv.weight.dtype == torch.float32
            and conv.weight.device == x.device
            and not torch.is_autocast_enabled()
        ):
            return _conv3d_extension.conv3d_k331_cuda(
                x.contiguous(),
                conv.weight.contiguous(),
                conv.bias,
            )
        return conv(x)
```