```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 gamma,
    torch::Tensor beta,
    torch::Tensor scale,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor num_batches,
    double eps,
    double momentum,
    bool training);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <vector>
#include <cstdint>

__global__ void training_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int64_t* __restrict__ num_batches,
    int rows,
    int cols,
    float eps,
    float momentum)
{
    __shared__ float scratch[256];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int group = tid >> 5;
    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];
        }
    }

    scratch[tid] = sum;
    __syncthreads();

    if (group == 0) {
        float total = 0.0f;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            total += scratch[i * 32 + lane];
        }
        const float m = total / rows;
        scratch[lane] = m;
        if (col < cols) {
            mean[col] = m;
        }
    }
    __syncthreads();

    const float m = scratch[lane];
    __syncthreads();

    float squared_sum = 0.0f;
    if (col < cols) {
        for (int row = group; row < rows; row += 8) {
            const float d = x[(int64_t)row * cols + col] - m;
            squared_sum = fmaf(d, d, squared_sum);
        }
    }

    scratch[tid] = squared_sum;
    __syncthreads();

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

        const float variance = total / rows;
        invstd[col] = rsqrtf(variance + eps);
        running_mean[col] =
            (1.0f - momentum) * running_mean[col] + momentum * m;
        running_var[col] =
            (1.0f - momentum) * running_var[col]
            + momentum * (total / (rows - 1));
    }

    if (blockIdx.x == 0 && tid == 0) {
        num_batches[0] += 1;
    }
}

__global__ void evaluation_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);
    }
}

template <bool Maximum>
__device__ __forceinline__ float block_reduce(float value, float* scratch)
{
    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) {
        scratch[warp] = value;
    }
    __syncthreads();

    if (warp == 0) {
        value = lane < 8
            ? scratch[lane]
            : (Maximum ? -CUDART_INF_F : 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) {
            scratch[0] = value;
        }
    }
    __syncthreads();

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

template <int Items>
__global__ void bn_scale_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ mean,
    const float* __restrict__ invstd,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ scale,
    float* __restrict__ output,
    int cols)
{
    __shared__ float scratch[8];

    const int tid = threadIdx.x;
    const int64_t row_offset = (int64_t)blockIdx.x * cols;
    const float s = scale[0];

    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) {
            const float normalized =
                (x[row_offset + col] - mean[col]) * invstd[col];
            value = fmaf(normalized, gamma[col], beta[col]) * s;
        }
        values[i] = value;
        local_max = fmaxf(local_max, value);
    }

    const float maximum = block_reduce<true>(local_max, scratch);
    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 denominator = block_reduce<false>(local_sum, scratch);
    const float reciprocal = 1.0f / denominator;

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

std::vector<torch::Tensor> bn_scale_softmax_cuda(
    torch::Tensor x,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor scale,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor num_batches,
    double eps,
    double momentum,
    bool training)
{
    const c10::cuda::CUDAGuard device_guard(x.device());
    const auto stream = c10::cuda::getCurrentCUDAStream().stream();

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

    auto output = torch::empty_like(x);
    auto mean = torch::empty({cols}, x.options());
    auto invstd = torch::empty({cols}, x.options());

    if (training) {
        training_stats_kernel<<<(cols + 31) / 32, 256, 0, stream>>>(
            x.data_ptr<float>(),
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            num_batches.data_ptr<int64_t>(),
            rows, cols,
            static_cast<float>(eps),
            static_cast<float>(momentum));
    } else {
        evaluation_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,
            static_cast<float>(eps));
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    #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>(), \
            gamma.data_ptr<float>(), \
            beta.data_ptr<float>(), \
            scale.data_ptr<float>(), \
            output.data_ptr<float>(), \
            cols)

    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 {output, mean, invstd};
}
"""

_fused_ops = load_inline(
    name="gemm_bn_scale_softmax_fused_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["bn_scale_softmax_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _BatchNormScaleSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, x, gamma, beta, scale, running_mean, running_var,
        num_batches, eps, momentum, training
    ):
        output, mean, invstd = _fused_ops.bn_scale_softmax_cuda(
            x, gamma, beta, scale, running_mean, running_var,
            num_batches, eps, momentum, training
        )
        ctx.training = training
        ctx.save_for_backward(x, gamma, beta, scale, output, mean, invstd)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        x, gamma, beta, scale, output, mean, invstd = ctx.saved_tensors

        grad_logits = output * (
            grad_output - (grad_output * output).sum(dim=1, keepdim=True)
        )
        normalized = (x - mean) * invstd
        grad_affine = grad_logits * scale
        grad_normalized = grad_affine * gamma

        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

        grad_gamma = (grad_affine * normalized).sum(dim=0)
        grad_beta = grad_affine.sum(dim=0)
        grad_scale = (
            grad_logits * (normalized * gamma + beta)
        ).sum().reshape_as(scale)

        return (
            grad_x, grad_gamma, grad_beta, grad_scale,
            None, 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)
        bn = self.bn

        tensors = (
            bn.weight, bn.bias, self.scale,
            bn.running_mean, 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 bn.training else 0)
            and self.scale.numel() == 1
            and self.scale.ndim <= 2
            and bn.momentum is not None
            and bn.num_batches_tracked is not None
            and bn.num_batches_tracked.device == x.device
            and bn.num_batches_tracked.dtype == torch.int64
            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 * bn(x))

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