```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 conv331_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_runtime.h>

constexpr int CHANNEL_TILE = 8;
constexpr int BLOCK_SIZE = 128;

__global__ void conv331_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int64_t height,
    int64_t width,
    int64_t depth,
    int64_t output_width,
    int64_t spatial_size,
    int out_channels)
{
    const int64_t p =
        static_cast<int64_t>(blockIdx.x) * BLOCK_SIZE + threadIdx.x;
    if (p >= spatial_size) {
        return;
    }

    const int first_channel = blockIdx.y * CHANNEL_TILE;
    const int64_t batch = blockIdx.z;
    const int64_t row = p / (output_width * depth);
    const int64_t remainder = p - row * output_width * depth;
    const int64_t column = remainder / depth;
    const int64_t z = remainder - column * depth;

    const int64_t input_channel_stride = height * width * depth;
    const int64_t input_base =
        batch * 3 * input_channel_stride +
        (row * width + column) * depth + z;

    float acc[CHANNEL_TILE];

    #pragma unroll
    for (int o = 0; o < CHANNEL_TILE; ++o) {
        acc[o] = 0.0f;
    }

    #pragma unroll
    for (int c = 0; c < 3; ++c) {
        #pragma unroll
        for (int kh = 0; kh < 3; ++kh) {
            #pragma unroll
            for (int kw = 0; kw < 3; ++kw) {
                const float value = __ldg(
                    x + input_base +
                    c * input_channel_stride +
                    (kh * width + kw) * depth);

                const int k = c * 9 + kh * 3 + kw;

                #pragma unroll
                for (int o = 0; o < CHANNEL_TILE; ++o) {
                    const int oc = first_channel + o;
                    if (oc < out_channels) {
                        acc[o] = fmaf(
                            value, __ldg(weight + oc * 27 + k), acc[o]);
                    }
                }
            }
        }
    }

    #pragma unroll
    for (int o = 0; o < CHANNEL_TILE; ++o) {
        const int oc = first_channel + o;
        if (oc < out_channels) {
            float result = acc[o];
            if (bias != nullptr) {
                result += __ldg(bias + oc);
            }
            output[(batch * out_channels + oc) * spatial_size + p] = result;
        }
    }
}

torch::Tensor conv331_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    c10::optional<torch::Tensor> bias)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "Expected CUDA tensors");
    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.dim() == 5 && x.size(1) == 3 &&
        weight.dim() == 5 && weight.size(1) == 3 &&
        weight.size(2) == 3 && weight.size(3) == 3 &&
        weight.size(4) == 1,
        "Unsupported convolution shape");
    TORCH_CHECK(x.device() == weight.device(), "Device mismatch");
    TORCH_CHECK(x.size(2) >= 3 && x.size(3) >= 3 && x.size(4) > 0,
                "Input is too small");

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

    const int64_t batch = x.size(0);
    const int64_t height = x.size(2);
    const int64_t width = x.size(3);
    const int64_t depth = x.size(4);
    const int out_channels = static_cast<int>(weight.size(0));
    const int64_t output_height = height - 2;
    const int64_t output_width = width - 2;
    const int64_t spatial_size = output_height * output_width * depth;

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

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

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

    TORCH_CHECK(batch <= 65535, "Batch exceeds CUDA grid limit");
    TORCH_CHECK((out_channels + CHANNEL_TILE - 1) / CHANNEL_TILE <= 65535,
                "Channels exceed CUDA grid limit");
    TORCH_CHECK((spatial_size + BLOCK_SIZE - 1) / BLOCK_SIZE <= 2147483647LL,
                "Spatial size exceeds CUDA grid limit");

    const dim3 grid(
        static_cast<unsigned>((spatial_size + BLOCK_SIZE - 1) / BLOCK_SIZE),
        (out_channels + CHANNEL_TILE - 1) / CHANNEL_TILE,
        static_cast<unsigned>(batch));

    conv331_kernel<<<
        grid, BLOCK_SIZE, 0, at::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(),
            weight.data_ptr<float>(),
            bias_ptr,
            output.data_ptr<float>(),
            height,
            width,
            depth,
            output_width,
            spatial_size,
            out_channels);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_conv331 = load_inline(
    name="conv331_channel_tiled_cuda_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["conv331_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _Conv331Function(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias):
        ctx.save_for_backward(x, weight)
        return _conv331.conv331_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,
                stride=1, padding=0, dilation=1, groups=1
            )
        if ctx.needs_input_grad[1]:
            grad_weight = torch.nn.grad.conv3d_weight(
                x, weight.shape, grad_output,
                stride=1, padding=0, dilation=1, groups=1
            )
        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, 1),
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv3d
        use_custom = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 5
            and x.is_contiguous()
            and conv.weight.is_contiguous()
            and conv.weight.dtype == torch.float32
            and conv.weight.device == x.device
            and conv.in_channels == 3
            and x.shape[1] == 3
            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 x.shape[2] >= 3
            and x.shape[3] >= 3
            and x.shape[4] > 0
            and x.shape[0] <= 65535
            and conv.out_channels <= 8 * 65535
            and not torch.is_autocast_enabled()
        )

        if use_custom:
            return _Conv331Function.apply(x, conv.weight, conv.bias)
        return conv(x)
```