```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> scaled_bn_cuda(
    torch::Tensor x,
    torch::Tensor scale,
    torch::Tensor weight,
    torch::Tensor bias,
    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_runtime.h>
#include <vector>
#include <algorithm>

constexpr int ROWS_PER_CHUNK = 256;

__global__ void partial_moments_kernel(
    const float* __restrict__ x,
    const float* __restrict__ scale,
    float* __restrict__ partial_mean,
    float* __restrict__ partial_m2,
    int rows,
    int cols)
{
    __shared__ float sm[256];
    __shared__ float sv[256];
    __shared__ int sn[256];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int group = tid >> 5;
    const int col = blockIdx.x * 32 + lane;
    const int start = blockIdx.y * ROWS_PER_CHUNK;
    const int end = min(start + ROWS_PER_CHUNK, rows);

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

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

    sm[tid] = mean;
    sv[tid] = m2;
    sn[tid] = count;
    __syncthreads();

    for (int stride = 4; stride > 0; stride >>= 1) {
        if (group < stride) {
            const int other = tid + stride * 32;
            const int na = sn[tid];
            const int nb = sn[other];
            const int total = na + nb;

            if (nb > 0) {
                const float delta = sm[other] - sm[tid];
                const float fraction = float(nb) / float(total);
                sv[tid] += sv[other] +
                           delta * delta * float(na) * fraction;
                sm[tid] += delta * fraction;
                sn[tid] = total;
            }
        }
        __syncthreads();
    }

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

__global__ void finalize_moments_kernel(
    const float* __restrict__ partial_mean,
    const float* __restrict__ partial_m2,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int rows,
    int cols,
    int chunks,
    float momentum,
    float eps)
{
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

    float sum = 0.0f;
    for (int chunk = 0; chunk < chunks; ++chunk) {
        const int count = min(ROWS_PER_CHUNK,
                              rows - chunk * ROWS_PER_CHUNK);
        sum += partial_mean[(int64_t)chunk * cols + col] * float(count);
    }

    const float mu = sum / float(rows);
    float m2 = 0.0f;

    for (int chunk = 0; chunk < chunks; ++chunk) {
        const int count = min(ROWS_PER_CHUNK,
                              rows - chunk * ROWS_PER_CHUNK);
        const int64_t index = (int64_t)chunk * cols + col;
        const float delta = partial_mean[index] - mu;
        m2 += partial_m2[index] + delta * delta * float(count);
    }

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

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

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

__global__ void normalize_kernel(
    const float* __restrict__ x,
    const float* __restrict__ scale,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ mean,
    const float* __restrict__ invstd,
    float* __restrict__ output,
    int64_t elements,
    int cols)
{
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
         i < elements; i += stride) {
        const int col = i % cols;
        const float scaled = __fmul_rn(x[i], scale[col]);
        const float normalized = (scaled - mean[col]) * invstd[col];
        output[i] = normalized * weight[col] + bias[col];
    }
}

std::vector<torch::Tensor> scaled_bn_cuda(
    torch::Tensor x,
    torch::Tensor scale,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    bool training,
    double momentum,
    double eps)
{
    const c10::cuda::CUDAGuard device_guard(x.device());
    const int rows = static_cast<int>(x.size(0));
    const int cols = static_cast<int>(x.size(1));

    TORCH_CHECK(!training || rows > 1,
                "Batch normalization requires more than one training sample");

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

    if (training) {
        const int chunks = (rows + ROWS_PER_CHUNK - 1) / ROWS_PER_CHUNK;
        auto partial = torch::empty({2, chunks, cols}, x.options());
        float* partial_mean = partial.data_ptr<float>();
        float* partial_m2 = partial_mean + (int64_t)chunks * cols;

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

        finalize_moments_kernel<<<(cols + 255) / 256, 256, 0, stream>>>(
            partial_mean,
            partial_m2,
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            rows,
            cols,
            chunks,
            static_cast<float>(momentum),
            static_cast<float>(eps));
    } else {
        inference_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));
    }

    const int blocks = static_cast<int>(
        std::min<int64_t>((x.numel() + 255) / 256, 4096));

    normalize_kernel<<<blocks, 256, 0, stream>>>(
        x.data_ptr<float>(),
        scale.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        mean.data_ptr<float>(),
        invstd.data_ptr<float>(),
        output.data_ptr<float>(),
        x.numel(),
        cols);

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

_scaled_bn = load_inline(
    name="fused_scaled_batchnorm_welford",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["scaled_bn_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _ScaledBatchNorm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, scale, weight, bias,
                running_mean, running_var, training, momentum, eps):
        output, mean, invstd = _scaled_bn.scaled_bn_cuda(
            x, scale, weight, bias,
            running_mean, running_var, training, momentum, eps
        )
        ctx.training = training
        ctx.save_for_backward(x, scale, weight, mean, invstd)
        return output

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

        grad_x = grad_scale = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
            grad_scaled = grad_output * weight
            if ctx.training:
                grad_scaled = (
                    grad_scaled
                    - grad_scaled.mean(dim=0)
                    - normalized
                    * (grad_scaled * normalized).mean(dim=0)
                )
            grad_scaled = grad_scaled * invstd

            if ctx.needs_input_grad[0]:
                grad_x = grad_scaled * scale
            if ctx.needs_input_grad[1]:
                grad_scale = (grad_scaled * x).sum(dim=0).reshape_as(scale)

        if ctx.needs_input_grad[2]:
            grad_weight = (grad_output * normalized).sum(dim=0)
        if ctx.needs_input_grad[3]:
            grad_bias = grad_output.sum(dim=0)

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


class ModelNew(nn.Module):
    def __init__(self, in_features, out_features, scale_shape,
                 eps=1e-5, momentum=0.1):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.scale = nn.Parameter(torch.randn(scale_shape))
        self.bn = nn.BatchNorm1d(
            out_features, eps=eps, momentum=momentum
        )

    def forward(self, x):
        x = self.gemm(x)
        bn = self.bn

        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and x.numel() > 0
            and tuple(self.scale.shape) in (
                (x.shape[1],), (1, x.shape[1])
            )
            and bn.affine
            and bn.track_running_stats
        )

        if supported:
            tensors = (
                self.scale, bn.weight, bn.bias,
                bn.running_mean, bn.running_var
            )
            supported = 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 bn(x * self.scale)

        if bn.training and x.shape[0] <= 1:
            return bn(x * self.scale)

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

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