```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Pipeline (everything stays in channels_last_3d end to end):
#   0) Custom NCHW->NHWC transpose (tiled 32x32 smem) for the NCHW input.
#   1) stride-1/pad-0 ConvTranspose3d evaluated as an equivalent forward
#      Conv3d (flipped + channel-swapped kernel, full padding k-1) so cuDNN
#      uses the fast NHWC implicit-GEMM path instead of the dgrad path with
#      layout-conversion kernels. Verified once at runtime, permanent
#      fallback if the check fails.
#   2) ReLU+GroupNorm fused into custom kernels (stats with ReLU on the fly,
#      tiny affine finalize, vectorized apply) consuming the NHWC conv output
#      directly; ReLU is never materialized.
#   3) The batch is split into chunks pipelined across two CUDA streams:
#      transposes+convs on the caller stream, fused GN on a side stream, so
#      the memory-bound GN of chunk i overlaps the compute-bound conv of
#      chunk i+1. Cross-stream ordering is enforced with events.
fastpath_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <climits>

#define STATS_BLOCK 256
#define APPLY_BLOCK 256
#define APPLY_ROWS 8

__global__ void nchw2nhwc_kernel(const float* __restrict__ src,
                                 float* __restrict__ dst,
                                 int C, int S) {
    __shared__ float tile[32][33];
    const int n  = blockIdx.z;
    const int c0 = blockIdx.x * 32;
    const int s0 = blockIdx.y * 32;
    const size_t base = (size_t)n * C * S;
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int c = c0 + ty + i;
        const int s = s0 + tx;
        if (c < C && s < S) tile[ty + i][tx] = src[base + (size_t)c * S + s];
    }
    __syncthreads();
    #pragma unroll
    for (int i = 0; i < 32; i += 8) {
        const int s = s0 + ty + i;
        const int c = c0 + tx;
        if (s < S && c < C) dst[base + (size_t)s * C + c] = tile[tx][ty + i];
    }
}

torch::Tensor nchw_to_nhwc(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat && x.dim() == 5,
                "bad input");
    auto xc = x.contiguous();
    const int64_t N = xc.size(0), C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    TORCH_CHECK(S <= INT_MAX && C <= INT_MAX && N <= 65535, "too big");
    auto y = at::empty_like(xc, at::MemoryFormat::ChannelsLast3d);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const dim3 block(32, 8, 1);
    const dim3 grid((unsigned)((C + 31) / 32), (unsigned)((S + 31) / 32),
                    (unsigned)N);
    nchw2nhwc_kernel<<<grid, block, 0, stream>>>(
        xc.data_ptr<float>(), y.data_ptr<float>(), (int)C, (int)S);
    return y;
}

template <int BLOCK>
__device__ __forceinline__ void block_reduce2(float fs, float fq,
                                              double* out_s, double* out_q) {
    double ds = (double)fs;
    double dq = (double)fq;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        ds += __shfl_down_sync(0xffffffffu, ds, off);
        dq += __shfl_down_sync(0xffffffffu, dq, off);
    }
    __shared__ double shs[BLOCK / 32];
    __shared__ double shq[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) { shs[wid] = ds; shq[wid] = dq; }
    __syncthreads();
    if (wid == 0) {
        constexpr int NW = BLOCK / 32;
        ds = (lane < NW) ? shs[lane] : 0.0;
        dq = (lane < NW) ? shq[lane] : 0.0;
        #pragma unroll
        for (int off = NW / 2; off > 0; off >>= 1) {
            ds += __shfl_down_sync(0xffffffffu, ds, off);
            dq += __shfl_down_sync(0xffffffffu, dq, off);
        }
        if (lane == 0) { *out_s = ds; *out_q = dq; }
    }
}

__device__ __forceinline__ void relu4_accum(const float4 v, float& fs, float& fq) {
    const float a0 = fmaxf(v.x, 0.0f);
    const float a1 = fmaxf(v.y, 0.0f);
    const float a2 = fmaxf(v.z, 0.0f);
    const float a3 = fmaxf(v.w, 0.0f);
    fs += (a0 + a1) + (a2 + a3);
    fq += (a0 * a0 + a1 * a1) + (a2 * a2 + a3 * a3);
}

template <int VSHIFT>
__global__ void stats_pow2_kernel(const float* __restrict__ x,
                                  double* __restrict__ partials,
                                  int S, int C, int Cg, int G,
                                  int rows_per_split, int splits) {
    constexpr int V = 1 << VSHIFT;
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    const int c4 = threadIdx.x & (V - 1);
    const int r0 = threadIdx.x >> VSHIFT;
    const int rstride = STATS_BLOCK >> VSHIFT;
    float fs = 0.0f, fq = 0.0f;
    for (int r = r0; r < s1 - s0; r += rstride) {
        const float4 v = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + r) * C + (c4 << 2));
        relu4_accum(v, fs, fq);
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void stats_vec4_kernel(const float* __restrict__ x,
                                  double* __restrict__ partials,
                                  int S, int C, int Cg, int G, int V,
                                  int rows_per_split, int splits) {
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    float fs = 0.0f, fq = 0.0f;
    const int count = (s1 - s0) * V;
    for (int idx = threadIdx.x; idx < count; idx += STATS_BLOCK) {
        const int row = idx / V;
        const int c4  = idx - row * V;
        const float4 v = *reinterpret_cast<const float4*>(
            base + (size_t)(s0 + row) * C + (c4 << 2));
        relu4_accum(v, fs, fq);
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void stats_scalar_kernel(const float* __restrict__ x,
                                    double* __restrict__ partials,
                                    int S, int C, int Cg, int G,
                                    int rows_per_split, int splits) {
    const int n  = blockIdx.z;
    const int g  = blockIdx.y;
    const int s0 = blockIdx.x * rows_per_split;
    const int s1 = ::min(S, s0 + rows_per_split);
    const float* base = x + ((size_t)n * S) * C + (size_t)g * Cg;
    float fs = 0.0f, fq = 0.0f;
    for (int r = s0; r < s1; ++r) {
        const float* rowp = base + (size_t)r * C;
        for (int c = threadIdx.x; c < Cg; c += STATS_BLOCK) {
            const float v = fmaxf(rowp[c], 0.0f);
            fs += v;
            fq += v * v;
        }
    }
    double rs, rq;
    block_reduce2<STATS_BLOCK>(fs, fq, &rs, &rq);
    if (threadIdx.x == 0) {
        double* p = partials + ((size_t)(n * G + g) * splits + blockIdx.x) * 2;
        p[0] = rs;
        p[1] = rq;
    }
}

__global__ void finalize_affine_kernel(const double* __restrict__ partials,
                                       const float* __restrict__ gamma,
                                       const float* __restrict__ beta,
                                       float* __restrict__ a,
                                       float* __restrict__ b,
                                       int NC, int C, int Cg, int G,
                                       int splits, double inv_gs, double eps) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= NC) return;
    const int n = idx / C;
    const int c = idx - n * C;
    const int g = c / Cg;
    const double* p = partials + (size_t)(n * G + g) * splits * 2;
    double s = 0.0, q = 0.0;
    for (int j = 0; j < splits; ++j) {
        s += p[2 * j];
        q += p[2 * j + 1];
    }
    const double m = s * inv_gs;
    double var = q * inv_gs - m * m;
    if (var < 0.0) var = 0.0;
    const double rstd = 1.0 / sqrt(var + eps);
    const float av = (float)(rstd * (double)gamma[c]);
    a[idx] = av;
    b[idx] = (float)((double)beta[c] - m * (double)av);
}

template <int VSHIFT>
__global__ void apply_pow2_kernel(const float* __restrict__ x,
                                  const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int S, int C) {
    constexpr int V = 1 << VSHIFT;
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c4  = tid & (V - 1);
    const int s0  = tid >> VSHIFT;
    const int sstride = (gridDim.x * APPLY_BLOCK) >> VSHIFT;
    if (s0 >= S) return;
    const size_t img = (size_t)n * S * V;
    const float4 av = *reinterpret_cast<const float4*>(a + (size_t)n * C + (c4 << 2));
    const float4 bv = *reinterpret_cast<const float4*>(b + (size_t)n * C + (c4 << 2));
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * V + c4;
        const float4 v = *reinterpret_cast<const float4*>(x + (off << 2));
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), av.x, bv.x);
        o.y = fmaf(fmaxf(v.y, 0.0f), av.y, bv.y);
        o.z = fmaf(fmaxf(v.z, 0.0f), av.z, bv.z);
        o.w = fmaf(fmaxf(v.w, 0.0f), av.w, bv.w);
        *reinterpret_cast<float4*>(y + (off << 2)) = o;
    }
}

__global__ void apply_vec4_kernel(const float* __restrict__ x,
                                  const float* __restrict__ a,
                                  const float* __restrict__ b,
                                  float* __restrict__ y,
                                  int S, int C, int V) {
    const int n   = blockIdx.y;
    const int tid = blockIdx.x * APPLY_BLOCK + threadIdx.x;
    const int c4  = tid % V;
    const int s0  = tid / V;
    const int sstride = (gridDim.x * APPLY_BLOCK) / V;
    if (s0 >= S || sstride <= 0) return;
    const size_t img = (size_t)n * S * V;
    const float4 av = *reinterpret_cast<const float4*>(a + (size_t)n * C + (c4 << 2));
    const float4 bv = *reinterpret_cast<const float4*>(b + (size_t)n * C + (c4 << 2));
    for (int s = s0; s < S; s += sstride) {
        const size_t off = img + (size_t)s * V + c4;
        const float4 v = *reinterpret_cast<const float4*>(x + (off << 2));
        float4 o;
        o.x = fmaf(fmaxf(v.x, 0.0f), av.x, bv.x);
        o.y = fmaf(fmaxf(v.y, 0.0f), av.y, bv.y);
        o.z = fmaf(fmaxf(v.z, 0.0f), av.z, bv.z);
        o.w = fmaf(fmaxf(v.w, 0.0f), av.w, bv.w);
        *reinterpret_cast<float4*>(y + (off << 2)) = o;
    }
}

__global__ void apply_scalar_kernel(const float* __restrict__ x,
                                    const float* __restrict__ a,
                                    const float* __restrict__ b,
                                    float* __restrict__ y,
                                    size_t total, int C) {
    const int n = blockIdx.y;
    const size_t base = (size_t)n * total;
    const size_t ab   = (size_t)n * C;
    const size_t stride = (size_t)gridDim.x * APPLY_BLOCK;
    for (size_t p = (size_t)blockIdx.x * APPLY_BLOCK + threadIdx.x;
         p < total; p += stride) {
        const int c = (int)(p % (size_t)C);
        const float v = fmaxf(x[base + p], 0.0f);
        y[base + p] = fmaf(v, a[ab + c], b[ab + c]);
    }
}

static inline int log2_int(int v) {
    int s = 0;
    while (v > 1) { v >>= 1; ++s; }
    return s;
}

torch::Tensor fused_relu_groupnorm_nhwc(torch::Tensor x, torch::Tensor gamma,
                                        torch::Tensor beta, int64_t groups,
                                        double eps, torch::Tensor out) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 5, "x must have shape (N, C, D, H, W)");
    TORCH_CHECK(out.numel() == x.numel() && out.scalar_type() == at::kFloat,
                "bad out");

    auto xc = x.contiguous(at::MemoryFormat::ChannelsLast3d);
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();

    const int64_t N = xc.size(0);
    const int64_t C = xc.size(1);
    const int64_t S = xc.size(2) * xc.size(3) * xc.size(4);
    const int64_t G = groups;
    TORCH_CHECK(G > 0 && C % G == 0, "C must be divisible by groups");
    const int64_t Cg = C / G;
    TORCH_CHECK(S > 0 && S <= INT_MAX && C <= INT_MAX && Cg <= INT_MAX,
                "bad tensor dims");
    TORCH_CHECK(N <= 65535 && G <= 65535, "N/G too large for grid");
    TORCH_CHECK(gc.numel() == C && bc.numel() == C, "gamma/beta size mismatch");

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t NG = N * G;
    int splits = (int)((8LL * sm_count + NG - 1) / NG);
    if (splits < 1) splits = 1;
    if (splits > 32) splits = 32;
    if ((int64_t)splits > S) splits = (int)S;
    int rows_per_split = (int)((S + splits - 1) / splits);
    splits = (int)((S + rows_per_split - 1) / rows_per_split);

    auto partials = at::empty({NG * splits * 2}, xc.options().dtype(at::kDouble));
    const dim3 sgrid(splits, (unsigned)G, (unsigned)N);
    const float* xp = xc.data_ptr<float>();
    double* pp = partials.data_ptr<double>();

    const bool vec_ok = (C % 4 == 0) && (Cg % 4 == 0);
    const int Vs = (int)(Cg / 4);
    const bool vpow2 = vec_ok && Vs > 0 && (Vs & (Vs - 1)) == 0;
    if (vec_ok && vpow2) {
        const int sh = log2_int(Vs);
        switch (sh) {
            case 0: stats_pow2_kernel<0><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 1: stats_pow2_kernel<1><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 2: stats_pow2_kernel<2><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 3: stats_pow2_kernel<3><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 4: stats_pow2_kernel<4><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 5: stats_pow2_kernel<5><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            case 6: stats_pow2_kernel<6><<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits); break;
            default: stats_vec4_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, Vs, rows_per_split, splits); break;
        }
    } else if (vec_ok) {
        stats_vec4_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, Vs, rows_per_split, splits);
    } else {
        stats_scalar_kernel<<<sgrid, STATS_BLOCK, 0, stream>>>(xp, pp, (int)S, (int)C, (int)Cg, (int)G, rows_per_split, splits);
    }

    const int NC = (int)(N * C);
    auto a = at::empty({NC}, xc.options());
    auto b = at::empty({NC}, xc.options());
    finalize_affine_kernel<<<(NC + 255) / 256, 256, 0, stream>>>(
        pp, gc.data_ptr<float>(), bc.data_ptr<float>(),
        a.data_ptr<float>(), b.data_ptr<float>(),
        NC, (int)C, (int)Cg, (int)G, splits,
        1.0 / (double)(Cg * S), eps);

    float* yp = out.data_ptr<float>();
    const float* ap = a.data_ptr<float>();
    const float* bp = b.data_ptr<float>();
    if (C % 4 == 0) {
        const int V = (int)(C / 4);
        int64_t blocks_x = (S * V + (int64_t)APPLY_BLOCK * APPLY_ROWS - 1) /
                           ((int64_t)APPLY_BLOCK * APPLY_ROWS);
        const int64_t min_blocks = (V + APPLY_BLOCK - 1) / APPLY_BLOCK;
        if (blocks_x < min_blocks) blocks_x = min_blocks;
        if (blocks_x < 1) blocks_x = 1;
        const dim3 agrid((unsigned)blocks_x, (unsigned)N, 1);
        const bool vp2 = (V & (V - 1)) == 0;
        if (vp2) {
            const int sh = log2_int(V);
            switch (sh) {
                case 0: apply_pow2_kernel<0><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 1: apply_pow2_kernel<1><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 2: apply_pow2_kernel<2><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 3: apply_pow2_kernel<3><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 4: apply_pow2_kernel<4><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 5: apply_pow2_kernel<5><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 6: apply_pow2_kernel<6><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                case 7: apply_pow2_kernel<7><<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C); break;
                default: apply_vec4_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C, V); break;
            }
        } else {
            apply_vec4_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, (int)S, (int)C, V);
        }
    } else {
        const size_t total = (size_t)S * (size_t)C;
        int64_t blocks_x = ((int64_t)total + (int64_t)APPLY_BLOCK * APPLY_ROWS - 1) /
                           ((int64_t)APPLY_BLOCK * APPLY_ROWS);
        if (blocks_x < 1) blocks_x = 1;
        const dim3 agrid((unsigned)blocks_x, (unsigned)N, 1);
        apply_scalar_kernel<<<agrid, APPLY_BLOCK, 0, stream>>>(xp, ap, bp, yp, total, (int)C);
    }

    return out;
}
"""

fastpath_cpp_source = (
    "torch::Tensor nchw_to_nhwc(torch::Tensor x);\n"
    "torch::Tensor fused_relu_groupnorm_nhwc(torch::Tensor x, torch::Tensor gamma, "
    "torch::Tensor beta, int64_t groups, double eps, torch::Tensor out);"
)

fastpath = load_inline(
    name="convgn_fastpath_ext_v4",
    cpp_sources=fastpath_cpp_source,
    cuda_sources=fastpath_cuda_source,
    functions=["nchw_to_nhwc", "fused_relu_groupnorm_nhwc"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * everything runs in channels_last_3d (cuDNN NHWC tensor-core conv,
        no layout conversions),
      * stride-1 ConvTranspose3d evaluated as equivalent forward Conv3d
        (verified once at runtime, permanent fallback on mismatch),
      * ReLU + GroupNorm fused into custom CUDA kernels reading the NHWC
        conv output directly,
      * batch chunked and pipelined over two CUDA streams so the
        memory-bound fused norm overlaps the compute-bound convolutions.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cuda.matmul.allow_tf32 = True
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, bias=bias
        )
        with torch.no_grad():
            self.conv_transpose.weight.data = (
                self.conv_transpose.weight.data.to(
                    memory_format=torch.channels_last_3d
                )
            )
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)
        self.groups = groups
        self.eps = self.group_norm.eps
        self.fused = fastpath
        if isinstance(kernel_size, (tuple, list)):
            kd, kh, kw = (int(kernel_size[0]), int(kernel_size[1]), int(kernel_size[2]))
        else:
            kd = kh = kw = int(kernel_size)
        self.pad = (kd - 1, kh - 1, kw - 1)
        self._equiv_ok = (
            self.conv_transpose.stride == (1, 1, 1)
            and self.conv_transpose.dilation == (1, 1, 1)
            and self.conv_transpose.groups == 1
            and self.conv_transpose.padding == (0, 0, 0)
            and self.conv_transpose.output_padding == (0, 0, 0)
        )
        self._w_key = None
        self._w_conv = None
        self._use_conv3d = False
        self._checked = False
        self._stream = None

    def _get_conv_weight(self):
        w = self.conv_transpose.weight
        key = (w.data_ptr(), w._version, w.device, w.shape)
        if self._w_conv is not None and key == self._w_key:
            return self._w_conv
        with torch.no_grad():
            wc = w.detach().permute(1, 0, 2, 3, 4).flip([2, 3, 4]).contiguous(
                memory_format=torch.channels_last_3d
            )
        self._w_conv = wc
        self._w_key = key
        return wc

    def _to_nhwc(self, x):
        if x.is_contiguous(memory_format=torch.channels_last_3d):
            return x
        return self.fused.nchw_to_nhwc(x)

    def _forward_fast(self, x):
        f = self.fused
        w = self._get_conv_weight()
        bias = self.conv_transpose.bias
        gnw, gnb = self.group_norm.weight, self.group_norm.bias
        G, eps = self.groups, self.eps
        N = x.shape[0]
        nch = 4 if N >= 4 else 1
        bnd = [(i * N) // nch for i in range(nch + 1)]
        need_tp = not x.is_contiguous(memory_format=torch.channels_last_3d)
        main = torch.cuda.current_stream()
        if self._stream is None:
            self._stream = torch.cuda.Stream()
        side = self._stream

        ev0 = torch.cuda.Event()
        ev0.record(main)
        side.wait_event(ev0)

        ts = [None] * nch
        evs_t = [None] * nch
        if need_tp:
            with torch.cuda.stream(side):
                for k in range(1, nch):
                    ts[k] = f.nchw_to_nhwc(x[bnd[k]:bnd[k + 1]])
                    e = torch.cuda.Event()
                    e.record(side)
                    evs_t[k] = e
            ts[0] = f.nchw_to_nhwc(x[bnd[0]:bnd[1]])
        else:
            for k in range(nch):
                ts[k] = x[bnd[k]:bnd[k + 1]]

        cs = [None] * nch
        evs_c = [None] * nch
        y = None
        for k in range(nch):
            if evs_t[k] is not None:
                main.wait_event(evs_t[k])
            c = F.conv3d(ts[k], w, bias, (1, 1, 1), self.pad)
            cs[k] = c
            if y is None:
                y = torch.empty(
                    (N,) + tuple(c.shape[1:]),
                    device=c.device, dtype=c.dtype,
                    memory_format=torch.channels_last_3d,
                )
            e = torch.cuda.Event()
            e.record(main)
            evs_c[k] = e

        with torch.cuda.stream(side):
            for k in range(nch):
                side.wait_event(evs_c[k])
                f.fused_relu_groupnorm_nhwc(
                    cs[k], gnw, gnb, G, eps, y[bnd[k]:bnd[k + 1]]
                )
            evg = torch.cuda.Event()
            evg.record(side)
        main.wait_event(evg)
        return y

    def forward(self, x):
        if self._use_conv3d:
            return self._forward_fast(x)
        f = self.fused
        gnw, gnb = self.group_norm.weight, self.group_norm.bias
        if self._equiv_ok and not self._checked:
            self._checked = True
            try:
                xn = self._to_nhwc(x)
                wc = self._get_conv_weight()
                ref = self.conv_transpose(xn)
                alt = F.conv3d(xn, wc, self.conv_transpose.bias,
                               (1, 1, 1), self.pad)
                if torch.allclose(alt, ref, atol=1e-2, rtol=1e-2):
                    self._use_conv3d = True
                    y = torch.empty_like(
                        alt, memory_format=torch.channels_last_3d
                    )
                    f.fused_relu_groupnorm_nhwc(
                        alt, gnw, gnb, self.groups, self.eps, y
                    )
                    return y
                self._w_conv = None
                self._w_key = None
            except Exception:
                self._w_conv = None
                self._w_key = None
        xn = self._to_nhwc(x)
        c = self.conv_transpose(xn)
        y = torch.empty_like(c, memory_format=torch.channels_last_3d)
        f.fused_relu_groupnorm_nhwc(c, gnw, gnb, self.groups, self.eps, y)
        return y
```