```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 bn_scale_softmax_cuda(
    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 <cmath>
#include <cstdint>

constexpr int ROW_CHUNK = 128;

// Each block computes Welford statistics for 128 rows and 32 columns.
__global__ void partial_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ partial_mean,
    float* __restrict__ partial_m2,
    int rows,
    int cols)
{
    const int lane = threadIdx.x & 31;
    const int group = threadIdx.x >> 5;
    const int col = blockIdx.x * 32 + lane;
    const int begin = blockIdx.y * ROW_CHUNK;
    const int end = min(begin + ROW_CHUNK, rows);

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

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

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

    const int tid = threadIdx.x;
    means[tid] = mean;
    m2s[tid] = m2;
    counts[tid] = count;
    __syncthreads();

    for (int step = 4; step > 0; step >>= 1) {
        if (group < 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) {
                float delta = means[other] - means[tid];
                means[tid] += delta * (float(nb) / total);
                m2s[tid] += m2s[other]
                    + delta * delta * (float(na) * nb / total);
            }
            counts[tid] = total;
        }
        __syncthreads();
    }

    if (group == 0 && col < cols) {
        int64_t index = (int64_t)blockIdx.y * cols + col;
        partial_mean[index] = means[tid];
        partial_m2[index] = m2s[tid];
    }
}

__global__ void finish_stats_kernel(
    const float* __restrict__ partial_mean,
    const float* __restrict__ partial_m2,
    float* __restrict__ mean_out,
    float* __restrict__ invstd_out,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    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(ROW_CHUNK, rows - part * ROW_CHUNK);
        int64_t index = (int64_t)part * cols + col;
        float pm = partial_mean[index];
        float delta = pm - mean;
        int total = count + n;

        mean += delta * (float(n) / total);
        m2 += partial_m2[index]
            + delta * delta * (float(count) * n / total);
        count = total;
    }

    mean_out[col] = mean;
    invstd_out[col] = rsqrtf(m2 / rows + eps);

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

__device__ __forceinline__ float warp_max(float value)
{
    for (int offset = 16; offset > 0; offset >>= 1)
        value = fmaxf(value, __shfl_down_sync(0xffffffff, value, offset));
    return value;
}

__device__ __forceinline__ float warp_sum(float value)
{
    for (int offset = 16; offset > 0; offset >>= 1)
        value += __shfl_down_sync(0xffffffff, value, offset);
    return value;
}

template<int ITEMS, bool TRAINING>
__global__ void fused_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ mean,
    const float* __restrict__ variance_or_invstd,
    const float* __restrict__ scale,
    float* __restrict__ output,
    int cols,
    int64_t scale_row_stride,
    int64_t scale_col_stride,
    float eps)
{
    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int64_t base = (int64_t)row * cols;

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

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int col = tid + i * 256;
        float value = -INFINITY;
        if (col < cols) {
            float invstd = TRAINING
                ? variance_or_invstd[col]
                : rsqrtf(variance_or_invstd[col] + eps);

            float normalized = (x[base + col] - mean[col]) * invstd;
            float affine = normalized * weight[col] + bias[col];
            value = affine * scale[
                (int64_t)row * scale_row_stride
                + (int64_t)col * scale_col_stride];
        }
        values[i] = value;
        local_max = fmaxf(local_max, value);
    }

    __shared__ float reduction[8];
    local_max = warp_max(local_max);
    if (lane == 0) reduction[warp] = local_max;
    __syncthreads();

    if (warp == 0) {
        float value = lane < 8 ? reduction[lane] : -INFINITY;
        value = warp_max(value);
        if (lane == 0) reduction[0] = value;
    }
    __syncthreads();

    const float row_max = reduction[0];
    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];
    }

    // All threads must finish reading row_max before reusing reduction.
    __syncthreads();
    local_sum = warp_sum(local_sum);
    if (lane == 0) reduction[warp] = local_sum;
    __syncthreads();

    if (warp == 0) {
        float value = lane < 8 ? reduction[lane] : 0.0f;
        value = warp_sum(value);
        if (lane == 0) reduction[0] = value;
    }
    __syncthreads();

    const float reciprocal_sum = 1.0f / reduction[0];

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

torch::Tensor bn_scale_softmax_cuda(
    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 guard(x.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

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

    torch::Tensor partial, stats;
    const float* mean_ptr = running_mean.data_ptr<float>();
    const float* variance_ptr = running_var.data_ptr<float>();

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

        float* partial_mean = partial.data_ptr<float>();
        float* partial_m2 = partial_mean + (int64_t)parts * cols;
        float* mean = stats.data_ptr<float>();
        float* invstd = mean + cols;

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

        finish_stats_kernel<<<(cols + 255) / 256, 256, 0, stream>>>(
            partial_mean, partial_m2, mean, invstd,
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            rows, cols, parts, static_cast<float>(eps),
            static_cast<float>(momentum));

        mean_ptr = mean;
        variance_ptr = invstd;
    }

    #define LAUNCH(ITEMS, TRAINING) \
        fused_softmax_kernel<ITEMS, TRAINING><<<rows, 256, 0, stream>>>( \
            x.data_ptr<float>(), weight.data_ptr<float>(), \
            bias.data_ptr<float>(), mean_ptr, variance_ptr, \
            scale.data_ptr<float>(), output.data_ptr<float>(), cols, \
            scale.stride(0), scale.stride(1), static_cast<float>(eps))

    #define DISPATCH(ITEMS) \
        do { \
            if (training) { LAUNCH(ITEMS, true); } \
            else { LAUNCH(ITEMS, false); } \
        } while (0)

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

    #undef DISPATCH
    #undef LAUNCH

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""


_fused_ops = load_inline(
    name="gemm_bn_scale_softmax_welford_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["bn_scale_softmax_cuda"],
    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)

        eligible = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and 0 < x.shape[1] <= 16384
            and self.bn.affine
            and self.bn.track_running_stats
            and (not self.bn.training or x.shape[0] > 1)
            and self.scale.ndim <= 2
            and all(
                tensor is not None
                and tensor.dtype == torch.float32
                and tensor.device == x.device
                and tensor.is_contiguous()
                for tensor in (
                    self.bn.weight,
                    self.bn.bias,
                    self.bn.running_mean,
                    self.bn.running_var,
                    self.scale,
                )
            )
        )

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

        scale = self.scale.expand_as(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()

        return _fused_ops.bn_scale_softmax_cuda(
            x,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            scale,
            self.bn.training,
            self.bn.eps,
            0.0 if momentum is None else momentum,
        )
```