Looking at the profile, the single fused kernel is correct at ~398 µs. The dominant cost is the dense FP32 FMA work for the conv-transpose (~2.7 G FMA). Since TF32 tensor-core math is explicitly allowed (the reference uses it), I'll rewrite the kernel around `mma.sync.m16n8k8.tf32`, using the sub-convolution (parity) decomposition of the stride-2 conv-transpose, which also maps the 2×2×2 max-pool to a simple elementwise max across parity fragments. The softmax/swish/channel-max epilogue is done with in-register shuffle reductions, so nothing intermediate touches global memory. A safe generic fallback remains for non-target configs.

```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
from torch.utils.cpp_extension import load_inline

# Fully fused TF32 tensor-core kernel:
#   ConvTranspose3d (k=3, s=2, p=1, op=1) -> MaxPool3d (2,2,0) -> Softmax(dim=1)
#   -> per-channel subtract -> Swish -> Max over channels
#
# Design notes:
#  * The stride-2 conv-transpose is decomposed into 8 stride-1 "parity"
#    sub-convolutions on the input grid. The 2x2x2 max-pool window corresponds
#    exactly to the 8 parity classes at one grid cell, so pooling becomes an
#    elementwise max across the 8 accumulator fragments, and the pooled spatial
#    size equals the input spatial size.
#  * Each warp computes a 16(co) x 8(cells) output tile per parity with
#    mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 (TF32 tensor-core math,
#    same precision class as the TF32 reference). All 14 weight fragments are
#    preloaded into registers once per warp and reused across every tile.
#  * B-fragments (gathered inputs) are loaded straight from global memory with
#    per-thread address computation -- no shared-memory staging, no barriers.
#  * Softmax denominator exponentials are reused for the probabilities; the
#    subtract + swish + channel-max epilogue uses __shfl_xor reductions within
#    8-lane groups, so no intermediate tensor is ever materialized.

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cfloat>

#define CI 3
#define CO 16

#define PAR_KD(p) (((p) >> 2) & 1)
#define PAR_KH(p) (((p) >> 1) & 1)
#define PAR_KW(p) ((p) & 1)

__global__ __launch_bounds__(128) void fused_convT_pool_softmax_swish_cmax_mma_kernel(
    const float* __restrict__ x,        // [B, CI, D, H, W]
    const float* __restrict__ weight,   // [CI, CO, 3, 3, 3]
    const float* __restrict__ bias,     // [CO]
    const float* __restrict__ sub,      // [CO]
    float* __restrict__ out,            // [B, Dp=D, Hp=H, Wp=W]
    const int D, const int H, const int W,
    const int nwt,                      // W/8 tiles per (b,d,h)
    const int tiles_per_b,              // D*H*nwt
    const int nblocks, const int T)     // tiles per warp chunk
{
    __shared__ int   s_ci[8][24];       // input channel per (parity, k)
    __shared__ int   s_tap[8][24];      // flattened 3x3x3 tap index per (parity, k)
    __shared__ short s_off[8][24][3];   // input offsets (d,h,w) per (parity, k)

    // Build the parity -> (ci, tap, offsets) tables (single warp, once).
    if (threadIdx.x < 32) {
        const int lane = threadIdx.x;
        for (int idx = lane; idx < 8 * 24; idx += 32) {
            const int p = idx >> 5;                 // 24 rounded to 32 per parity
            const int k = idx & 31;
            int ci = 0, tap = 0, od = 0, oh = 0, ow = 0;
            if (k < 24) {
                const int kd = PAR_KD(p), kh = PAR_KH(p), kw = PAR_KW(p);
                const int TW = kw + 1, TH = kh + 1, TD = kd + 1;
                const int TPS = TD * TH * TW;
                const int Kp  = CI * TPS;
                if (k < Kp) {
                    ci = k / TPS;
                    int r = k - ci * TPS;
                    const int di = r / (TH * TW);
                    r -= di * (TH * TW);
                    const int hi = r / TW;
                    const int wi = r - hi * TW;
                    const int tdk = kd ? (di == 0 ? 0 : 2) : 1;
                    const int thk = kh ? (hi == 0 ? 0 : 2) : 1;
                    const int twk = kw ? (wi == 0 ? 0 : 2) : 1;
                    od = kd ? (1 - di) : 0;
                    oh = kh ? (1 - hi) : 0;
                    ow = kw ? (1 - wi) : 0;
                    tap = tdk * 9 + thk * 3 + twk;
                }
            }
            s_ci[p][k] = ci;
            s_tap[p][k] = tap;
            s_off[p][k][0] = (short)od;
            s_off[p][k][1] = (short)oh;
            s_off[p][k][2] = (short)ow;
        }
    }
    __syncthreads();

    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    const int g  = lane >> 2;    // mma group id (0..7)
    const int tg = lane & 3;     // thread id in group (0..3)

    const int HW  = H * W;
    const int DHW = D * HW;

    // ---- Preload all 14 A (weight) fragments into registers. ----
    float a0[14], a1[14], a2[14], a3[14];
    {
        int f = 0;
        #pragma unroll
        for (int p = 0; p < 8; ++p) {
            const int TW = PAR_KW(p) + 1, TH = PAR_KH(p) + 1, TD = PAR_KD(p) + 1;
            const int Kp = CI * TD * TH * TW;
            const int steps = (Kp + 7) >> 3;
            for (int st = 0; st < steps; ++st, ++f) {
                const int ka = 8 * st + tg, kb = ka + 4;
                const float* wg = weight + g * 27;
                a0[f] = (ka < Kp) ? __ldg(wg + s_ci[p][ka] * 432 + s_tap[p][ka]) : 0.0f;
                a1[f] = (ka < Kp) ? __ldg(wg + 216 + s_ci[p][ka] * 432 + s_tap[p][ka]) : 0.0f;
                a2[f] = (kb < Kp) ? __ldg(wg + s_ci[p][kb] * 432 + s_tap[p][kb]) : 0.0f;
                a3[f] = (kb < Kp) ? __ldg(wg + 216 + s_ci[p][kb] * 432 + s_tap[p][kb]) : 0.0f;
            }
        }
    }

    const float b0c = __ldg(bias + g);
    const float b8c = __ldg(bias + g + 8);
    const float s0c = __ldg(sub + g);
    const float s8c = __ldg(sub + g + 8);

    // ---- Tile loop: each warp handles a contiguous chunk of T tiles. ----
    for (int blk = blockIdx.x; blk < nblocks; blk += gridDim.x) {
        const int base = (blk * 4 + wid) * T;
        if (base >= tiles_per_b * gridDim.z) continue;
        const int b   = blockIdx.z;
        const int rem = tiles_per_b - base;
        const int tcount = rem < T ? rem : T;

        const float* xb = x + (long long)b * CI * DHW;
        float* outb = out + (long long)b * D * HW;

        for (int t = 0; t < tcount; ++t) {
            const int tile = base + t;
            const int wt = tile % nwt;
            const int dh = tile / nwt;
            const int h = dh % H;
            const int d = dh / H;
            const int w0 = wt * 8;

            // Accumulators: [parity][c0..c3]
            float acc[8][4];
            #pragma unroll
            for (int p = 0; p < 8; ++p) {
                acc[p][0] = 0.0f; acc[p][1] = 0.0f;
                acc[p][2] = 0.0f; acc[p][3] = 0.0f;
            }

            int f = 0;
            #pragma unroll
            for (int p = 0; p < 8; ++p) {
                const int TW = PAR_KW(p) + 1, TH = PAR_KH(p) + 1, TD = PAR_KD(p) + 1;
                const int Kp = CI * TD * TH * TW;
                const int steps = (Kp + 7) >> 3;
                #pragma unroll
                for (int st = 0; st < 3; ++st) {
                    if (st < steps) {
                        const int ka = 8 * st + tg, kb = ka + 4;
                        // B fragment: row k, column n=g -> input w index w0 + g + offw
                        const int wa = w0 + g + s_off[p][ka < 24 ? ka : 0][2];
                        const int wb = w0 + g + s_off[p][kb < 24 ? kb : 0][2];
                        const bool va = (ka < Kp) && ((unsigned)wa < (unsigned)W);
                        const bool vb = (kb < Kp) && ((unsigned)wb < (unsigned)W);
                        const int da = d + s_off[p][ka < 24 ? ka : 0][0];
                        const int ha = h + s_off[p][ka < 24 ? ka : 0][1];
                        const int db = d + s_off[p][kb < 24 ? kb : 0][0];
                        const int hb = h + s_off[p][kb < 24 ? kb : 0][1];
                        const bool vda = va && ((unsigned)da < (unsigned)D)
                                       && ((unsigned)ha < (unsigned)H);
                        const bool vdb = vb && ((unsigned)db < (unsigned)D)
                                       && ((unsigned)hb < (unsigned)H);
                        const long long offa = (long long)s_ci[p][ka < 24 ? ka : 0] * DHW
                                             + (long long)da * HW + (long long)ha * W + wa;
                        const long long offb = (long long)s_ci[p][kb < 24 ? kb : 0] * DHW
                                             + (long long)db * HW + (long long)hb * W + wb;
                        const float bv0 = vda ? __ldg(xb + offa) : 0.0f;
                        const float bv1 = vdb ? __ldg(xb + offb) : 0.0f;
                        asm volatile(
                            "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
                            "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
                            : "+f"(acc[p][0]), "+f"(acc[p][1]),
                              "+f"(acc[p][2]), "+f"(acc[p][3])
                            : "r"(__float_as_uint(a0[f])), "r"(__float_as_uint(a1[f])),
                              "r"(__float_as_uint(a2[f])), "r"(__float_as_uint(a3[f])),
                              "r"(__float_as_uint(bv0)), "r"(__float_as_uint(bv1)));
                        ++f;
                    }
                }
            }

            // MaxPool3d across the 8 parity classes, then add the conv bias.
            float r0 = acc[0][0], r1 = acc[0][1], r2 = acc[0][2], r3 = acc[0][3];
            #pragma unroll
            for (int p = 1; p < 8; ++p) {
                r0 = fmaxf(r0, acc[p][0]);
                r1 = fmaxf(r1, acc[p][1]);
                r2 = fmaxf(r2, acc[p][2]);
                r3 = fmaxf(r3, acc[p][3]);
            }
            r0 += b0c; r1 += b0c; r2 += b8c; r3 += b8c;

            // Softmax across the 16 channels (rows) of each column.
            float m01 = fmaxf(r0, r1);
            float m23 = fmaxf(r2, r3);
            m01 = fmaxf(m01, __shfl_xor_sync(0xffffffffu, m01, 4));
            m23 = fmaxf(m23, __shfl_xor_sync(0xffffffffu, m23, 4));
            m01 = fmaxf(m01, __shfl_xor_sync(0xffffffffu, m01, 8));
            m23 = fmaxf(m23, __shfl_xor_sync(0xffffffffu, m23, 8));
            m01 = fmaxf(m01, __shfl_xor_sync(0xffffffffu, m01, 16));
            m23 = fmaxf(m23, __shfl_xor_sync(0xffffffffu, m23, 16));

            float e0 = __expf(r0 - m01), e1 = __expf(r1 - m01);
            float e2 = __expf(r2 - m23), e3 = __expf(r3 - m23);
            float sA = e0 + e1, sB = e2 + e3;
            sA += __shfl_xor_sync(0xffffffffu, sA, 4);
            sB += __shfl_xor_sync(0xffffffffu, sB, 4);
            sA += __shfl_xor_sync(0xffffffffu, sA, 8);
            sB += __shfl_xor_sync(0xffffffffu, sB, 8);
            sA += __shfl_xor_sync(0xffffffffu, sA, 16);
            sB += __shfl_xor_sync(0xffffffffu, sB, 16);
            const float invA = __fdividef(1.0f, sA);
            const float invB = __fdividef(1.0f, sB);

            // Subtract per-channel constant, swish (y * sigmoid(y)), max over channels.
            float y0 = fmaf(e0, invA, -s0c);
            float y1 = fmaf(e1, invA, -s0c);
            float y2 = fmaf(e2, invB, -s8c);
            float y3 = fmaf(e3, invB, -s8c);
            float z0 = __fdividef(y0, 1.0f + __expf(-y0));
            float z1 = __fdividef(y1, 1.0f + __expf(-y1));
            float z2 = __fdividef(y2, 1.0f + __expf(-y2));
            float z3 = __fdividef(y3, 1.0f + __expf(-y3));
            float zA = fmaxf(z0, z1);
            float zB = fmaxf(z2, z3);
            zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 4));
            zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 4));
            zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 8));
            zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 8));
            zA = fmaxf(zA, __shfl_xor_sync(0xffffffffu, zA, 16));
            zB = fmaxf(zB, __shfl_xor_sync(0xffffffffu, zB, 16));

            if (g == 0) {
                outb[d * HW + h * W + w0 + 2 * tg]     = zA;
                outb[d * HW + h * W + w0 + 2 * tg + 1] = zB;
            }
        }
    }
}

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                            torch::Tensor sub, int64_t Dp, int64_t Hp, int64_t Wp)
{
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda() && sub.is_cuda(),
                "all tensors must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.size(0) == CI && weight.size(1) == CO && weight.size(2) == 3,
                "unexpected weight shape");

    const int64_t B = x.size(0);
    const int64_t D = x.size(2);
    const int64_t H = x.size(3);
    const int64_t W = x.size(4);

    auto out = at::empty({B, Dp, Hp, Wp}, x.options());
    const long long total = (long long)B * Dp * Hp * Wp;
    if (total == 0) return out;

    const int nwt = (int)((W + 7) / 8);
    const int tiles_per_b = (int)(D * H) * nwt;

    // Choose the per-warp chunk so that all SMs are saturated.
    int dev = 0;
    cudaGetDevice(&dev);
    static int nsm = 0;
    if (nsm == 0) cudaDeviceGetAttribute(&nsm, cudaDevAttrMultiProcessorCount, dev);
    const int want_blocks = nsm * 16;
    int T = (tiles_per_b + want_blocks - 1) / want_blocks;
    if (T < 1) T = 1;
    const int nblocks = (tiles_per_b + 4 * T - 1) / (4 * T);
    const int gx = nblocks > 65535 ? 65535 : nblocks;

    const dim3 grid((unsigned)gx, 1u, (unsigned)B);
    const dim3 block(128, 1, 1);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_convT_pool_softmax_swish_cmax_mma_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        sub.data_ptr<float>(), out.data_ptr<float>(),
        (int)D, (int)H, (int)W, nwt, tiles_per_b, nblocks, T);
    return out;
}
"""

_fused_cpp_src = (
    "torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, "
    "torch::Tensor sub, int64_t Dp, int64_t Hp, int64_t Wp);"
)

_fused_ext = load_inline(
    name="fused_convT3d_pool_softmax_sub_swish_cmax_mma_v4",
    cpp_sources=_fused_cpp_src,
    cuda_sources=_fused_cuda_src,
    functions=["fused_forward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "-arch=sm_80"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Uses a single fully-fused TF32 tensor-core CUDA
    kernel that performs ConvTranspose3d + MaxPool3d + Softmax + Subtract +
    Swish + channel-wise Max without materializing any intermediate tensor in
    global memory. Falls back to the original PyTorch ops for unsupported
    configurations.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding,
                 output_padding, pool_kernel_size, pool_stride, pool_padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding,
                                                 output_padding=output_padding)
        self.max_pool = nn.MaxPool3d(kernel_size=pool_kernel_size, stride=pool_stride,
                                     padding=pool_padding)
        self.subtract = nn.Parameter(torch.randn(out_channels))

        def _i(v):
            return v[0] if isinstance(v, (tuple, list)) else v

        cfg = (in_channels, out_channels, _i(kernel_size), _i(stride), _i(padding),
               _i(output_padding), _i(pool_kernel_size), _i(pool_stride), _i(pool_padding))
        self._fused_ok = (cfg == (3, 16, 3, 2, 1, 1, 2, 2, 0))
        self._zero_bias = None

    def forward(self, x):
        if self._fused_ok and x.is_cuda and x.dtype == torch.float32:
            B, D, H, W = x.shape[0], x.shape[2], x.shape[3], x.shape[4]
            # conv-transpose output size: (n-1)*2 - 2 + 3 + 1 = 2n
            Dc = (D - 1) * 2 - 2 + 3 + 1
            Hc = (H - 1) * 2 - 2 + 3 + 1
            Wc = (W - 1) * 2 - 2 + 3 + 1
            # max-pool (k=2, s=2, p=0) output size
            Dp = (Dc - 2) // 2 + 1
            Hp = (Hc - 2) // 2 + 1
            Wp = (Wc - 2) // 2 + 1
            if (x.numel() < 2**31 and B <= 65535 and Dp == D and Hp == H and Wp == W
                    and B * D * H * W < 2**31):
                bias = self.conv_transpose.bias
                if bias is None:
                    if self._zero_bias is None or self._zero_bias.device != x.device:
                        self._zero_bias = torch.zeros(16, device=x.device, dtype=torch.float32)
                    bias = self._zero_bias
                return _fused_ext.fused_forward(
                    x.contiguous(),
                    self.conv_transpose.weight.contiguous(),
                    bias.contiguous(),
                    self.subtract.contiguous(),
                    Dp, Hp, Wp)

        # Generic fallback: original PyTorch operator sequence.
        x = self.conv_transpose(x)
        x = self.max_pool(x)
        x = torch.softmax(x, dim=1)
        x = x - self.subtract.view(1, -1, 1, 1, 1)
        x = torch.sigmoid(x) * x
        x = torch.max(x, dim=1)[0]
        return x
```