```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

# Enable cuDNN autotuning and TF32 (same numerics mode as the reference config).
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_tf32 = True
try:
    torch.backends.cudnn.benchmark_limit = 0
except Exception:
    pass

_fused_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>

namespace {

// ===========================================================================
// Fully fused kernel: 3x3x3 conv (in_ch=3, stride=1, pad=0) + min over depth
// + per-channel bias + softmax over channels, all in one pass. The 228 MB
// conv output is never materialized: each warp computes a (32ch x 24d) tile
// of conv outputs in tensor-core accumulators (TF32 mma.m16n8k8), reduces the
// min over depth in registers, and applies the softmax epilogue directly.
// ===========================================================================
#define G_POS 6          // output positions per block (one per warp)
#define DCT 3            // 8-wide column tiles covering D_out <= 24
#define B_STRIDE (G_POS * 8 * DCT + 8)
#define TILE_D 26
#define TILE_W (G_POS + 2)

__device__ __forceinline__ unsigned f2tf32(float x) {
    unsigned r;
    asm("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(r) : "f"(x));
    return r;
}

__device__ __forceinline__ void mma1688tf32(float c[4], unsigned a0, unsigned a1,
                                            unsigned a2, unsigned a3,
                                            unsigned b0, unsigned b1) {
    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"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3])
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
__global__ __launch_bounds__(G_POS * 32) void fused_conv_min_softmax_kernel(
        const float* __restrict__ x,
        const float* __restrict__ w,
        const float* __restrict__ bias,
        float* __restrict__ out,
        int N, int C, int D_in, int H_in, int W_in,
        int D_out, int H_out, int W_out,
        int hblocks, int wblocks, int has_bias) {
    extern __shared__ unsigned smem[];
    unsigned* const As = smem;                          // 32 x 88
    unsigned* const Bs = As + 32 * 88;                  // 88 x B_STRIDE
    unsigned* const tile = Bs + 88 * B_STRIDE;          // 3*TILE_D*3*TILE_W

    const int bid = blockIdx.x;
    const int n = bid / (hblocks * wblocks);
    const int rem = bid - n * (hblocks * wblocks);
    const int h = rem / wblocks;
    const int wg = rem - h * wblocks;
    const int w0 = wg * G_POS;

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int g = lane >> 2;
    const int tg = lane & 3;
    const int p = tid >> 5;

    // ---- load zero-padded input tile (as tf32 bits) ----
    {
        const int hw0 = h - 1;
        const int ww0 = w0 - 1;
        const float* xb = x + (((size_t)n * 3) * (size_t)D_in) * (size_t)(H_in * W_in);
        const int plane = H_in * W_in;
        const int dplane = 3 * TILE_W;
        const int cplane = TILE_D * dplane;
        const int tot = 3 * cplane;
        for (int e = tid; e < tot; e += G_POS * 32) {
            const int ci = e / cplane;
            const int r1 = e - ci * cplane;
            const int dp = r1 / dplane;
            const int r2 = r1 - dp * dplane;
            const int khp = r2 / TILE_W;
            const int wp = r2 - khp * TILE_W;
            const int hh = hw0 + khp;
            const int ww = ww0 + wp;
            float v = 0.0f;
            if (dp < D_in && hh >= 0 && hh < H_in && ww >= 0 && ww < W_in) {
                v = xb[((size_t)ci * D_in + dp) * plane + hh * W_in + ww];
            }
            tile[e] = f2tf32(v);
        }
    }

    // ---- load weight matrix A (32 x 88, zero-padded, tf32 bits) ----
    for (int e = tid; e < 32 * 88; e += G_POS * 32) {
        const int c = e / 88;
        const int k = e - (c * 88);
        float v = 0.0f;
        if (c < C && k < 81) v = w[c * 81 + k];
        As[e] = f2tf32(v);
    }
    __syncthreads();

    // ---- build im2col matrix B (88 x B_STRIDE, tf32 bits) ----
    {
        const int d = tid % 24;
        const int pp = tid / 24;
        const bool dval = (d < D_out) && (pp < G_POS);
        const int dbase = d * 3 * TILE_W + pp;
        for (int k = 0; k < 88; ++k) {
            unsigned val = 0u;
            if (k < 81 && dval) {
                const int ci = k / 27;
                const int r1 = k - ci * 27;
                const int kd = r1 / 9;
                const int r2 = r1 - kd * 9;
                const int kh = r2 / 3;
                const int kw = r2 - kh * 3;
                val = tile[ci * (TILE_D * 3 * TILE_W) + dbase + kd * (3 * TILE_W)
                           + kh * TILE_W + kw];
            }
            Bs[k * B_STRIDE + tid] = val;
        }
    }
    __syncthreads();

    // ---- tensor-core contraction ----
    float acc[DCT][2][4];
    #pragma unroll
    for (int ct = 0; ct < DCT; ++ct)
        #pragma unroll
        for (int rt = 0; rt < 2; ++rt)
            #pragma unroll
            for (int i = 0; i < 4; ++i) acc[ct][rt][i] = 0.0f;

    const int colbase = p * 24 + g;
    #pragma unroll
    for (int ks = 0; ks < 11; ++ks) {
        const int k0 = ks * 8 + tg;
        const unsigned a00 = As[(g) * 88 + k0];
        const unsigned a01 = As[(g) * 88 + k0 + 4];
        const unsigned a02 = As[(g + 8) * 88 + k0];
        const unsigned a03 = As[(g + 8) * 88 + k0 + 4];
        const unsigned a10 = As[(g + 16) * 88 + k0];
        const unsigned a11 = As[(g + 16) * 88 + k0 + 4];
        const unsigned a12 = As[(g + 24) * 88 + k0];
        const unsigned a13 = As[(g + 24) * 88 + k0 + 4];
        const unsigned* bptr = Bs + k0 * B_STRIDE + colbase;
        #pragma unroll
        for (int ct = 0; ct < DCT; ++ct) {
            const unsigned b0 = bptr[ct * 8];
            const unsigned b1 = bptr[4 * B_STRIDE + ct * 8];
            mma1688tf32(acc[ct][0], a00, a01, a02, a03, b0, b1);
            mma1688tf32(acc[ct][1], a10, a11, a12, a13, b0, b1);
        }
    }

    // ---- min over depth (per channel), excluding padded d ----
    const int dd = 2 * tg;
    float m0 = CUDART_INF_F, m1 = CUDART_INF_F, m2 = CUDART_INF_F, m3 = CUDART_INF_F;
    #pragma unroll
    for (int ct = 0; ct < DCT; ++ct) {
        const bool v0 = (ct * 8 + dd) < D_out;
        const bool v1 = (ct * 8 + dd + 1) < D_out;
        if (v0) {
            m0 = fminf(m0, acc[ct][0][0]);
            m1 = fminf(m1, acc[ct][0][1]);
            m2 = fminf(m2, acc[ct][0][2]);
            m3 = fminf(m3, acc[ct][0][3]);
        }
        if (v1) {
            m0 = fminf(m0, acc[ct][1][0]);
            m1 = fminf(m1, acc[ct][1][1]);
            m2 = fminf(m2, acc[ct][1][2]);
            m3 = fminf(m3, acc[ct][1][3]);
        }
    }
    // reduce across the 4 threads of each quad (all d columns)
    m0 = fminf(m0, __shfl_xor_sync(0xffffffffu, m0, 1));
    m1 = fminf(m1, __shfl_xor_sync(0xffffffffu, m1, 1));
    m2 = fminf(m2, __shfl_xor_sync(0xffffffffu, m2, 1));
    m3 = fminf(m3, __shfl_xor_sync(0xffffffffu, m3, 1));
    m0 = fminf(m0, __shfl_xor_sync(0xffffffffu, m0, 2));
    m1 = fminf(m1, __shfl_xor_sync(0xffffffffu, m1, 2));
    m2 = fminf(m2, __shfl_xor_sync(0xffffffffu, m2, 2));
    m3 = fminf(m3, __shfl_xor_sync(0xffffffffu, m3, 2));

    // ---- fold bias (exact: min(y+b) == min(y)+b), mask invalid channels ----
    const bool hv0 = g < C;
    const bool hv1 = (g + 8) < C;
    const bool hv2 = (g + 16) < C;
    float v0 = hv0 ? m0 : -CUDART_INF_F;
    float v1 = hv1 ? m1 : -CUDART_INF_F;
    float v2 = hv2 ? m2 : -CUDART_INF_F;
    (void)m3;  // channels >= 24: never valid
    if (has_bias) {
        if (hv0) v0 += __ldg(bias + g);
        if (hv1) v1 += __ldg(bias + g + 8);
        if (hv2) v2 += __ldg(bias + g + 16);
    }

    // ---- softmax over channels (warp reduction; quads hold replicas) ----
    float mx = fmaxf(v0, fmaxf(v1, v2));
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, off));
    float e0 = hv0 ? __expf(v0 - mx) : 0.0f;
    float e1 = hv1 ? __expf(v1 - mx) : 0.0f;
    float e2 = hv2 ? __expf(v2 - mx) : 0.0f;
    float s = e0 + e1 + e2;
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        s += __shfl_xor_sync(0xffffffffu, s, off);
    const float inv = 1.0f / s;

    if (tg == 0) {
        float* ob = out + (((size_t)n * C) * H_out + h) * W_out + (w0 + p);
        const int plane = H_out * W_out;
        if (hv0) ob[(size_t)g * plane] = e0 * inv;
        if (hv1) ob[(size_t)(g + 8) * plane] = e1 * inv;
        if (hv2) ob[(size_t)(g + 16) * plane] = e2 * inv;
    }
}
#endif  // __CUDA_ARCH__ >= 800

// ---------------------------------------------------------------------------
// NCHW -> NHWC transpose (fallback path). Fully vectorized.
// ---------------------------------------------------------------------------
template <int C>
__global__ __launch_bounds__(256) void nchw_to_nhwc_kernel(
        const float* __restrict__ in, float* __restrict__ out, int DS) {
    const int n = blockIdx.x;
    const int s0 = (blockIdx.y * blockDim.x + threadIdx.x) * 4;
    if (s0 >= DS) return;
    const float* in_n = in + (size_t)n * (size_t)C * (size_t)DS;
    float* out_n = out + (size_t)n * (size_t)DS * (size_t)C;
    if (s0 + 4 <= DS) {
        float v[C][4];
        if ((DS & 3) == 0) {
            #pragma unroll
            for (int c = 0; c < C; ++c) {
                const float4 t = *reinterpret_cast<const float4*>(
                    in_n + (size_t)c * (size_t)DS + s0);
                v[c][0] = t.x; v[c][1] = t.y; v[c][2] = t.z; v[c][3] = t.w;
            }
        } else {
            #pragma unroll
            for (int c = 0; c < C; ++c)
                #pragma unroll
                for (int j = 0; j < 4; ++j)
                    v[c][j] = in_n[(size_t)c * (size_t)DS + s0 + j];
        }
        if (((DS * C) & 3) == 0) {
            float4* o = reinterpret_cast<float4*>(out_n + (size_t)s0 * C);
            #pragma unroll
            for (int k = 0; k < C; ++k) {
                float4 t;
                t.x = v[(4 * k + 0) % C][(4 * k + 0) / C];
                t.y = v[(4 * k + 1) % C][(4 * k + 1) / C];
                t.z = v[(4 * k + 2) % C][(4 * k + 2) / C];
                t.w = v[(4 * k + 3) % C][(4 * k + 3) / C];
                o[k] = t;
            }
        } else {
            #pragma unroll
            for (int j = 0; j < 4; ++j)
                #pragma unroll
                for (int c = 0; c < C; ++c)
                    out_n[(size_t)(s0 + j) * C + c] = v[c][j];
        }
    } else {
        for (int s = s0; s < DS; ++s)
            #pragma unroll
            for (int c = 0; c < C; ++c)
                out_n[(size_t)s * C + c] = in_n[(size_t)c * (size_t)DS + s];
    }
}

// ---------------------------------------------------------------------------
// NHWC fused min-over-depth + bias + softmax-over-channels (fallback path).
// ---------------------------------------------------------------------------
template <int C>
__global__ __launch_bounds__(256) void fused_min_softmax_nhwc_kernel(
        const float* __restrict__ in, const float* __restrict__ bias,
        float* __restrict__ out, int N, int D, int HW) {
    constexpr int C4 = C / 4;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * HW) return;
    const int n = idx / HW;
    const int s = idx - n * HW;
    const float4* base = reinterpret_cast<const float4*>(in)
                       + ((size_t)n * (size_t)D * (size_t)HW + (size_t)s) * C4;
    const size_t dstride = (size_t)HW * C4;
    float m[C];
    {
        const float4* p = base;
        #pragma unroll
        for (int i = 0; i < C4; ++i) {
            const float4 v = __ldcs(p + i);
            m[4*i+0] = v.x; m[4*i+1] = v.y; m[4*i+2] = v.z; m[4*i+3] = v.w;
        }
    }
    for (int d = 1; d < D; ++d) {
        const float4* p = base + (size_t)d * dstride;
        #pragma unroll
        for (int i = 0; i < C4; ++i) {
            const float4 v = __ldcs(p + i);
            m[4*i+0] = fminf(m[4*i+0], v.x);
            m[4*i+1] = fminf(m[4*i+1], v.y);
            m[4*i+2] = fminf(m[4*i+2], v.z);
            m[4*i+3] = fminf(m[4*i+3], v.w);
        }
    }
    if (bias != nullptr) {
        #pragma unroll
        for (int c = 0; c < C; ++c) m[c] += __ldg(bias + c);
    }
    float maxv = m[0];
    #pragma unroll
    for (int c = 1; c < C; ++c) maxv = fmaxf(maxv, m[c]);
    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < C; ++c) { m[c] = expf(m[c] - maxv); sum += m[c]; }
    const float inv = 1.0f / sum;
    float* op = out + (size_t)n * C * HW + s;
    #pragma unroll
    for (int c = 0; c < C; ++c) op[(size_t)c * HW] = m[c] * inv;
}

__global__ void fused_min_softmax_nhwc_generic_kernel(
        const float* __restrict__ in, const float* __restrict__ bias,
        float* __restrict__ out, int N, int C, int D, int HW) {
    extern __shared__ float smem[];
    const int n = blockIdx.x / HW;
    const int s = blockIdx.x - n * HW;
    const int tid = threadIdx.x;
    const int nt = blockDim.x;
    const float* base = in + ((size_t)n * (size_t)D * (size_t)HW + (size_t)s) * (size_t)C;
    const size_t dstride = (size_t)HW * (size_t)C;
    for (int c = tid; c < C; c += nt) {
        float m = base[c];
        for (int d = 1; d < D; ++d) m = fminf(m, base[(size_t)d * dstride + c]);
        if (bias != nullptr) m += __ldg(bias + c);
        smem[c] = m;
    }
    __syncthreads();
    float* red = smem + C;
    float lmax = -INFINITY;
    for (int c = tid; c < C; c += nt) lmax = fmaxf(lmax, smem[c]);
    red[tid] = lmax;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] = fmaxf(red[tid], red[tid + st]);
        __syncthreads();
    }
    const float maxv = red[0];
    __syncthreads();
    float lsum = 0.0f;
    for (int c = tid; c < C; c += nt) {
        const float e = expf(smem[c] - maxv);
        smem[c] = e;
        lsum += e;
    }
    red[tid] = lsum;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] += red[tid + st];
        __syncthreads();
    }
    const float inv = 1.0f / red[0];
    float* op = out + (size_t)n * (size_t)C * (size_t)HW + s;
    for (int c = tid; c < C; c += nt) op[(size_t)c * HW] = smem[c] * inv;
}

template <int C>
__global__ void fused_min_softmax_nchw_kernel(
        const float* __restrict__ in, const float* __restrict__ bias,
        float* __restrict__ out, int N, int D, int HW) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * HW) return;
    const int n = idx / HW;
    const int s = idx - n * HW;
    const float* in_n = in + (size_t)n * C * D * HW + s;
    float vals[C];
    #pragma unroll
    for (int c = 0; c < C; ++c) {
        const float* p = in_n + (size_t)c * D * HW;
        float m = p[0];
        #pragma unroll 4
        for (int d = 1; d < D; ++d) m = fminf(m, p[(size_t)d * HW]);
        if (bias != nullptr) m += __ldg(bias + c);
        vals[c] = m;
    }
    float maxv = vals[0];
    #pragma unroll
    for (int c = 1; c < C; ++c) maxv = fmaxf(maxv, vals[c]);
    float sum = 0.0f;
    #pragma unroll
    for (int c = 0; c < C; ++c) { vals[c] = expf(vals[c] - maxv); sum += vals[c]; }
    const float inv = 1.0f / sum;
    float* out_n = out + (size_t)n * C * HW + s;
    #pragma unroll
    for (int c = 0; c < C; ++c) out_n[(size_t)c * HW] = vals[c] * inv;
}

__global__ void fused_min_softmax_nchw_generic_kernel(
        const float* __restrict__ in, const float* __restrict__ bias,
        float* __restrict__ out, int N, int C, int D, int HW) {
    extern __shared__ float smem[];
    const int n = blockIdx.x / HW;
    const int s = blockIdx.x - n * HW;
    const int tid = threadIdx.x;
    const int nt = blockDim.x;
    const float* in_n = in + (size_t)n * C * D * HW + s;
    for (int c = tid; c < C; c += nt) {
        const float* p = in_n + (size_t)c * D * HW;
        float m = p[0];
        for (int d = 1; d < D; ++d) m = fminf(m, p[(size_t)d * HW]);
        if (bias != nullptr) m += __ldg(bias + c);
        smem[c] = m;
    }
    __syncthreads();
    float* red = smem + C;
    float lmax = -INFINITY;
    for (int c = tid; c < C; c += nt) lmax = fmaxf(lmax, smem[c]);
    red[tid] = lmax;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] = fmaxf(red[tid], red[tid + st]);
        __syncthreads();
    }
    const float maxv = red[0];
    __syncthreads();
    float lsum = 0.0f;
    for (int c = tid; c < C; c += nt) {
        const float e = expf(smem[c] - maxv);
        smem[c] = e;
        lsum += e;
    }
    red[tid] = lsum;
    __syncthreads();
    for (int st = nt >> 1; st > 0; st >>= 1) {
        if (tid < st) red[tid] += red[tid + st];
        __syncthreads();
    }
    const float inv = 1.0f / red[0];
    float* out_n = out + (size_t)n * C * HW + s;
    for (int c = tid; c < C; c += nt) out_n[(size_t)c * HW] = smem[c] * inv;
}

const float* get_bias_ptr(const c10::optional<torch::Tensor>& bias_opt, int C) {
    if (!bias_opt.has_value() || !bias_opt->defined()) return nullptr;
    const auto& b = *bias_opt;
    TORCH_CHECK(b.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(b.scalar_type() == at::kFloat, "bias must be float32");
    TORCH_CHECK(b.numel() >= C, "bias must have at least C elements");
    TORCH_CHECK(b.is_contiguous(), "bias must be contiguous");
    return b.data_ptr<float>();
}

}  // namespace

torch::Tensor fused_conv_min_softmax_cuda(
        torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> bias) {
#if defined(__CUDA_ARCH__) || defined(__CUDACC__)
    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(w.scalar_type() == at::kFloat, "w must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(w.is_contiguous(), "w must be contiguous");
    TORCH_CHECK(x.dim() == 5 && w.dim() == 5, "x and w must be 5D");
    TORCH_CHECK(w.size(1) == 3 && w.size(2) == 3 && w.size(3) == 3 && w.size(4) == 3,
                "fused kernel requires in_channels=3, kernel_size=3");

    const int N = (int)x.size(0);
    const int C = (int)w.size(0);
    const int D_in = (int)x.size(2);
    const int H_in = (int)x.size(3);
    const int W_in = (int)x.size(4);
    const int D_out = D_in - 2;
    const int H_out = H_in - 2;
    const int W_out = W_in - 2;
    TORCH_CHECK(C <= 32, "fused kernel requires out_channels <= 32");
    TORCH_CHECK(D_out >= 1 && D_out <= 24, "fused kernel requires 1 <= D_out <= 24");

    auto out = at::empty({N, C, H_out, W_out}, x.options());
    if (N == 0 || H_out == 0 || W_out == 0 || C == 0) return out;

    const float* bias_ptr = get_bias_ptr(bias, C);
    const int has_bias = bias_ptr != nullptr ? 1 : 0;

    const int wblocks = (W_out + G_POS - 1) / G_POS;
    const int hblocks = H_out;
    const long long grid = (long long)N * hblocks * wblocks;
    TORCH_CHECK(grid <= 0x7fffffff, "grid too large");

    const size_t shmem = (size_t)(32 * 88 + 88 * B_STRIDE + 3 * TILE_D * 3 * TILE_W)
                         * sizeof(unsigned);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute(fused_conv_min_softmax_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             (int)(128 * 1024));
        configured = true;
    }
    fused_conv_min_softmax_kernel<<<(int)grid, G_POS * 32, shmem, stream>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr,
        out.data_ptr<float>(), N, C, D_in, H_in, W_in,
        D_out, H_out, W_out, hblocks, wblocks, has_bias);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
#else
    TORCH_CHECK(false, "fused kernel unavailable");
#endif
}

torch::Tensor nchw_to_nhwc_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 5, "input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(x.is_contiguous(), "input must be NCHW contiguous");
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    TORCH_CHECK(C >= 1 && C <= 4, "channel count must be in [1, 4]");
    auto out = at::empty({N, C, D, H, W},
                         x.options().memory_format(at::MemoryFormat::ChannelsLast3d));
    const long long DS = (long long)D * H * W;
    if (DS == 0 || N == 0) return out;
    const long long blocks_y = (DS + 256 * 4 - 1) / (256 * 4);
    TORCH_CHECK(blocks_y <= 65535, "spatial size too large for transpose kernel");
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const dim3 grid(N, (unsigned)blocks_y);
    const float* in_ptr = x.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();
    switch (C) {
        case 1: nchw_to_nhwc_kernel<1><<<grid, 256, 0, stream>>>(in_ptr, out_ptr, (int)DS); break;
        case 2: nchw_to_nhwc_kernel<2><<<grid, 256, 0, stream>>>(in_ptr, out_ptr, (int)DS); break;
        case 3: nchw_to_nhwc_kernel<3><<<grid, 256, 0, stream>>>(in_ptr, out_ptr, (int)DS); break;
        default: nchw_to_nhwc_kernel<4><<<grid, 256, 0, stream>>>(in_ptr, out_ptr, (int)DS); break;
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}

torch::Tensor fused_min_softmax_nhwc_cuda(torch::Tensor input,
                                          c10::optional<torch::Tensor> bias) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 5, "input must be 5D (N, C, D, H, W)");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
    TORCH_CHECK(input.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "input must be channels_last_3d contiguous");
    const int N = (int)input.size(0);
    const int C = (int)input.size(1);
    const int D = (int)input.size(2);
    const int H = (int)input.size(3);
    const int W = (int)input.size(4);
    TORCH_CHECK(D >= 1, "depth must be >= 1");
    auto out = at::empty({N, C, H, W}, input.options());
    const int HW = H * W;
    const long long total = (long long)N * HW;
    if (total == 0 || C == 0) return out;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const float* in_ptr = input.data_ptr<float>();
    const float* bias_ptr = get_bias_ptr(bias, C);
    float* out_ptr = out.data_ptr<float>();
    const bool aligned16 = ((reinterpret_cast<uintptr_t>(in_ptr)) & 0xF) == 0;
    const int threads = 256;
    const int blocks = (int)((total + threads - 1) / threads);
    bool launched = false;
    if (aligned16 && (C % 4) == 0) {
        #define FUSED_NHWC_CASE(CC)                                                \
            case CC:                                                               \
                fused_min_softmax_nhwc_kernel<CC><<<blocks, threads, 0, stream>>>( \
                    in_ptr, bias_ptr, out_ptr, N, D, HW);                          \
                launched = true;                                                   \
                break;
        switch (C) {
            FUSED_NHWC_CASE(4) FUSED_NHWC_CASE(8) FUSED_NHWC_CASE(12)
            FUSED_NHWC_CASE(16) F