You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # 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


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # 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

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

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

class Model(nn.Module):
    """
    Simple model that performs a 3D convolution, applies minimum operation along a specific dimension, 
    and then applies softmax.
    """
    def __init__(self, in_channels, out_channels, kernel_size, dim):
        super(Model, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.dim = dim

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W)
        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, H, W)
        """
        x = self.conv(x)
        x = torch.min(x, dim=self.dim)[0]  # Apply minimum along the specified dimension
        x = torch.softmax(x, dim=1)  # Apply softmax along the channel dimension
        return x

batch_size = 128
in_channels = 3
out_channels = 24  # Increased output channels
D, H, W = 24, 32, 32  # Increased depth
kernel_size = 3
dim = 2  # Dimension along which to apply minimum operation (e.g., depth)

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, dim]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 6 of 6 for this task. Your previous solution (turn 5) was:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA kernels
#
# 1) conv3d_min_softmax_smem_kernel  (NEW)
#    Fully fused, single pass, zero intermediates:
#       Conv3d (IC=3, K=3, COUT=24, stride 1, no padding)
#       -> min over depth (dim=2) -> softmax over channels (dim=1)
#    Implicit GEMM on TF32 tensor cores (mma.m16n8k8.tf32; TF32 is allowed by
#    the task precision rule - the PyTorch reference itself runs cuDNN with
#    TF32).  Each block owns a tile of 2 output rows x W_OUT columns (<= 64
#    positions) of one batch image: 4 warps each compute a 16(pos) x 24(ch)
#    tile via 3 mmas and sweep all output depths keeping a running minimum in
#    registers.  The whole required input region (4 input rows x W x 3 ch x D
#    depth, <= 48 KB) is staged into dynamic shared memory once per block with
#    coalesced 128-bit loads, converting to TF32 *once* during staging, so the
#    depth sweep issues only low-latency shared loads (no global traffic, no
#    cvt instructions, no __syncthreads).  Weight fragments are preloaded into
#    registers once per block through the same shared buffer (aliased: weights
#    are staged first, consumed into registers, then overwritten by the input
#    region).  The channel softmax is a per-quad reduction (two shfl_xor), so
#    no cross-warp traffic is needed.
#
# 2) conv3d_min_softmax_tf32_kernel  (previously validated fused kernel that
#    reads A-fragments straight from global memory via __ldg; kept as a
#    known-good fast path / fallback for shapes the smem kernel does not
#    support).
#
# 3) min_softmax_kernel
#    Fused min-over-depth + softmax-over-channels used on the cuDNN fallback
#    path.
# ---------------------------------------------------------------------------

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

#define DEVINL __device__ __forceinline__

namespace {

constexpr int COUT = 24;
constexpr int KP   = 88;   // 81 real reduction terms + 7 zero padding
constexpr int IC   = 3;
constexpr int NR2  = 4;    // staged input rows (2 output rows + kernel size - 1)
constexpr int LDSW = COUT * 81 / 4;  // float4s of weights (486)

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

DEVINL void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                     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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// x: (B, 3, D, H, W) with 3 <= D <= 32, 4 <= W <= 32, W % 4 == 0, H >= 3
// wt: (24, 3, 3, 3, 3), bias: (24)
// out: (B, 24, H-2, W-2) == softmax_c min_d conv3d(x)
__global__ void __launch_bounds__(128)
conv3d_min_softmax_smem_kernel(const float* __restrict__ x,
                               const float* __restrict__ wt,
                               const float* __restrict__ bias,
                               float* __restrict__ out,
                               int B, int D, int H, int W,
                               int H_OUT, int W_OUT, int HWo) {
    extern __shared__ __align__(16) float s_raw[];  // [D][IC][NR2][W] tf32 slab

    __shared__ int   sk[KP];
    __shared__ float sbias[COUT];

    const int tid = threadIdx.x;
    const int HW  = H * W;
    const int DHW = D * HW;
    const int SD  = NR2 * IC * W;

    const int hps = (H_OUT + 1) >> 1;                 // 2-row strips per image
    const int b   = blockIdx.x / hps;
    const int h0  = (blockIdx.x - b * hps) << 1;

    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g    = lane >> 2;   // groupID within the mma tile
    const int t    = lane & 3;    // threadID-in-group

    const int rows_here = min(2, H_OUT - h0);
    const int npos      = rows_here * W_OUT;          // <= 64

    const int rp0 = (warp << 4) + g;
    const int rp1 = rp0 + 8;
    const bool v0 = rp0 < npos;
    const bool v1 = rp1 < npos;
    const int p0  = min(rp0, npos - 1);
    const int p1  = min(rp1, npos - 1);

    // local offsets inside the staged region: [ic][row][w]
    const int rl0 = p0 / W_OUT;
    const int rl1 = p1 / W_OUT;
    const int ro0 = rl0 * W + (p0 - rl0 * W_OUT);
    const int ro1 = rl1 * W + (p1 - rl1 * W_OUT);

    // ---- stage weights -> shared (tf32, c-major) ------------------------
    {
        const int4* w4 = reinterpret_cast<const int4*>(wt);
        int4* s4 = reinterpret_cast<int4*>(s_raw);
        for (int i = tid; i < LDSW; i += 128) {
            const int4 pk = w4[i];
            int4 o;
            o.x = (int)f2tf32(__int_as_float(pk.x));
            o.y = (int)f2tf32(__int_as_float(pk.y));
            o.z = (int)f2tf32(__int_as_float(pk.z));
            o.w = (int)f2tf32(__int_as_float(pk.w));
            s4[i] = o;
        }
    }
    for (int i = tid; i < KP; i += 128) {
        int off = 0;
        if (i < 81) {
            const int ic = i / 27, rem = i % 27;
            const int kd = rem / 9, kh = (rem % 9) / 3, kw = rem % 3;
            off = kd * SD + ic * (NR2 * W) + kh * W + kw;
        }
        sk[i] = off;
    }
    if (tid < COUT) sbias[tid] = bias[tid];
    __syncthreads();

    // ---- preload weight fragments into registers (depth-independent) ----
    unsigned breg[11][3][2];
    {
        const unsigned* sw = reinterpret_cast<const unsigned*>(s_raw);
        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            const int k0 = ks * 8 + t;
            const int k1 = k0 + 4;
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                const int cbase = (j * 8 + g) * 81;
                breg[ks][j][0] = (k0 < 81) ? sw[cbase + k0] : 0u;
                breg[ks][j][1] = (k1 < 81) ? sw[cbase + k1] : 0u;
            }
        }
    }
    __syncthreads();   // done with the weight staging area

    // ---- stage the input region (tf32) over the same shared memory ------
    {
        const float* xb = x + (long)b * IC * DHW;
        const int W4  = W >> 2;              // W % 4 == 0 (guaranteed by guard)
        const int per = D * IC * W4;         // float4s per staged row
        for (int r = 0; r < NR2; ++r) {
            const float* srcb = xb + (long)min(h0 + r, H - 1) * W;
            float* dstb = s_raw + r * W;
            for (int e = tid; e < per; e += 128) {
                const int d  = e / (IC * W4);
                const int rm = e - d * (IC * W4);
                const int ic = rm / W4;
                const int wo = (rm - ic * W4) << 2;
                const float4 pk = *reinterpret_cast<const float4*>(
                    srcb + (long)ic * DHW + (long)d * HW + wo);
                float4 o;
                o.x = __uint_as_float(f2tf32(pk.x));
                o.y = __uint_as_float(f2tf32(pk.y));
                o.z = __uint_as_float(f2tf32(pk.z));
                o.w = __uint_as_float(f2tf32(pk.w));
                *reinterpret_cast<float4*>(dstb + d * SD + ic * (NR2 * W) + wo) = o;
            }
        }
    }
    __syncthreads();

    // ---- main sweep over output depth: conv + running min ----------------
    const unsigned* s_u = reinterpret_cast<const unsigned*>(s_raw);

    float minv[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j)
        #pragma unroll
        for (int i = 0; i < 4; ++i) minv[j][i] = __int_as_float(0x7f800000);

    const int D_OUT = D - 2;
    for (int d = 0; d < D_OUT; ++d) {
        const unsigned* a0p = s_u + d * SD + ro0;
        const unsigned* a1p = s_u + d * SD + ro1;
        float acc[3][4];
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) acc[j][i] = 0.f;

        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            const int k0 = ks * 8 + t;
            const int k1 = k0 + 4;
            const int o0 = sk[k0];
            const int o1 = sk[k1];
            const unsigned a0 = a0p[o0];
            const unsigned a1 = a1p[o0];
            const unsigned a2 = a0p[o1];
            const unsigned a3 = a1p[o1];
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                mma_tf32(acc[j][0], acc[j][1], acc[j][2], acc[j][3],
                         a0, a1, a2, a3, breg[ks][j][0], breg[ks][j][1]);
            }
        }
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) minv[j][i] = fminf(minv[j][i], acc[j][i]);
    }

    // ---- bias + softmax over the 24 channels -----------------------------
    float v[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const float bj0 = sbias[j * 8 + 2 * t];
        const float bj1 = sbias[j * 8 + 2 * t + 1];
        v[j][0] = minv[j][0] + bj0;
        v[j][1] = minv[j][1] + bj1;
        v[j][2] = minv[j][2] + bj0;
        v[j][3] = minv[j][3] + bj1;
    }

    // each position's 24 channel values live in a quad of lanes (t = 0..3)
    float m0 = fmaxf(fmaxf(v[0][0], v[0][1]), fmaxf(v[1][0], v[1][1]));
    m0 = fmaxf(m0, fmaxf(v[2][0], v[2][1]));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2));
    float m1 = fmaxf(fmaxf(v[0][2], v[0][3]), fmaxf(v[1][2], v[1][3]));
    m1 = fmaxf(m1, fmaxf(v[2][2], v[2][3]));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2));

    float s0 = 0.f, s1 = 0.f;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        v[j][0] = expf(v[j][0] - m0); s0 += v[j][0];
        v[j][1] = expf(v[j][1] - m0); s0 += v[j][1];
        v[j][2] = expf(v[j][2] - m1); s1 += v[j][2];
        v[j][3] = expf(v[j][3] - m1); s1 += v[j][3];
    }
    s0 += __shfl_xor_sync(0xffffffffu, s0, 1);
    s0 += __shfl_xor_sync(0xffffffffu, s0, 2);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 1);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 2);
    const float r0 = 1.f / s0;
    const float r1 = 1.f / s1;

    const long ob = (long)b * COUT * HWo + h0 * W_OUT;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const int c0 = j * 8 + 2 * t;
        const int c1 = c0 + 1;
        if (v0) {
            out[ob + (long)c0 * HWo + rp0] = v[j][0] * r0;
            out[ob + (long)c1 * HWo + rp0] = v[j][1] * r0;
        }
        if (v1) {
            out[ob + (long)c0 * HWo + rp1] = v[j][2] * r1;
            out[ob + (long)c1 * HWo + rp1] = v[j][3] * r1;
        }
    }
}

// Previously validated fused kernel (global-memory A fragments).
// x: (B, 3, D, H, W), wt: (24, 3, 3, 3, 3), bias: (24)
// out: (B, 24, H-2, W-2)  == softmax_c min_d conv3d(x)
__global__ void __launch_bounds__(128)
conv3d_min_softmax_tf32_kernel(const float* __restrict__ x,
                               const float* __restrict__ wt,
                               const float* __restrict__ bias,
                               float* __restrict__ out,
                               int B, int D, int H, int W,
                               int H_OUT, int W_OUT) {
    __shared__ float sw[KP][COUT];
    __shared__ int   sk[KP];
    __shared__ float sb[COUT];

    const int tid = threadIdx.x;
    const int HW  = H * W;
    const int DHW = D * HW;

    for (int i = tid; i < KP * COUT; i += blockDim.x) {
        const int k = i / COUT, c = i % COUT;
        sw[k][c] = (k < 81) ? wt[c * 81 + k] : 0.f;
    }
    for (int i = tid; i < KP; i += blockDim.x) {
        int off = 0;
        if (i < 81) {
            const int ic = i / 27, rem = i % 27;
            const int kd = rem / 9, kh = (rem % 9) / 3, kw = rem % 3;
            off = ic * DHW + kd * HW + kh * W + kw;
        }
        sk[i] = off;
    }
    if (tid < COUT) sb[tid] = bias[tid];
    __syncthreads();

    const int HWo    = H_OUT * W_OUT;
    const int bpb    = (HWo + 63) >> 6;
    const int b      = blockIdx.x / bpb;
    const int p_base = (blockIdx.x - b * bpb) << 6;

    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g    = lane >> 2;
    const int t    = lane & 3;
    const int wp0  = p_base + (warp << 4);

    const int rp0 = wp0 + g;
    const int rp1 = wp0 + g + 8;
    const bool v0 = rp0 < HWo;
    const bool v1 = rp1 < HWo;
    const int p0  = min(rp0, HWo - 1);
    const int p1  = min(rp1, HWo - 1);
    const int ro0 = (p0 / W_OUT) * W + (p0 % W_OUT);
    const int ro1 = (p1 / W_OUT) * W + (p1 % W_OUT);

    const float* xb = x + (long)b * 3 * DHW;

    unsigned breg[11][3][2];
    #pragma unroll
    for (int ks = 0; ks < 11; ++ks) {
        #pragma unroll
        for (int j = 0; j < 3; ++j) {
            breg[ks][j][0] = f2tf32(sw[ks * 8 + t][j * 8 + g]);
            breg[ks][j][1] = f2tf32(sw[ks * 8 + t + 4][j * 8 + g]);
        }
    }

    float minv[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j)
        #pragma unroll
        for (int i = 0; i < 4; ++i) minv[j][i] = __int_as_float(0x7f800000);

    const int D_OUT = D - 2;
    for (int d = 0; d < D_OUT; ++d) {
        const float* xd = xb + (long)d * HW;
        float acc[3][4];
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) acc[j][i] = 0.f;

        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            const int k0 = ks * 8 + t;
            const int k1 = k0 + 4;
            const unsigned a0 = f2tf32(__ldg(xd + ro0 + sk[k0]));
            const unsigned a1 = f2tf32(__ldg(xd + ro1 + sk[k0]));
            const unsigned a2 = f2tf32(__ldg(xd + ro0 + sk[k1]));
            const unsigned a3 = f2tf32(__ldg(xd + ro1 + sk[k1]));
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                mma_tf32(acc[j][0], acc[j][1], acc[j][2], acc[j][3],
                         a0, a1, a2, a3, breg[ks][j][0], breg[ks][j][1]);
            }
        }
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) minv[j][i] = fminf(minv[j][i], acc[j][i]);
    }

    float v[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const float bj0 = sb[j * 8 + 2 * t];
        const float bj1 = sb[j * 8 + 2 * t + 1];
        v[j][0] = minv[j][0] + bj0;
        v[j][1] = minv[j][1] + bj1;
        v[j][2] = minv[j][2] + bj0;
        v[j][3] = minv[j][3] + bj1;
    }

    float m0 = fmaxf(fmaxf(v[0][0], v[0][1]), fmaxf(v[1][0], v[1][1]));
    m0 = fmaxf(m0, fmaxf(v[2][0], v[2][1]));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2));
    float m1 = fmaxf(fmaxf(v[0][2], v[0][3]), fmaxf(v[1][2], v[1][3]));
    m1 = fmaxf(m1, fmaxf(v[2][2], v[2][3]));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2));

    float s0 = 0.f, s1 = 0.f;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        v[j][0] = expf(v[j][0] - m0); s0 += v[j][0];
        v[j][1] = expf(v[j][1] - m0); s0 += v[j][1];
        v[j][2] = expf(v[j][2] - m1); s1 += v[j][2];
        v[j][3] = expf(v[j][3] - m1); s1 += v[j][3];
    }
    s0 += __shfl_xor_sync(0xffffffffu, s0, 1);
    s0 += __shfl_xor_sync(0xffffffffu, s0, 2);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 1);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 2);
    const float r0 = 1.f / s0;
    const float r1 = 1.f / s1;

    const long ob = (long)b * COUT * HWo;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const int c0 = j * 8 + 2 * t;
        const int c1 = c0 + 1;
        if (v0) {
            out[ob + (long)c0 * HWo + rp0] = v[j][0] * r0;
            out[ob + (long)c1 * HWo + rp0] = v[j][1] * r0;
        }
        if (v1) {
            out[ob + (long)c0 * HWo + rp1] = v[j][2] * r1;
            out[ob + (long)c1 * HWo + rp1] = v[j][3] * r1;
        }
    }
}

// Fused min over depth -> softmax over channels (cuDNN fallback path).
// t: (B, 24, D2, H2, W2), out: (B, 24, H2, W2), W2 <= 32
__global__ void __launch_bounds__(256)
min_softmax_kernel(const float* __restrict__ t,
                   float* __restrict__ out,
                   int B, int D2, int H2, int W2) {
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int gw   = blockIdx.x * (blockDim.x >> 5) + warp;
    const int rows = B * H2;
    if (gw >= rows) return;

    const int h      = gw % H2;
    const int b      = gw / H2;
    const int w      = min(lane, W2 - 1);
    const bool valid = (lane < W2);

    const long HW2 = (long)H2 * W2;
    const float* base = t + (long)b * COUT * D2 * HW2 + (long)h * W2 + w;

    float vals[COUT];
    #pragma unroll
    for (int c = 0; c < COUT; ++c) vals[c] = __ldg(base + (long)c * D2 * HW2);

    for (int d = 1; d < D2; ++d) {
        #pragma unroll
        for (int c = 0; c < COUT; ++c) {
            const float vv = __ldg(base + ((long)c * D2 + d) * HW2);
            vals[c] = fminf(vals[c], vv);
        }
    }

    float mx = vals[0];
    #pragma unroll
    for (int c = 1; c < COUT; ++c) mx = fmaxf(mx, vals[c]);
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < COUT; ++c) { vals[c] = expf(vals[c] - mx); sum += vals[c]; }

    if (valid) {
        const long ob = (long)b * COUT * HW2 + (long)h * W2 + w;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) out[ob + (long)c * HW2] = vals[c] / sum;
    }
}

}  // namespace

torch::Tensor conv3d_min_softmax_smem(torch::Tensor x, torch::Tensor weight,
                                      torch::Tensor bias) {
    const int B = x.size(0);
    const int D = x.size(2), H = x.size(3), W = x.size(4);
    const int D_OUT = D - 2, H_OUT = H - 2, W_OUT = W - 2;
    TORCH_CHECK(D_OUT >= 1 && H_OUT >= 1 && W_OUT >= 1, "bad shape");
    TORCH_CHECK(D <= 32 && W <= 32 && (W % 4) == 0, "shape unsupported");

    auto out = torch::empty({B, COUT, H_OUT, W_OUT}, x.options());
    const int HWo  = H_OUT * W_OUT;
    const long hps = (H_OUT + 1) / 2;
    const long grid = (long)B * hps;

    const int region = D * IC * NR2 * W;             // staged input floats
    const int smem   = (region > COUT * 81 ? region : COUT * 81) * 4;

    if (smem > 48 * 1024) {
        static bool configured = false;
        if (!configured) {
            cudaFuncSetAttribute(conv3d_min_softmax_smem_kernel,
                                 cudaFuncAttributeMaxDynamicSharedMemorySize,
                                 64 * 1024);
            configured = true;
        }
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_smem_kernel<<<(int)grid, 128, smem, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), B, D, H, W, H_OUT, W_OUT, HWo);
    return out;
}

torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight,
                                      torch::Tensor bias) {
    const int B = x.size(0);
    const int D = x.size(2), H = x.size(3), W = x.size(4);
    const int D_OUT = D - 2, H_OUT = H - 2, W_OUT = W - 2;
    TORCH_CHECK(D_OUT >= 1 && H_OUT >= 1 && W_OUT >= 1, "bad shape");

    auto out = torch::empty({B, 24, H_OUT, W_OUT}, x.options());
    const long HWo = (long)H_OUT * W_OUT;
    const long bpb = (HWo + 63) / 64;
    const long grid = (long)B * bpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_tf32_kernel<<<(int)grid, 128, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), B, D, H, W, H_OUT, W_OUT);
    return out;
}

torch::Tensor min_softmax(torch::Tensor t) {
    const int B  = t.size(0);
    const int D2 = t.size(2), H2 = t.size(3), W2 = t.size(4);
    TORCH_CHECK(W2 <= 32, "unsupported width");

    auto out = torch::empty({B, 24, H2, W2}, t.options());
    const int rows   = B * H2;
    const int wpb    = 8;
    const int blocks = (rows + wpb - 1) / wpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    min_softmax_kernel<<<blocks, wpb * 32, 0, stream>>>(
        t.data_ptr<float>(), out.data_ptr<float>(), B, D2, H2, W2);
    return out;
}
"""

cpp_source = r"""
torch::Tensor conv3d_min_softmax_smem(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor min_softmax(torch::Tensor t);
"""

_fused_ops = load_inline(
    name="fused_conv3d_min_softmax_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_min_softmax_smem", "conv3d_min_softmax_tf32", "min_softmax"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Fast paths (both fully fused Conv3d + min-over-depth + channel-softmax,
    writing only the final 4-D output so the large 5-D conv result is never
    materialized):
      * 'smem' : TF32 tensor-core kernel that stages the input region in
                 shared memory (used when the shape supports it);
      * 'tf32' : TF32 tensor-core kernel reading A-fragments via __ldg.
    Fallbacks: cuDNN Conv3d + fused min/softmax kernel, then pure PyTorch.
    On the first CUDA forward all applicable paths are verified against the
    cuDNN path and timed; the fastest correct one is cached and reused.
    """

    def __init__(self, in_channels, out_channels, kernel_size, dim):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.dim = dim
        self._path = None  # 'smem', 'tf32' or 'cudnn', chosen on first forward
        self._fast_ok = (
            in_channels == 3
            and out_channels == 24
            and kernel_size == 3
            and dim == 2
        )
        self._mma_ok = None

    def _reference(self, x):
        x = self.conv(x)
        x = torch.min(x, dim=self.dim)[0]
        x = torch.softmax(x, dim=1)
        return x

    def _cudnn_path(self, x, w, b):
        return _fused_ops.min_softmax(F.conv3d(x, w, b))

    def _tf32_path(self, x, w, b):
        return _fused_ops.conv3d_min_softmax_tf32(x, w, b)

    def _smem_path(self, x, w, b):
        return _fused_ops.conv3d_min_softmax_smem(x, w, b)

    def _supports_mma(self):
        if self._mma_ok is None:
            try:
                cap = torch.cuda.get_device_capability()
                self._mma_ok = bool(cap[0] >= 8)
            except Exception:
                self._mma_ok = False
        return self._mma_ok

    def _run(self, path, x, w, b):
        if path == "smem":
            return self._smem_path(x, w, b)
        if path == "tf32":
            return self._tf32_path(x, w, b)
        return self._cudnn_path(x, w, b)

```
Result of evaluating it:
The solution compiled but failed at run time. Error output:
```
Module [ModelNew] is missing the required "forward" function
```

The fastest correct solution so far (6.398x speedup) is:
```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

# ---------------------------------------------------------------------------
# Custom CUDA kernels
#
# 1) conv3d_min_softmax_tf32_kernel:
#    Fully fused kernel computing in ONE pass, with zero intermediates:
#       Conv3d (IC=3, K=3, COUT=24, stride 1, no padding)
#       -> min over depth (dim=2)
#       -> softmax over channels (dim=1)
#    The convolution is evaluated as an implicit GEMM on TF32 tensor cores
#    (mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32), which is allowed by
#    the task's precision rule (the PyTorch reference itself runs cuDNN with
#    TF32 enabled).  Each block handles 64 consecutive output positions
#    (flattened h*w) of one batch image; each of its 4 warps owns a 16x24
#    (position x channel) tile and sweeps all output depths, keeping a running
#    minimum in registers.  The 88-entry (81 real + 7 padding) patch-offset
#    table and the k-major weight matrix live in shared memory; weight
#    fragments are preloaded into registers once per block (they do not depend
#    on depth).  The channel softmax is a per-row reduction over the 4 lanes
#    of each accumulator quad (two __shfl_xor_sync steps), so no cross-warp
#    communication is ever needed.
#
# 2) min_softmax_kernel:
#    Fused min-over-depth + softmax-over-channels used on the cuDNN fallback
#    path (kept from the previous, known-correct solution).
# ---------------------------------------------------------------------------

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

#define DEVINL __device__ __forceinline__

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

DEVINL void mma_tf32(float &d0, float &d1, float &d2, float &d3,
                     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"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1));
}

// x:  (B, 3, D, H, W), wt: (24, 3, 3, 3, 3), bias: (24)
// out: (B, 24, H-2, W-2)  == softmax_c min_d conv3d(x)
__global__ void __launch_bounds__(128)
conv3d_min_softmax_tf32_kernel(const float* __restrict__ x,
                               const float* __restrict__ wt,
                               const float* __restrict__ bias,
                               float* __restrict__ out,
                               int B, int D, int H, int W,
                               int H_OUT, int W_OUT) {
    constexpr int COUT = 24;
    constexpr int KP = 88;  // 81 real k + 7 zero padding

    __shared__ float sw[KP][COUT];
    __shared__ int   sk[KP];
    __shared__ float sb[COUT];

    const int tid = threadIdx.x;
    const int HW  = H * W;
    const int DHW = D * HW;

    for (int i = tid; i < KP * COUT; i += blockDim.x) {
        const int k = i / COUT, c = i % COUT;
        sw[k][c] = (k < 81) ? wt[c * 81 + k] : 0.f;
    }
    for (int i = tid; i < KP; i += blockDim.x) {
        int off = 0;
        if (i < 81) {
            const int ic = i / 27, rem = i % 27;
            const int kd = rem / 9, kh = (rem % 9) / 3, kw = rem % 3;
            off = ic * DHW + kd * HW + kh * W + kw;
        }
        sk[i] = off;
    }
    if (tid < COUT) sb[tid] = bias[tid];
    __syncthreads();

    const int HWo    = H_OUT * W_OUT;
    const int bpb    = (HWo + 63) >> 6;
    const int b      = blockIdx.x / bpb;
    const int p_base = (blockIdx.x - b * bpb) << 6;

    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int g    = lane >> 2;   // groupID
    const int t    = lane & 3;    // threadID_in_group
    const int wp0  = p_base + (warp << 4);

    const int rp0 = wp0 + g;
    const int rp1 = wp0 + g + 8;
    const bool v0 = rp0 < HWo;
    const bool v1 = rp1 < HWo;
    const int p0  = min(rp0, HWo - 1);
    const int p1  = min(rp1, HWo - 1);
    const int ro0 = (p0 / W_OUT) * W + (p0 % W_OUT);
    const int ro1 = (p1 / W_OUT) * W + (p1 % W_OUT);

    const float* xb = x + (long)b * 3 * DHW;

    // Preload weight fragments (independent of depth) into registers.
    unsigned breg[11][3][2];
    #pragma unroll
    for (int ks = 0; ks < 11; ++ks) {
        #pragma unroll
        for (int j = 0; j < 3; ++j) {
            breg[ks][j][0] = f2tf32(sw[ks * 8 + t][j * 8 + g]);
            breg[ks][j][1] = f2tf32(sw[ks * 8 + t + 4][j * 8 + g]);
        }
    }

    float minv[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j)
        #pragma unroll
        for (int i = 0; i < 4; ++i) minv[j][i] = __int_as_float(0x7f800000);

    const int D_OUT = D - 2;
    for (int d = 0; d < D_OUT; ++d) {
        const float* xd = xb + (long)d * HW;
        float acc[3][4];
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) acc[j][i] = 0.f;

        #pragma unroll
        for (int ks = 0; ks < 11; ++ks) {
            const int k0 = ks * 8 + t;
            const int k1 = k0 + 4;
            const unsigned a0 = f2tf32(__ldg(xd + ro0 + sk[k0]));
            const unsigned a1 = f2tf32(__ldg(xd + ro1 + sk[k0]));
            const unsigned a2 = f2tf32(__ldg(xd + ro0 + sk[k1]));
            const unsigned a3 = f2tf32(__ldg(xd + ro1 + sk[k1]));
            #pragma unroll
            for (int j = 0; j < 3; ++j) {
                mma_tf32(acc[j][0], acc[j][1], acc[j][2], acc[j][3],
                         a0, a1, a2, a3, breg[ks][j][0], breg[ks][j][1]);
            }
        }
        #pragma unroll
        for (int j = 0; j < 3; ++j)
            #pragma unroll
            for (int i = 0; i < 4; ++i) minv[j][i] = fminf(minv[j][i], acc[j][i]);
    }

    // bias (constant over depth, so it can be added after the min)
    float v[3][4];
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const float bj0 = sb[j * 8 + 2 * t];
        const float bj1 = sb[j * 8 + 2 * t + 1];
        v[j][0] = minv[j][0] + bj0;
        v[j][1] = minv[j][1] + bj1;
        v[j][2] = minv[j][2] + bj0;
        v[j][3] = minv[j][3] + bj1;
    }

    // softmax over the 24 channels; each position's values live in a quad of
    // lanes (t = 0..3) -> reduce with two xor-shuffles.
    float m0 = fmaxf(fmaxf(v[0][0], v[0][1]), fmaxf(v[1][0], v[1][1]));
    m0 = fmaxf(m0, fmaxf(v[2][0], v[2][1]));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 1));
    m0 = fmaxf(m0, __shfl_xor_sync(0xffffffffu, m0, 2));
    float m1 = fmaxf(fmaxf(v[0][2], v[0][3]), fmaxf(v[1][2], v[1][3]));
    m1 = fmaxf(m1, fmaxf(v[2][2], v[2][3]));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 1));
    m1 = fmaxf(m1, __shfl_xor_sync(0xffffffffu, m1, 2));

    float s0 = 0.f, s1 = 0.f;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        v[j][0] = expf(v[j][0] - m0); s0 += v[j][0];
        v[j][1] = expf(v[j][1] - m0); s0 += v[j][1];
        v[j][2] = expf(v[j][2] - m1); s1 += v[j][2];
        v[j][3] = expf(v[j][3] - m1); s1 += v[j][3];
    }
    s0 += __shfl_xor_sync(0xffffffffu, s0, 1);
    s0 += __shfl_xor_sync(0xffffffffu, s0, 2);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 1);
    s1 += __shfl_xor_sync(0xffffffffu, s1, 2);
    const float r0 = 1.f / s0;
    const float r1 = 1.f / s1;

    const long ob = (long)b * COUT * HWo;
    #pragma unroll
    for (int j = 0; j < 3; ++j) {
        const int c0 = j * 8 + 2 * t;
        const int c1 = c0 + 1;
        if (v0) {
            out[ob + (long)c0 * HWo + rp0] = v[j][0] * r0;
            out[ob + (long)c1 * HWo + rp0] = v[j][1] * r0;
        }
        if (v1) {
            out[ob + (long)c0 * HWo + rp1] = v[j][2] * r1;
            out[ob + (long)c1 * HWo + rp1] = v[j][3] * r1;
        }
    }
}

// Fused min over depth -> softmax over channels (cuDNN fallback path).
// t: (B, 24, D2, H2, W2), out: (B, 24, H2, W2), W2 <= 32
__global__ void __launch_bounds__(256)
min_softmax_kernel(const float* __restrict__ t,
                   float* __restrict__ out,
                   int B, int D2, int H2, int W2) {
    constexpr int COUT = 24;
    const int tid  = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int gw   = blockIdx.x * (blockDim.x >> 5) + warp;
    const int rows = B * H2;
    if (gw >= rows) return;

    const int h      = gw % H2;
    const int b      = gw / H2;
    const int w      = min(lane, W2 - 1);
    const bool valid = (lane < W2);

    const long HW2 = (long)H2 * W2;
    const float* base = t + (long)b * COUT * D2 * HW2 + (long)h * W2 + w;

    float vals[COUT];
    #pragma unroll
    for (int c = 0; c < COUT; ++c) vals[c] = __ldg(base + (long)c * D2 * HW2);

    for (int d = 1; d < D2; ++d) {
        #pragma unroll
        for (int c = 0; c < COUT; ++c) {
            const float vv = __ldg(base + ((long)c * D2 + d) * HW2);
            vals[c] = fminf(vals[c], vv);
        }
    }

    float mx = vals[0];
    #pragma unroll
    for (int c = 1; c < COUT; ++c) mx = fmaxf(mx, vals[c]);
    float sum = 0.f;
    #pragma unroll
    for (int c = 0; c < COUT; ++c) { vals[c] = expf(vals[c] - mx); sum += vals[c]; }

    if (valid) {
        const long ob = (long)b * COUT * HW2 + (long)h * W2 + w;
        #pragma unroll
        for (int c = 0; c < COUT; ++c) out[ob + (long)c * HW2] = vals[c] / sum;
    }
}

torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight,
                                      torch::Tensor bias) {
    const int B = x.size(0);
    const int D = x.size(2), H = x.size(3), W = x.size(4);
    const int D_OUT = D - 2, H_OUT = H - 2, W_OUT = W - 2;
    TORCH_CHECK(D_OUT >= 1 && H_OUT >= 1 && W_OUT >= 1, "bad shape");

    auto out = torch::empty({B, 24, H_OUT, W_OUT}, x.options());
    const long HWo = (long)H_OUT * W_OUT;
    const long bpb = (HWo + 63) / 64;
    const long grid = (long)B * bpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    conv3d_min_softmax_tf32_kernel<<<(int)grid, 128, 0, stream>>>(
        x.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(), B, D, H, W, H_OUT, W_OUT);
    return out;
}

torch::Tensor min_softmax(torch::Tensor t) {
    const int B  = t.size(0);
    const int D2 = t.size(2), H2 = t.size(3), W2 = t.size(4);
    TORCH_CHECK(W2 <= 32, "unsupported width");

    auto out = torch::empty({B, 24, H2, W2}, t.options());
    const int rows   = B * H2;
    const int wpb    = 8;
    const int blocks = (rows + wpb - 1) / wpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    min_softmax_kernel<<<blocks, wpb * 32, 0, stream>>>(
        t.data_ptr<float>(), out.data_ptr<float>(), B, D2, H2, W2);
    return out;
}
"""

cpp_source = r"""
torch::Tensor conv3d_min_softmax_tf32(torch::Tensor x, torch::Tensor weight, torch::Tensor bias);
torch::Tensor min_softmax(torch::Tensor t);
"""

_fused_ops = load_inline(
    name="fused_conv3d_min_softmax_tf32",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["conv3d_min_softmax_tf32", "min_softmax"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    Fast path: a single fused TF32 tensor-core CUDA kernel that performs
    Conv3d + min-over-depth + softmax-over-channels in one pass, writing only
    the final 4-D output (the large 5-D conv result is never materialized).
    Fallbacks: cuDNN Conv3d + fused min/softmax kernel, then pure PyTorch.
    The fastest correct path is selected once at runtime.
    """

    def __init__(self, in_channels, out_channels, kernel_size, dim):
        super(ModelNew, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size)
        self.dim = dim
        self._path = None  # 'tf32' or 'cudnn', chosen on first CUDA forward
        self._fast_ok = (
            in_channels == 3
            and out_channels == 24
            and kernel_size == 3
            and dim == 2
        )
        self._tf32_ok = None

    def _reference(self, x):
        x = self.conv(x)
        x = torch.min(x, dim=self.dim)[0]
        x = torch.softmax(x, dim=1)
        return x

    def _cudnn_path(self, x, w, b):
        return _fused_ops.min_softmax(F.conv3d(x, w, b))

    def _tf32_path(self, x, w, b):
        return _fused_ops.conv3d_min_softmax_tf32(x, w, b)

    def _supports_tf32(self):
        if self._tf32_ok is None:
            try:
                cap = torch.cuda.get_device_capability()
                self._tf32_ok = bool(cap[0] >= 8)
            except Exception:
                self._tf32_ok = False
        return self._tf32_ok

    def forward(self, x):
        if (
            not self._fast_ok
            or not x.is_cuda
            or x.dtype != torch.float32
            or x.dim() != 5
            or x.size(1) != 3
            or x.size(2) < 3
            or x.size(3) < 3
            or x.size(4) < 3
            or x.size(4) - 2 > 32  # cudnn fallback kernel needs W_OUT <= 32
        ):
            return self._reference(x)

        xc = x.contiguous()
        wc = self.conv.weight.contiguous()
        bc = self.conv.bias.contiguous()

        if self._path == "tf32":
            return self._tf32_path(xc, wc, bc)
        if self._path == "cudnn":
            return self._cudnn_path(xc, wc, bc)

        # One-time selection of the fastest correct path.
        r_cudnn = self._cudnn_path(xc, wc, bc)
        r_tf32 = None
        if self._supports_tf32():
            r_tf32 = self._tf32_path(xc, wc, bc)
            torch.cuda.synchronize()
            if not torch.allclose(r_tf32, r_cudnn, rtol=1e-2, atol=1e-2):
                r_tf32 = None  # fused kernel not trustworthy on this device

        if r_tf32 is None:
            self._path = "cudnn"
            return r_cudnn

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        iters = 5

        start.record()
        for _ in range(iters):
            self._tf32_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_tf32 = start.elapsed_time(end)

        start.record()
        for _ in range(iters):
            self._cudnn_path(xc, wc, bc)
        end.record()
        torch.cuda.synchronize()
        t_cudnn = start.elapsed_time(end)

        if t_tf32 <= t_cudnn:
            self._path = "tf32"
            return r_tf32
        self._path = "cudnn"
        return r_cudnn
```

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.


Precision rule: the timed forward must compute in FP32 (TF32 tensor-core math is allowed, as in the PyTorch reference). Do not cast inputs, weights or intermediate results to float16 or bfloat16 and do not call half-precision kernels in forward().
