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


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

std::vector<torch::Tensor> bn_scale_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_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 <vector>

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

__device__ __forceinline__ void merge_welford(
    float& mean, float& m2, int& count,
    float other_mean, float other_m2, int other_count) {
    if (other_count == 0) return;
    int total = count + other_count;
    float delta = other_mean - mean;
    m2 += other_m2 + delta * delta *
          (float(count) * float(other_count) / float(total));
    mean += delta * (float(other_count) / float(total));
    count = total;
}

__global__ void partial_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ partial,
    int n, int d) {
    const int col = blockIdx.x * 32 + threadIdx.x;
    const int row_begin = blockIdx.y * ROW_TILE;
    const int row_end = min(row_begin + ROW_TILE, n);
    const int tid = threadIdx.y * 32 + threadIdx.x;

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

    for (int row = row_begin + threadIdx.y;
         row < row_end; row += 8) {
        float value = col < d ? x[(int64_t)row * d + col] : 0.0f;
        ++count;
        float delta = value - mean;
        mean += delta / float(count);
        m2 += delta * (value - mean);
    }

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

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

    for (int stride = 4; stride > 0; stride >>= 1) {
        if (threadIdx.y < stride) {
            int other = tid + stride * 32;
            merge_welford(
                mean, m2, count,
                means[other], m2s[other], counts[other]);
            means[tid] = mean;
            m2s[tid] = m2;
            counts[tid] = count;
        }
        __syncthreads();
    }

    if (threadIdx.y == 0 && col < d) {
        int64_t base = (int64_t)blockIdx.y * 2 * d;
        partial[base + col] = mean;
        partial[base + d + col] = m2;
    }
}

__global__ void finish_stats_kernel(
    const float* __restrict__ partial,
    float* __restrict__ stats,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int n, int d, int chunks,
    float momentum, float eps) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= d) return;

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

    for (int chunk = 0; chunk < chunks; ++chunk) {
        int chunk_count = min(ROW_TILE, n - chunk * ROW_TILE);
        int64_t base = (int64_t)chunk * 2 * d;
        merge_welford(
            mean, m2, count,
            partial[base + col],
            partial[base + d + col],
            chunk_count);
    }

    stats[col] = mean;
    stats[d + col] = rsqrtf(m2 / float(n) + eps);

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

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

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

    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) warp_values[warp] = value;
    __syncthreads();

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

        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) warp_values[0] = value;
    }

    __syncthreads();
    return warp_values[0];
}

template<int ITEMS>
__global__ void fused_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    const float* __restrict__ stats,
    float* __restrict__ output,
    int d, bool scalar_scale) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int64_t base = (int64_t)row * d;
    float values[ITEMS];
    float maximum = -CUDART_INF_F;

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int col = tid + i * SOFTMAX_THREADS;
        float value = -CUDART_INF_F;
        if (col < d) {
            float normalized = (x[base + col] - stats[col]) * stats[d + col];
            float affine = normalized * weight[col] + bias[col];
            value = affine * scale[scalar_scale ? 0 : col];
        }
        values[i] = value;
        maximum = fmaxf(maximum, value);
    }

    maximum = block_reduce<true>(maximum);
    // All threads must finish reading the shared reduction result.
    __syncthreads();

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

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

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

std::vector<torch::Tensor> bn_scale_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) {
    c10::cuda::CUDAGuard device_guard(x.device());

    const int n = static_cast<int>(x.size(0));
    const int d = static_cast<int>(x.size(1));
    auto output = torch::empty_like(x);
    auto stats = torch::empty({2, d}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        int chunks = (n + ROW_TILE - 1) / ROW_TILE;
        auto partial = torch::empty({chunks, 2, d}, x.options());

        partial_stats_kernel<<<
            dim3((d + 31) / 32, chunks), dim3(32, 8), 0, stream>>>(
            x.data_ptr<float>(), partial.data_ptr<float>(), n, d);

        finish_stats_kernel<<<(d + 255) / 256, 256, 0, stream>>>(
            partial.data_ptr<float>(), stats.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            n, d, chunks, static_cast<float>(momentum),
            static_cast<float>(eps));
    } else {
        eval_stats_kernel<<<(d + 255) / 256, 256, 0, stream>>>(
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            stats.data_ptr<float>(), d, static_cast<float>(eps));
    }

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

    #define LAUNCH_SOFTMAX(ITEMS)                                      \
        fused_softmax_kernel<ITEMS><<<n, 256, 0, stream>>>(            \
            x.data_ptr<float>(), weight.data_ptr<float>(),            \
            bias.data_ptr<float>(), scale.data_ptr<float>(),          \
            stats.data_ptr<float>(), output.data_ptr<float>(),        \
            d, scalar_scale)

    if (d <= 1024) {
        LAUNCH_SOFTMAX(4);
    } else if (d <= 2048) {
        LAUNCH_SOFTMAX(8);
    } else if (d <= 4096) {
        LAUNCH_SOFTMAX(16);
    } else {
        LAUNCH_SOFTMAX(32);
    }

    #undef LAUNCH_SOFTMAX
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {output, stats};
}
"""

_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 _BNScaleSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, x, weight, bias, scale,
        running_mean, running_var, training, momentum, eps
    ):
        output, stats = _fused_ops.bn_scale_softmax_cuda(
            x, weight, bias, scale,
            running_mean, running_var, training, momentum, eps
        )
        ctx.training = training
        ctx.save_for_backward(x, weight, bias, scale, output, stats)
        return output

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

        grad_logits = output * (
            grad_output - (grad_output * output).sum(dim=1, keepdim=True)
        )

        grad_scale = None
        if ctx.needs_input_grad[3]:
            affine = normalized * weight + bias
            grad_scale = (grad_logits * affine).sum_to_size(scale.shape)

        grad_affine = grad_logits * scale
        grad_weight = (
            (grad_affine * normalized).sum(dim=0)
            if ctx.needs_input_grad[1] else None
        )
        grad_bias = (
            grad_affine.sum(dim=0)
            if ctx.needs_input_grad[2] else None
        )

        grad_x = None
        if ctx.needs_input_grad[0]:
            grad_normalized = grad_affine * weight
            if ctx.training:
                grad_x = invstd * (
                    grad_normalized
                    - grad_normalized.mean(dim=0)
                    - normalized
                    * (grad_normalized * normalized).mean(dim=0)
                )
            else:
                grad_x = grad_normalized * 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)
        d = x.shape[1] if x.ndim == 2 else 0

        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 x.shape[0] > 0
            and 0 < d <= 8192
            and (not self.bn.training or x.shape[0] > 1)
            and tuple(self.scale.shape)
            in ((), (1,), (1, 1), (d,), (1, d))
            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()

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