```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_bn_scale_softmax(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches,
    torch::Tensor scale,
    bool training,
    double momentum,
    double eps);
"""

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

__global__ void increment_batches(int64_t* batches) {
    if (threadIdx.x == 0) ++batches[0];
}

__global__ void batch_statistics(
    const float* __restrict__ x,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    const int64_t* __restrict__ batches,
    float* __restrict__ mean_out,
    float* __restrict__ invstd_out,
    int rows,
    int cols,
    bool training,
    float momentum,
    float eps)
{
    const int lane = threadIdx.x;
    const int part = threadIdx.y;
    const int tid = part * 32 + lane;
    const int col = blockIdx.x * 32 + lane;

    if (!training) {
        if (part == 0 && col < cols) {
            mean_out[col] = running_mean[col];
            invstd_out[col] = rsqrtf(running_var[col] + eps);
        }
        return;
    }

    __shared__ float means[256];
    __shared__ float moments[256];
    __shared__ int counts[256];

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

    if (col < cols) {
        for (int row = part; row < rows; row += 8) {
            const float value = x[(int64_t)row * cols + col];
            ++count;
            const float delta = value - mean;
            mean += delta / (float)count;
            m2 += delta * (value - mean);
        }
    }

    means[tid] = mean;
    moments[tid] = m2;
    counts[tid] = count;
    __syncthreads();

    #pragma unroll
    for (int step = 4; step > 0; step >>= 1) {
        if (part < step) {
            const int other = tid + step * 32;
            const int na = counts[tid];
            const int nb = counts[other];
            const int total = na + nb;
            if (nb != 0) {
                const float delta = means[other] - means[tid];
                means[tid] += delta * ((float)nb / (float)total);
                moments[tid] += moments[other]
                    + delta * delta * ((float)na * nb / (float)total);
            }
            counts[tid] = total;
        }
        __syncthreads();
    }

    if (part == 0 && col < cols) {
        const float batch_mean = means[tid];
        const float batch_m2 = moments[tid];
        const float factor = momentum < 0.0f
            ? 1.0f / (float)batches[0]
            : momentum;

        mean_out[col] = batch_mean;
        invstd_out[col] = rsqrtf(batch_m2 / (float)rows + eps);

        running_mean[col] =
            (1.0f - factor) * running_mean[col] + factor * batch_mean;
        running_var[col] =
            (1.0f - factor) * running_var[col]
            + factor * (batch_m2 / (float)(rows - 1));
    }
}

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

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        const float other = __shfl_down_sync(0xffffffff, value, offset);
        value = IsMax ? fmaxf(value, other) : value + other;
    }

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

    if (warp == 0) {
        value = lane < 8 ? scratch[lane]
                        : (IsMax ? -CUDART_INF_F : 0.0f);
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            const float other = __shfl_down_sync(0xffffffff, value, offset);
            value = IsMax ? fmaxf(value, other) : value + other;
        }
        if (lane == 0) scratch[0] = value;
    }
    __syncthreads();
    return scratch[0];
}

template<int Items>
__global__ void normalize_scale_softmax(
    const float* __restrict__ x,
    const float* __restrict__ mean,
    const float* __restrict__ invstd,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    float* __restrict__ output,
    int cols,
    int64_t scale_row_stride,
    int64_t scale_col_stride)
{
    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    const int64_t base = (int64_t)row * cols;
    const int64_t scale_base = (int64_t)row * scale_row_stride;

    float values[Items];
    float local_max = -CUDART_INF_F;

    #pragma unroll
    for (int i = 0; i < Items; ++i) {
        const int col = tid + i * 256;
        float value = -CUDART_INF_F;
        if (col < cols) {
            value = (x[base + col] - mean[col]) * invstd[col];
            value = value * weight[col] + bias[col];
            value *= scale[scale_base + (int64_t)col * scale_col_stride];
        }
        values[i] = value;
        local_max = fmaxf(local_max, value);
    }

    const float row_max = block_reduce<true>(local_max);
    float local_sum = 0.0f;

    #pragma unroll
    for (int i = 0; i < Items; ++i) {
        values[i] = __expf(values[i] - row_max);
        local_sum += values[i];
    }

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

    #pragma unroll
    for (int i = 0; i < Items; ++i) {
        const int col = tid + i * 256;
        if (col < cols)
            output[base + col] = values[i] * inverse_sum;
    }
}

torch::Tensor fused_bn_scale_softmax(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches,
    torch::Tensor scale,
    bool training,
    double momentum,
    double eps)
{
    const c10::cuda::CUDAGuard guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int rows = (int)x.size(0);
    const int cols = (int)x.size(1);
    auto output = torch::empty_like(x);
    if (rows == 0) return output;

    auto statistics = torch::empty({2, cols}, x.options());
    float* mean = statistics.data_ptr<float>();
    float* invstd = mean + cols;

    if (training) {
        increment_batches<<<1, 1, 0, stream>>>(
            batches.data_ptr<int64_t>());
    }

    batch_statistics<<<(cols + 31) / 32, dim3(32, 8), 0, stream>>>(
        x.data_ptr<float>(),
        running_mean.data_ptr<float>(),
        running_var.data_ptr<float>(),
        batches.data_ptr<int64_t>(),
        mean,
        invstd,
        rows,
        cols,
        training,
        (float)momentum,
        (float)eps);

    #define LAUNCH(ITEMS) \
        normalize_scale_softmax<ITEMS><<<rows, 256, 0, stream>>>( \
            x.data_ptr<float>(), mean, invstd, \
            weight.data_ptr<float>(), bias.data_ptr<float>(), \
            scale.data_ptr<float>(), output.data_ptr<float>(), \
            cols, scale.stride(0), scale.stride(1))

    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 {
        LAUNCH(32);
    }

    #undef LAUNCH
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_fused_ops = load_inline(
    name="gemm_bn_scale_softmax_cuda_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_bn_scale_softmax"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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):
        x = self.gemm(x)

        # Keep standard autograd behavior when gradients are requested.
        needs_grad = torch.is_grad_enabled() and (
            x.requires_grad
            or self.scale.requires_grad
            or (self.bn.weight is not None and self.bn.weight.requires_grad)
            or (self.bn.bias is not None and self.bn.bias.requires_grad)
        )

        use_fused = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and 0 < x.shape[1] <= 8192
            and x.shape[0] > 1
            and self.bn.affine
            and self.bn.track_running_stats
            and self.scale.dtype == torch.float32
            and self.scale.ndim <= 2
            and not needs_grad
        )

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

        scale = self.scale.expand_as(x)
        momentum = (
            -1.0 if self.bn.momentum is None else self.bn.momentum
        )

        return _fused_ops.fused_bn_scale_softmax(
            x,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.bn.num_batches_tracked,
            scale,
            self.bn.training,
            momentum,
            self.bn.eps,
        )
```