```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_source = r"""
#include <torch/extension.h>

torch::Tensor fused_forward(
    torch::Tensor input,
    torch::Tensor linear_weight,
    torch::Tensor linear_bias,
    torch::Tensor bn_weight,
    torch::Tensor bn_bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor scale,
    bool training,
    double eps,
    double momentum);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <math.h>

constexpr int TILE_ROWS = 128;
constexpr int THREADS = 256;

__global__ void partial_stats(
    const float* __restrict__ x,
    const float* __restrict__ linear_bias,
    float* __restrict__ partial,
    int rows, int cols, int parts)
{
    __shared__ float means[8][32];
    __shared__ float m2s[8][32];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int col = blockIdx.x * 32 + tx;
    int start = blockIdx.y * TILE_ROWS;
    int end = min(start + TILE_ROWS, rows);

    float mean = 0.0f;
    float m2 = 0.0f;
    int count = 0;

    if (col < cols) {
        float b = linear_bias[col];
        for (int row = start + ty; row < end; row += 8) {
            float v = __fadd_rn(x[(int64_t)row * cols + col], b);
            ++count;
            float delta = v - mean;
            mean += delta / count;
            m2 += delta * (v - mean);
        }
    }

    means[ty][tx] = mean;
    m2s[ty][tx] = m2;
    __syncthreads();

    if (ty == 0 && col < cols) {
        float total_mean = 0.0f;
        float total_m2 = 0.0f;
        int total_count = 0;

        #pragma unroll
        for (int lane = 0; lane < 8; ++lane) {
            int n = max(0, (end - start - lane + 7) / 8);
            if (n > 0) {
                int combined = total_count + n;
                float delta = means[lane][tx] - total_mean;
                total_m2 += m2s[lane][tx]
                    + delta * delta * ((float)total_count * n / combined);
                total_mean += delta * ((float)n / combined);
                total_count = combined;
            }
        }

        int64_t index = (int64_t)blockIdx.y * cols + col;
        partial[index] = total_mean;
        partial[(int64_t)parts * cols + index] = total_m2;
    }
}

__global__ void training_coefficients(
    const float* __restrict__ partial,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ coefficients,
    int rows, int cols, int parts,
    float eps, float momentum)
{
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

    float mean = 0.0f;
    float m2 = 0.0f;
    int count = 0;

    for (int part = 0; part < parts; ++part) {
        int n = min(TILE_ROWS, rows - part * TILE_ROWS);
        int combined = count + n;
        int64_t index = (int64_t)part * cols + col;
        float delta = partial[index] - mean;

        m2 += partial[(int64_t)parts * cols + index]
            + delta * delta * ((float)count * n / combined);
        mean += delta * ((float)n / combined);
        count = combined;
    }

    coefficients[col] = mean;
    coefficients[cols + col] =
        weight[col] * rsqrtf(fmaxf(m2 / rows, 0.0f) + eps);
    coefficients[2 * cols + col] = bias[col];

    running_mean[col] =
        (1.0f - momentum) * running_mean[col] + momentum * mean;
    running_var[col] =
        (1.0f - momentum) * running_var[col]
        + momentum * (m2 / (rows - 1));
}

__global__ void inference_coefficients(
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float* __restrict__ coefficients,
    int cols, float eps)
{
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

    coefficients[col] = running_mean[col];
    coefficients[cols + col] = weight[col] * rsqrtf(running_var[col] + eps);
    coefficients[2 * cols + col] = bias[col];
}

template<bool MAXIMUM>
__device__ __forceinline__ float block_reduce(float value)
{
    __shared__ float scratch[8];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, value, offset);
        if (MAXIMUM) value = fmaxf(value, other);
        else value += other;
    }

    if (lane == 0) scratch[warp] = value;
    __syncthreads();

    if (warp == 0) {
        value = lane < 8 ? scratch[lane] : (MAXIMUM ? -INFINITY : 0.0f);

        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            float other = __shfl_down_sync(0xffffffff, value, offset);
            if (MAXIMUM) value = fmaxf(value, other);
            else value += other;
        }

        if (lane == 0) scratch[0] = value;
    }

    __syncthreads();
    float result = scratch[0];
    __syncthreads();
    return result;
}

template<int ITEMS>
__global__ void normalized_softmax(
    const float* __restrict__ x,
    const float* __restrict__ linear_bias,
    const float* __restrict__ coefficients,
    const float* __restrict__ scale,
    float* __restrict__ out,
    int cols, bool scalar_scale)
{
    int tid = threadIdx.x;
    int64_t base = (int64_t)blockIdx.x * cols;

    float values[ITEMS];
    float maximum = -INFINITY;

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int col = tid + i * THREADS;
        float value = -INFINITY;

        if (col < cols) {
            float z = __fadd_rn(x[base + col], linear_bias[col]);
            value = fmaf(
                z - coefficients[col],
                coefficients[cols + col],
                coefficients[2 * cols + col]
            ) * scale[scalar_scale ? 0 : col];
        }

        values[i] = value;
        maximum = fmaxf(maximum, value);
    }

    maximum = block_reduce<true>(maximum);
    float sum = 0.0f;

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        values[i] = __expf(values[i] - maximum);
        sum += values[i];
    }

    float inverse_sum = 1.0f / block_reduce<false>(sum);

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int col = tid + i * THREADS;
        if (col < cols)
            out[base + col] = values[i] * inverse_sum;
    }
}

torch::Tensor fused_forward(
    torch::Tensor input,
    torch::Tensor linear_weight,
    torch::Tensor linear_bias,
    torch::Tensor bn_weight,
    torch::Tensor bn_bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor scale,
    bool training,
    double eps,
    double momentum)
{
    const c10::cuda::CUDAGuard guard(input.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    int rows = input.size(0);
    int inner = input.size(1);
    int cols = linear_weight.size(0);

    auto x = torch::empty({rows, cols}, input.options());
    auto out = torch::empty_like(x);
    if (rows == 0) return out;

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    TORCH_CHECK(
        cublasSetStream(handle, stream) == CUBLAS_STATUS_SUCCESS,
        "cublasSetStream failed");

    const float alpha = 1.0f;
    const float beta = 0.0f;

    // Row-major X @ W.T is column-major W @ X.T.
    // Explicitly enable TF32 tensor cores without changing PyTorch settings.
    cublasStatus_t status = cublasGemmEx(
        handle,
        CUBLAS_OP_T, CUBLAS_OP_N,
        cols, rows, inner,
        &alpha,
        linear_weight.data_ptr<float>(), CUDA_R_32F, inner,
        input.data_ptr<float>(), CUDA_R_32F, inner,
        &beta,
        x.data_ptr<float>(), CUDA_R_32F, cols,
        CUBLAS_COMPUTE_32F_FAST_TF32,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP);

    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
                "cublasGemmEx failed: ", static_cast<int>(status));

    auto coefficients = torch::empty({3, cols}, input.options());

    if (training) {
        int parts = (rows + TILE_ROWS - 1) / TILE_ROWS;
        auto partial = torch::empty({2, parts, cols}, input.options());

        partial_stats<<<
            dim3((cols + 31) / 32, parts), dim3(32, 8), 0, stream
        >>>(
            x.data_ptr<float>(), linear_bias.data_ptr<float>(),
            partial.data_ptr<float>(), rows, cols, parts
        );

        training_coefficients<<<(cols + 255) / 256, 256, 0, stream>>>(
            partial.data_ptr<float>(),
            bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            coefficients.data_ptr<float>(),
            rows, cols, parts, (float)eps, (float)momentum
        );
    } else {
        inference_coefficients<<<(cols + 255) / 256, 256, 0, stream>>>(
            bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            coefficients.data_ptr<float>(), cols, (float)eps
        );
    }

    bool scalar_scale = scale.numel() == 1;

    #define LAUNCH(N) \
        normalized_softmax<N><<<rows, THREADS, 0, stream>>>( \
            x.data_ptr<float>(), linear_bias.data_ptr<float>(), \
            coefficients.data_ptr<float>(), scale.data_ptr<float>(), \
            out.data_ptr<float>(), cols, scalar_scale)

    if      (cols <= 256)   { LAUNCH(1); }
    else if (cols <= 512)   { LAUNCH(2); }
    else if (cols <= 1024)  { LAUNCH(4); }
    else if (cols <= 2048)  { LAUNCH(8); }
    else if (cols <= 4096)  { LAUNCH(16); }
    else if (cols <= 8192)  { LAUNCH(32); }
    else                   { LAUNCH(64); }

    #undef LAUNCH

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ops = load_inline(
    name="tf32_gemm_bn_scale_softmax_v2",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(
        self,
        in_features,
        out_features,
        bn_eps=1e-5,
        bn_momentum=0.1,
        scale_shape=(1,),
    ):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(
            out_features, eps=bn_eps, momentum=bn_momentum
        )
        self.scale = nn.Parameter(torch.ones(scale_shape))
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        cols = self.gemm.out_features
        tensors = (
            self.gemm.weight,
            self.gemm.bias,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.scale,
        )

        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and x.size(1) == self.gemm.in_features
            and x.size(1) > 0
            and 0 < cols <= 16384
            and not torch.is_autocast_enabled()
            and self.bn.affine
            and self.bn.track_running_stats
            and (not self.bn.training or x.size(0) > 1)
            and tuple(self.scale.shape)
            in {(), (1,), (cols,), (1, 1), (1, cols)}
            and all(
                t is not None
                and t.device == x.device
                and t.dtype == torch.float32
                and t.is_contiguous()
                for t in tensors
            )
        )

        if not supported:
            return self.softmax(self.scale * self.bn(self.gemm(x)))

        momentum = self.bn.momentum
        if self.bn.training:
            self.bn.num_batches_tracked.add_(1)
            if momentum is None:
                momentum = 1.0 / self.bn.num_batches_tracked.item()
        elif momentum is None:
            momentum = 0.0

        return _ops.fused_forward(
            x,
            self.gemm.weight,
            self.gemm.bias,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.scale,
            self.bn.training,
            self.bn.eps,
            momentum,
        )
```