```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>

constexpr int TILE_ROWS = 256;

__global__ void partial_stats_kernel(
    const float* __restrict__ x,
    const float* __restrict__ scale,
    float* __restrict__ partial_mean,
    float* __restrict__ partial_m2,
    int rows,
    int features)
{
    const int lane = threadIdx.x;
    const int group = threadIdx.y;
    const int feature = blockIdx.x * 32 + lane;
    const int first_row = blockIdx.y * TILE_ROWS;

    float mean = 0.0f;
    float m2 = 0.0f;
    int count = 0;
    const float s = feature < features ? scale[feature] : 0.0f;

    #pragma unroll
    for (int i = 0; i < TILE_ROWS / 8; ++i) {
        const int row = first_row + group + i * 8;
        if (row < rows && feature < features) {
            const float value = x[(int64_t)row * features + feature] * s;
            ++count;
            const float delta = value - mean;
            mean += delta / float(count);
            m2 += delta * (value - mean);
        }
    }

    __shared__ float means[8][32];
    __shared__ float m2s[8][32];
    __shared__ int counts[8][32];

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

    if (group == 0 && feature < features) {
        float merged_mean = 0.0f;
        float merged_m2 = 0.0f;
        int merged_count = 0;

        #pragma unroll
        for (int g = 0; g < 8; ++g) {
            const int n = counts[g][lane];
            if (n != 0) {
                const int total = merged_count + n;
                const float delta = means[g][lane] - merged_mean;
                merged_mean += delta * (float(n) / float(total));
                merged_m2 += m2s[g][lane]
                    + delta * delta
                    * (float(merged_count) * float(n) / float(total));
                merged_count = total;
            }
        }

        const int64_t index = (int64_t)blockIdx.y * features + feature;
        partial_mean[index] = merged_mean;
        partial_m2[index] = merged_m2;
    }
}

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

    float merged_mean = 0.0f;
    float merged_m2 = 0.0f;
    int merged_count = 0;

    for (int chunk = 0; chunk < chunks; ++chunk) {
        const int remaining = rows - chunk * TILE_ROWS;
        const int n = remaining < TILE_ROWS ? remaining : TILE_ROWS;
        const int total = merged_count + n;
        const int64_t index = (int64_t)chunk * features + feature;
        const float delta = partial_mean[index] - merged_mean;

        merged_mean += delta * (float(n) / float(total));
        merged_m2 += partial_m2[index]
            + delta * delta
            * (float(merged_count) * float(n) / float(total));
        merged_count = total;
    }

    const float variance = fmaxf(merged_m2 / float(rows), 0.0f);
    mean[feature] = merged_mean;
    invstd[feature] = rsqrtf(variance + eps);

    running_mean[feature] =
        (1.0f - momentum) * running_mean[feature]
        + momentum * merged_mean;
    running_var[feature] =
        (1.0f - momentum) * running_var[feature]
        + momentum * (merged_m2 / float(rows - 1));
}

__global__ void eval_stats_kernel(
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    int features,
    float eps)
{
    const int feature = blockIdx.x * blockDim.x + threadIdx.x;
    if (feature < features) {
        mean[feature] = running_mean[feature];
        invstd[feature] = rsqrtf(running_var[feature] + 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 features)
{
    const int64_t start =
        (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t stride = (int64_t)gridDim.x * blockDim.x;

    for (int64_t index = start; index < elements; index += stride) {
        const int feature = index % features;
        const float value = x[index] * scale[feature];
        const float normalized =
            (value - mean[feature]) * invstd[feature];
        output[index] =
            normalized * weight[feature] + bias[feature];
    }
}

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 auto stream = at::cuda::getCurrentCUDAStream();

    const int rows = static_cast<int>(x.size(0));
    const int features = static_cast<int>(x.size(1));
    auto output = torch::empty_like(x);
    auto mean = torch::empty({features}, x.options());
    auto invstd = torch::empty({features}, x.options());

    if (training) {
        const int chunks = (rows + TILE_ROWS - 1) / TILE_ROWS;
        auto partial_mean = torch::empty({chunks, features}, x.options());
        auto partial_m2 = torch::empty({chunks, features}, x.options());

        partial_stats_kernel<<<
            dim3((features + 31) / 32, chunks), dim3(32, 8), 0, stream>>>(
                x.data_ptr<float>(),
                scale.data_ptr<float>(),
                partial_mean.data_ptr<float>(),
                partial_m2.data_ptr<float>(),
                rows, features);

        finalize_stats_kernel<<<(features + 255) / 256, 256, 0, stream>>>(
            partial_mean.data_ptr<float>(),
            partial_m2.data_ptr<float>(),
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            rows, features, chunks,
            static_cast<float>(momentum),
            static_cast<float>(eps));
    } else {
        eval_stats_kernel<<<(features + 255) / 256, 256, 0, stream>>>(
            running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(),
            mean.data_ptr<float>(),
            invstd.data_ptr<float>(),
            features, static_cast<float>(eps));
    }

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

    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>(),
        elements, features);

    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_bias = grad_output.sum(dim=0)
        grad_weight = (grad_output * normalized).sum(dim=0)
        grad_normalized = grad_output * weight

        if ctx.training:
            grad_scaled = (
                grad_normalized
                - grad_normalized.mean(dim=0)
                - normalized
                * (grad_normalized * normalized).mean(dim=0)
            ) * invstd
        else:
            grad_scaled = grad_normalized * invstd

        grad_x = grad_scaled * scale
        grad_scale = (grad_scaled * x).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)

        tensors = (
            self.scale, self.bn.weight, self.bn.bias,
            self.bn.running_mean, self.bn.running_var
        )
        use_custom = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and x.shape[0] > 1
            and x.shape[1] > 0
            and (x.shape[0] + 255) // 256 <= 65535
            and self.scale.shape == (x.shape[1],)
            and self.bn.track_running_stats
            and self.bn.affine
            and all(
                t is not None
                and t.dtype == torch.float32
                and t.device == x.device
                and t.is_contiguous()
                for t in tensors
            )
        )

        if not use_custom:
            return self.bn(x * self.scale)

        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 _ScaledBatchNorm.apply(
            x, self.scale, self.bn.weight, self.bn.bias,
            self.bn.running_mean, self.bn.running_var,
            self.bn.training,
            0.0 if momentum is None else momentum,
            self.bn.eps,
        )
```