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

_CPP = """
#include <torch/extension.h>
#include <vector>

std::vector<torch::Tensor> fused_bn_softmax_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor scale,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    bool training,
    double momentum,
    double eps);
"""

_CUDA = 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 <vector>

__global__ void batch_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int rows,
    int cols,
    float momentum,
    float eps)
{
    // Each warp reads adjacent features; eight warps cooperate on each
    // feature's reduction across the batch.
    __shared__ float partial[8][32];
    __shared__ float column_mean[32];

    const int lane = threadIdx.x;
    const int group = threadIdx.y;
    const int col = blockIdx.x * 32 + lane;

    float sum = 0.0f;
    if (col < cols) {
        for (int row = group; row < rows; row += 8) {
            sum += x[(int64_t)row * cols + col];
        }
    }
    partial[group][lane] = sum;
    __syncthreads();

    if (group == 0) {
        float total = 0.0f;
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            total += partial[k][lane];
        }
        column_mean[lane] = total / rows;
    }
    __syncthreads();

    const float mu = column_mean[lane];
    float squared_sum = 0.0f;
    if (col < cols) {
        for (int row = group; row < rows; row += 8) {
            const float delta = x[(int64_t)row * cols + col] - mu;
            squared_sum += delta * delta;
        }
    }
    partial[group][lane] = squared_sum;
    __syncthreads();

    if (group == 0 && col < cols) {
        float total = 0.0f;
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            total += partial[k][lane];
        }

        mean[col] = mu;
        invstd[col] = rsqrtf(total / rows + eps);

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

__global__ void eval_stats_kernel(
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    int cols,
    float eps)
{
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col < cols) {
        mean[col] = running_mean[col];
        invstd[col] = rsqrtf(running_var[col] + eps);
    }
}

__device__ __forceinline__ float warp_max(float value) {
    #pragma unroll
    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) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        value += __shfl_down_sync(0xffffffff, value, offset);
    }
    return value;
}

template<bool MAXIMUM>
__device__ __forceinline__ float block_reduce(
    float value, float* scratch)
{
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    __syncthreads();
    value = MAXIMUM ? warp_max(value) : warp_sum(value);
    if (lane == 0) {
        scratch[warp] = value;
    }
    __syncthreads();

    if (warp == 0) {
        value = lane < 8
            ? scratch[lane]
            : (MAXIMUM ? -CUDART_INF_F : 0.0f);
        value = MAXIMUM ? warp_max(value) : warp_sum(value);
        if (lane == 0) {
            scratch[0] = value;
        }
    }
    __syncthreads();
    return scratch[0];
}

template<int ITEMS>
__global__ void bn_scale_softmax_kernel(
    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__ out,
    int cols,
    bool scalar_scale)
{
    __shared__ float scratch[8];

    const int64_t row_offset = (int64_t)blockIdx.x * cols;
    float values[ITEMS];
    float maximum = -CUDART_INF_F;

    #pragma unroll
    for (int k = 0; k < ITEMS; ++k) {
        const int col = threadIdx.x + k * 256;
        float value = -CUDART_INF_F;
        if (col < cols) {
            const float normalized =
                (x[row_offset + col] - mean[col]) * invstd[col];
            value = (normalized * weight[col] + bias[col])
                * scale[scalar_scale ? 0 : col];
        }
        values[k] = value;
        maximum = fmaxf(maximum, value);
    }

    maximum = block_reduce<true>(maximum, scratch);

    float denominator = 0.0f;
    #pragma unroll
    for (int k = 0; k < ITEMS; ++k) {
        values[k] = expf(values[k] - maximum);
        denominator += values[k];
    }

    denominator = block_reduce<false>(denominator, scratch);
    const float reciprocal = 1.0f / denominator;

    #pragma unroll
    for (int k = 0; k < ITEMS; ++k) {
        const int col = threadIdx.x + k * 256;
        if (col < cols) {
            out[row_offset + col] = values[k] * reciprocal;
        }
    }
}

std::vector<torch::Tensor> fused_bn_softmax_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor scale,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    bool training,
    double momentum,
    double eps)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kFloat32,
                "Expected a CUDA float32 input");
    TORCH_CHECK(x.dim() == 2 && x.is_contiguous(),
                "Expected a contiguous matrix");

    const c10::cuda::CUDAGuard device_guard(x.device());
    const int rows = x.size(0);
    const int cols = x.size(1);

    TORCH_CHECK(rows > 0 && cols > 0 && cols <= 8192,
                "Unsupported input dimensions");
    TORCH_CHECK(!training || rows > 1,
                "Training batch normalization requires at least two rows");

    auto out = torch::empty_like(x);
    auto mean = torch::empty({cols}, x.options());
    auto invstd = torch::empty({cols}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        batch_stats_kernel<<<(cols + 31) / 32, dim3(32, 8), 0, stream>>>(
            x.data_ptr<float>(),
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            rows, cols, (float)momentum, (float)eps);
    } else {
        eval_stats_kernel<<<(cols + 255) / 256, 256, 0, stream>>>(
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            cols, (float)eps);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

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

    #define LAUNCH_SOFTMAX(ITEMS)                                      \
        bn_scale_softmax_kernel<ITEMS><<<rows, 256, 0, stream>>>(       \
            x.data_ptr<float>(), mean.data_ptr<float>(),              \
            invstd.data_ptr<float>(), weight.data_ptr<float>(),       \
            bias.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 {
        LAUNCH_SOFTMAX(32);
    }

    #undef LAUNCH_SOFTMAX

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {out, mean, invstd};
}
"""

_fused_ops = load_inline(
    name="fused_batchnorm_scale_softmax_v1",
    cpp_sources=_CPP,
    cuda_sources=_CUDA,
    functions=["fused_bn_softmax_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--fmad=false"],
    verbose=False,
)


class _FusedBatchNormScaleSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, x, weight, bias, scale,
        running_mean, running_var, training, momentum, eps
    ):
        out, mean, invstd = _fused_ops.fused_bn_softmax_cuda(
            x, weight, bias, scale,
            running_mean, running_var, training, momentum, eps
        )
        ctx.training = training
        ctx.save_for_backward(x, weight, bias, scale, out, mean, invstd)
        return out

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, bias, scale, out, mean, invstd = ctx.saved_tensors
        normalized = (x - mean) * invstd

        grad_logits = out * (
            grad_output - (grad_output * out).sum(dim=1, keepdim=True)
        )
        grad_affine = grad_logits * scale

        grad_scale = (
            grad_logits * (normalized * weight + bias)
        ).sum_to_size(scale.shape)
        grad_weight = (grad_affine * normalized).sum(dim=0)
        grad_bias = grad_affine.sum(dim=0)

        if ctx.training:
            grad_x = (
                grad_affine
                - grad_affine.mean(dim=0)
                - normalized * (grad_affine * normalized).mean(dim=0)
            ) * (weight * invstd)
        else:
            grad_x = grad_affine * (weight * invstd)

        return (
            grad_x, grad_weight, grad_bias, grad_scale,
            None, None, None, None, None
        )


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)

        feature_scale = (
            self.scale.ndim in (1, 2)
            and self.scale.shape[-1] == x.shape[-1]
            and self.scale.numel() == x.shape[-1]
        )
        scalar_scale = self.scale.ndim <= 2 and self.scale.numel() == 1

        tensors = (
            self.bn.weight,
            self.bn.bias,
            self.scale,
            self.bn.running_mean,
            self.bn.running_var,
        )
        supported = (
            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 if self.bn.training else 0)
            and (scalar_scale or feature_scale)
            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 / float(self.bn.num_batches_tracked)
        elif momentum is None:
            momentum = 0.0

        return _FusedBatchNormScaleSoftmax.apply(
            x,
            self.bn.weight,
            self.bn.bias,
            self.scale,
            self.bn.running_mean,
            self.bn.running_var,
            self.bn.training,
            momentum,
            self.bn.eps,
        )
```