```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 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 <cuda_runtime.h>
#include <math.h>

constexpr int ROW_TILE = 128;
constexpr int SOFTMAX_THREADS = 256;

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

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

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

    if (col < cols) {
        for (int row = start + ty; row < end; row += 8) {
            const float value = x[(int64_t)row * cols + col];
            ++count;
            const float delta = value - mean;
            mean += delta / count;
            m2 += delta * (value - 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) {
            const int n = max(0, (end - start - lane + 7) / 8);
            if (n > 0) {
                const int combined = total_count + n;
                const 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;
            }
        }

        const 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_kernel(
    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)
{
    const 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) {
        const int n = min(ROW_TILE, rows - part * ROW_TILE);
        const int combined = count + n;
        const int64_t index = (int64_t)part * cols + col;
        const float part_mean = partial[index];
        const float part_m2 = partial[(int64_t)parts * cols + index];
        const float delta = part_mean - mean;

        m2 += part_m2 + delta * delta * ((float)count * n / combined);
        mean += delta * ((float)n / combined);
        count = combined;
    }

    const float variance = fmaxf(m2 / rows, 0.0f);
    const float alpha = weight[col] * rsqrtf(variance + eps);

    coefficients[col] = alpha;
    coefficients[cols + col] = bias[col] - mean * alpha;

    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_kernel(
    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)
{
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

    const float alpha = weight[col] * rsqrtf(running_var[col] + eps);
    coefficients[col] = alpha;
    coefficients[cols + col] = bias[col] - running_mean[col] * alpha;
}

template<bool MAXIMUM>
__device__ __forceinline__ float block_reduce(float value)
{
    __shared__ float warp_results[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);
        if (MAXIMUM) value = fmaxf(value, other);
        else value += other;
    }

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

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

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

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

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

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

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

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        const int col = tid + i * SOFTMAX_THREADS;
        float value = -INFINITY;
        if (col < cols) {
            value = (
                x[row_offset + col] * coefficients[col]
                + coefficients[cols + col]
            ) * scale[scalar_scale ? 0 : col];
        }
        values[i] = value;
        local_max = fmaxf(local_max, value);
    }

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

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        values[i] = expf(values[i] - maximum);
        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 * SOFTMAX_THREADS;
        if (col < cols) {
            out[row_offset + 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 scale,
    bool training,
    double eps,
    double momentum)
{
    const c10::cuda::CUDAGuard device_guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int rows = x.size(0);
    const int cols = x.size(1);

    auto out = torch::empty_like(x);
    if (rows == 0) return out;

    auto coefficients = torch::empty({2, cols}, x.options());

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

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

        training_coefficients_kernel<<<
            (cols + 255) / 256, 256, 0, stream
        >>>(
            partial.data_ptr<float>(),
            weight.data_ptr<float>(),
            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_kernel<<<
            (cols + 255) / 256, 256, 0, stream
        >>>(
            weight.data_ptr<float>(),
            bias.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            coefficients.data_ptr<float>(),
            cols, (float)eps
        );
    }

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

    #define LAUNCH_SOFTMAX(N) \
        normalized_softmax_kernel<N><<<rows, SOFTMAX_THREADS, 0, stream>>>( \
            x.data_ptr<float>(), coefficients.data_ptr<float>(), \
            scale.data_ptr<float>(), out.data_ptr<float>(), \
            cols, scalar_scale)

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

    #undef LAUNCH_SOFTMAX

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_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)
        cols = self.gemm.out_features

        supported_scale = tuple(self.scale.shape) in {
            (), (1,), (cols,), (1, 1), (1, cols)
        }

        tensors = (
            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 0 < cols <= 16384
            and supported_scale
            and self.bn.affine
            and self.bn.track_running_stats
            and (not self.bn.training or x.size(0) > 1)
            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(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 _fused_ops.fused_bn_scale_softmax(
            x,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.scale,
            self.bn.training,
            self.bn.eps,
            momentum,
        )
```