I’ll focus on the output expansion, which accounts for most of the runtime, using vectorized stores and cheaper coordinate calculations.

```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 transform_weight(torch::Tensor weight);
torch::Tensor expand_output(
    torch::Tensor input, torch::Tensor bias, bool has_bias);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/Dispatch.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <cstdint>
#include <limits>

template <typename scalar_t>
__global__ void transform_weight_kernel(
    const scalar_t* __restrict__ weight,
    scalar_t* __restrict__ output,
    int in_channels,
    int out_channels,
    int64_t count)
{
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count;
         i += int64_t(blockDim.x) * gridDim.x) {
        const int k = i % 27;
        const int ic = (i / 27) % in_channels;
        const int oc = i / (27LL * in_channels);
        output[i] =
            weight[(int64_t(ic) * out_channels + oc) * 27 + 26 - k];
    }
}

// Each thread writes sixteen aligned bytes. Constant dimensions allow
// coordinate divisions to compile into multiplication and shifts.
template <int D, int H, int W, bool HAS_BIAS>
__global__ void expand_vector_kernel(
    const float* __restrict__ input,
    const float* __restrict__ bias,
    float4* __restrict__ output,
    int channels,
    uint32_t vector_count)
{
    constexpr uint32_t OD = 2 * D + 1;
    constexpr uint32_t OH = 2 * H + 1;
    constexpr uint32_t OW = 2 * W + 1;

    const uint32_t step = blockDim.x * gridDim.x;
    for (uint32_t v = blockIdx.x * blockDim.x + threadIdx.x;
         v < vector_count;
         v += step) {
        const uint32_t index = 4 * v;
        const uint32_t row = index / OW;
        const uint32_t w = index - row * OW;
        const uint32_t plane = row / OH;
        const uint32_t h = row - plane * OH;
        const uint32_t nc = plane / OD;
        const uint32_t d = plane - nc * OD;

        float4 result;

        if (w + 3 < OW) {
            float b = 0.0f;
            if (HAS_BIAS)
                b = bias[nc % channels];

            result = make_float4(b, b, b, b);
            if ((d & h & 1u) != 0) {
                const uint32_t source =
                    ((nc * D + d / 2) * H + h / 2) * W + w / 2;

                float a = input[source];
                float c = input[source + 1];
                if (HAS_BIAS) {
                    a += b;
                    c += b;
                }

                if (w & 1u) {
                    result.x = a;
                    result.z = c;
                } else {
                    result.y = a;
                    result.w = c;
                }
            }
        } else {
            // Handle vectors crossing row, plane, or channel boundaries.
            float values[4];
            #pragma unroll
            for (uint32_t j = 0; j < 4; ++j) {
                const uint32_t p = index + j;
                const uint32_t r = p / OW;
                const uint32_t ww = p - r * OW;
                const uint32_t pl = r / OH;
                const uint32_t hh = r - pl * OH;
                const uint32_t nncc = pl / OD;
                const uint32_t dd = pl - nncc * OD;

                float value = 0.0f;
                if ((dd & hh & ww & 1u) != 0) {
                    const uint32_t source =
                        ((nncc * D + dd / 2) * H + hh / 2) * W + ww / 2;
                    value = input[source];
                }
                if (HAS_BIAS)
                    value += bias[nncc % channels];
                values[j] = value;
            }
            result = make_float4(values[0], values[1], values[2], values[3]);
        }

        output[v] = result;
    }
}

template <typename scalar_t, bool HAS_BIAS>
__global__ void expand_generic_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ bias,
    scalar_t* __restrict__ output,
    int channels,
    int depth,
    int height,
    int width,
    int64_t count)
{
    const int od = 2 * depth + 1;
    const int oh = 2 * height + 1;
    const int ow = 2 * width + 1;

    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count;
         i += int64_t(blockDim.x) * gridDim.x) {
        int64_t q = i;
        const int w = q % ow;
        q /= ow;
        const int h = q % oh;
        q /= oh;
        const int d = q % od;
        const int64_t nc = q / od;

        scalar_t value = scalar_t(0);
        if ((d & h & w & 1) != 0) {
            value = input[
                ((nc * depth + d / 2) * height + h / 2) * width + w / 2];
        }
        if (HAS_BIAS)
            value = value + bias[nc % channels];
        output[i] = value;
    }
}

torch::Tensor transform_weight(torch::Tensor weight)
{
    const c10::cuda::CUDAGuard guard(weight.device());
    const int ic = weight.size(0);
    const int oc = weight.size(1);
    auto output = torch::empty({oc, ic, 3, 3, 3}, weight.options());

    const int threads = 256;
    const int blocks = static_cast<int>(
        std::min<int64_t>((output.numel() + threads - 1) / threads, 4096));
    const auto stream = c10::cuda::getCurrentCUDAStream();

    AT_DISPATCH_FLOATING_TYPES_AND2(
        at::ScalarType::Half, at::ScalarType::BFloat16,
        weight.scalar_type(), "transform_weight", [&] {
            transform_weight_kernel<scalar_t>
                <<<blocks, threads, 0, stream.stream()>>>(
                    weight.data_ptr<scalar_t>(),
                    output.data_ptr<scalar_t>(),
                    ic, oc, output.numel());
        });

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}

torch::Tensor expand_output(
    torch::Tensor input, torch::Tensor bias, bool has_bias)
{
    const c10::cuda::CUDAGuard guard(input.device());
    const int64_t batch = input.size(0);
    const int channels = input.size(1);
    const int depth = input.size(2);
    const int height = input.size(3);
    const int width = input.size(4);

    auto output = torch::empty(
        {batch, int64_t(channels),
         2LL * depth + 1, 2LL * height + 1, 2LL * width + 1},
        input.options());

    const int64_t count = output.numel();
    if (count == 0)
        return output;

    const int threads = 256;
    const auto stream = c10::cuda::getCurrentCUDAStream();

    if (input.scalar_type() == at::ScalarType::Float &&
        depth == 16 && height == 32 && width == 32 &&
        count % 4 == 0 &&
        count <= std::numeric_limits<int32_t>::max()) {
        const uint32_t vectors = static_cast<uint32_t>(count / 4);
        const int blocks = static_cast<int>(
            std::min<uint32_t>((vectors + threads - 1) / threads, 4096));

        if (has_bias) {
            expand_vector_kernel<16, 32, 32, true>
                <<<blocks, threads, 0, stream.stream()>>>(
                    input.data_ptr<float>(),
                    bias.data_ptr<float>(),
                    reinterpret_cast<float4*>(output.data_ptr<float>()),
                    channels, vectors);
        } else {
            expand_vector_kernel<16, 32, 32, false>
                <<<blocks, threads, 0, stream.stream()>>>(
                    input.data_ptr<float>(),
                    nullptr,
                    reinterpret_cast<float4*>(output.data_ptr<float>()),
                    channels, vectors);
        }
    } else {
        const int blocks = static_cast<int>(
            std::min<int64_t>((count + threads - 1) / threads, 4096));

        AT_DISPATCH_FLOATING_TYPES_AND2(
            at::ScalarType::Half, at::ScalarType::BFloat16,
            input.scalar_type(), "expand_output", [&] {
                if (has_bias) {
                    expand_generic_kernel<scalar_t, true>
                        <<<blocks, threads, 0, stream.stream()>>>(
                            input.data_ptr<scalar_t>(),
                            bias.data_ptr<scalar_t>(),
                            output.data_ptr<scalar_t>(),
                            channels, depth, height, width, count);
                } else {
                    expand_generic_kernel<scalar_t, false>
                        <<<blocks, threads, 0, stream.stream()>>>(
                            input.data_ptr<scalar_t>(),
                            nullptr,
                            output.data_ptr<scalar_t>(),
                            channels, depth, height, width, count);
                }
            });
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_ops = load_inline(
    name="sparse_transpose3d_vector_expand_v2",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["transform_weight", "expand_output"],
    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,
        bias: bool = False,
    ):
        super().__init__()
        self.conv_transpose3d = nn.ConvTranspose3d(
            in_channels,
            out_channels,
            kernel_size=(kernel_size, kernel_size, kernel_size),
            stride=stride,
            padding=padding,
            dilation=dilation,
            bias=bias,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv_transpose3d
        supported = (
            x.is_cuda
            and x.ndim == 5
            and x.dtype == conv.weight.dtype
            and x.dtype in (
                torch.float16,
                torch.bfloat16,
                torch.float32,
                torch.float64,
            )
            and conv.kernel_size == (3, 3, 3)
            and conv.stride == (2, 2, 2)
            and conv.padding == (1, 1, 1)
            and conv.dilation == (2, 2, 2)
            and conv.output_padding == (0, 0, 0)
            and conv.groups == 1
        )
        if not supported:
            return conv(x)

        weight = _ops.transform_weight(conv.weight.contiguous())
        compact = F.conv3d(
            x.contiguous(),
            weight,
            bias=None,
            stride=1,
            padding=1,
            dilation=1,
        ).contiguous()

        has_bias = conv.bias is not None
        bias_tensor = (
            conv.bias.to(dtype=compact.dtype).contiguous()
            if has_bias
            else compact
        )
        return _ops.expand_output(compact, bias_tensor, has_bias)
```