```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,
    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 <vector>
#include <cmath>

__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 eps,
    float momentum)
{
    const int col = blockIdx.x * 32 + threadIdx.x;
    const int lane = threadIdx.y;
    const int tid = lane * 32 + threadIdx.x;

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

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

    if (col < cols) {
        for (int row = lane; row < rows; row += 8) {
            float value = x[(int64_t)row * cols + col];
            ++count;
            float delta = value - m;
            m += delta / count;
            m2 += delta * (value - m);
        }
    }

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

    if (lane == 0 && col < cols) {
        m = means[threadIdx.x];
        m2 = m2s[threadIdx.x];
        count = counts[threadIdx.x];

        #pragma unroll
        for (int k = 1; k < 8; ++k) {
            int idx = k * 32 + threadIdx.x;
            int other_count = counts[idx];
            if (other_count > 0) {
                int total = count + other_count;
                float delta = means[idx] - m;
                m2 += m2s[idx]
                    + delta * delta * ((float)count * other_count / total);
                m += delta * ((float)other_count / total);
                count = total;
            }
        }

        float variance = fmaxf(m2 / rows, 0.0f);
        mean[col] = m;
        invstd[col] = rsqrtf(variance + eps);

        if (running_mean != nullptr) {
            running_mean[col] =
                (1.0f - momentum) * running_mean[col] + momentum * m;
            running_var[col] =
                (1.0f - momentum) * running_var[col]
                + momentum * (m2 / (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)
{
    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) {
    __shared__ float warp_values[8];
    __shared__ float result;

    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, value, offset);
        value = Maximum ? fmaxf(value, other) : value + other;
    }

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

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

        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            float other = __shfl_down_sync(0xffffffff, value, offset);
            value = Maximum ? fmaxf(value, other) : value + other;
        }

        if (lane == 0)
            result = value;
    }
    __syncthreads();
    return result;
}

template<int Items>
__global__ void fused_softmax_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ scale,
    const float* __restrict__ mean,
    const float* __restrict__ invstd,
    float* __restrict__ output,
    int cols,
    int64_t scale_row_stride,
    int64_t scale_col_stride)
{
    const int row = blockIdx.x;
    const int64_t row_offset = (int64_t)row * cols;
    const int64_t scale_offset = (int64_t)row * scale_row_stride;

    float values[Items];
    float maximum = -INFINITY;

    #pragma unroll
    for (int i = 0; i < Items; ++i) {
        int col = threadIdx.x + i * 256;
        float value = -INFINITY;
        if (col < cols) {
            float normalized = (x[row_offset + col] - mean[col]) * invstd[col];
            float affine = normalized * gamma[col] + beta[col];
            value = affine * scale[scale_offset + col * scale_col_stride];
        }
        values[i] = value;
        maximum = fmaxf(maximum, value);
    }

    maximum = block_reduce<true>(maximum);

    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 reciprocal = 1.0f / sum;

    #pragma unroll
    for (int i = 0; i < Items; ++i) {
        int col = threadIdx.x + 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,
    bool training,
    double eps,
    double momentum)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat,
                "Expected a CUDA float32 tensor");
    TORCH_CHECK(x.dim() == 2 && x.is_contiguous(),
                "Expected a contiguous matrix");

    const c10::cuda::CUDAGuard guard(x.device());
    const int rows = x.size(0);
    const int cols = x.size(1);
    TORCH_CHECK(rows > 0 && cols > 0 && cols <= 16384,
                "Unsupported matrix dimensions");
    TORCH_CHECK(!training || rows > 1,
                "Training BatchNorm requires more than one sample");

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        const bool track = running_mean.numel() != 0;
        batch_stats_kernel<<<(cols + 31) / 32, dim3(32, 8), 0, stream>>>(
            x.data_ptr<float>(),
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            track ? running_mean.data_ptr<float>() : nullptr,
            track ? running_var.data_ptr<float>() : nullptr,
            rows, cols, static_cast<float>(eps), static_cast<float>(momentum));
    } 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, static_cast<float>(eps));
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    #define LAUNCH(ITEMS) \
        fused_softmax_kernel<ITEMS><<<rows, 256, 0, stream>>>( \
            x.data_ptr<float>(), gamma.data_ptr<float>(), \
            beta.data_ptr<float>(), expanded_scale.data_ptr<float>(), \
            mean.data_ptr<float>(), invstd.data_ptr<float>(), \
            output.data_ptr<float>(), cols, \
            expanded_scale.stride(0), expanded_scale.stride(1))

    if (cols <= 1024) {
        LAUNCH(4);
    } else if (cols <= 4096) {
        LAUNCH(16);
    } else if (cols <= 8192) {
        LAUNCH(32);
    } else {
        LAUNCH(64);
    }
    #undef LAUNCH

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


_fused_extension = load_inline(
    name="fused_bn_scale_softmax_cuda_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,
        training, eps, momentum
    ):
        output, mean, invstd = _fused_extension.bn_scale_softmax_cuda(
            x, gamma, beta, scale, running_mean, running_var,
            training, eps, momentum
        )
        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_scale = None
        if ctx.needs_input_grad[3]:
            grad_scale = (
                grad_logits * (normalized * gamma + beta)
            ).sum_to_size(scale.shape)

        grad_gamma = None
        if ctx.needs_input_grad[1]:
            grad_gamma = (grad_affine * normalized).sum(dim=0)

        grad_beta = None
        if ctx.needs_input_grad[2]:
            grad_beta = grad_affine.sum(dim=0)

        grad_x = None
        if ctx.needs_input_grad[0]:
            if ctx.training:
                grad_x = (
                    grad_affine
                    - grad_affine.mean(dim=0)
                    - normalized * (grad_affine * normalized).mean(dim=0)
                ) * (gamma * invstd)
            else:
                grad_x = grad_affine * (gamma * invstd)

        return (
            grad_x, grad_gamma, grad_beta, 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)

        use_custom = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.shape[0] > 0
            and 0 < x.shape[1] <= 16384
            and self.scale.dtype == torch.float32
            and self.scale.ndim <= 2
            and self.bn.affine
            and self.bn.weight.dtype == torch.float32
            and self.bn.bias.dtype == torch.float32
        )
        if not use_custom:
            return self.softmax(self.scale * self.bn(x))

        training = self.bn.training or (
            self.bn.running_mean is None and self.bn.running_var is None
        )
        if training and x.shape[0] <= 1:
            return self.softmax(self.scale * self.bn(x))

        momentum = (
            self.bn.momentum if self.bn.momentum is not None else 0.0
        )
        if self.bn.training and self.bn.track_running_stats:
            self.bn.num_batches_tracked.add_(1)
            if self.bn.momentum is None:
                momentum = 1.0 / self.bn.num_batches_tracked.item()

        running_mean = self.bn.running_mean
        running_var = self.bn.running_var
        if running_mean is None:
            running_mean = x.new_empty(0)
            running_var = x.new_empty(0)

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