```python
import torch
import torch.nn as nn
import torch.nn.functional as F
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 input,
    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 ROWS_PER_PART = 256;

__global__ void partial_stats_kernel(
    const float* __restrict__ input,
    const float* __restrict__ scale,
    float* __restrict__ sums,
    float* __restrict__ squares,
    int rows,
    int channels)
{
    const int lane = threadIdx.x & 31;
    const int row_lane = threadIdx.x >> 5;
    const int channel = blockIdx.x * 32 + lane;
    const int first_row = blockIdx.y * ROWS_PER_PART;

    float sum = 0.0f;
    float square = 0.0f;
    if (channel < channels) {
        const float s = scale[channel];
        #pragma unroll
        for (int i = 0; i < ROWS_PER_PART / 8; ++i) {
            const int row = first_row + row_lane + i * 8;
            if (row < rows) {
                const float v = input[(int64_t)row * channels + channel] * s;
                sum += v;
                square = fmaf(v, v, square);
            }
        }
    }

    __shared__ float shared_sum[256];
    __shared__ float shared_square[256];
    shared_sum[threadIdx.x] = sum;
    shared_square[threadIdx.x] = square;
    __syncthreads();

    if (row_lane == 0 && channel < channels) {
        float total_sum = 0.0f;
        float total_square = 0.0f;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            total_sum += shared_sum[i * 32 + lane];
            total_square += shared_square[i * 32 + lane];
        }
        const int index = blockIdx.y * channels + channel;
        sums[index] = total_sum;
        squares[index] = total_square;
    }
}

__global__ void finish_stats_kernel(
    const float* __restrict__ sums,
    const float* __restrict__ squares,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    int rows,
    int channels,
    int parts,
    float momentum,
    float eps,
    bool training)
{
    const int channel = blockIdx.x * blockDim.x + threadIdx.x;
    if (channel >= channels)
        return;

    float m;
    float variance;

    if (training) {
        // Double precision here reduces cancellation in the final reduction.
        double sum = 0.0;
        double square = 0.0;
        for (int part = 0; part < parts; ++part) {
            const int index = part * channels + channel;
            sum += static_cast<double>(sums[index]);
            square += static_cast<double>(squares[index]);
        }

        const double average = sum / rows;
        m = static_cast<float>(average);
        variance = static_cast<float>(
            fmax(square / rows - average * average, 0.0));

        if (running_mean != nullptr) {
            running_mean[channel] =
                (1.0f - momentum) * running_mean[channel] + momentum * m;
            const float unbiased = variance * (float(rows) / float(rows - 1));
            running_var[channel] =
                (1.0f - momentum) * running_var[channel] + momentum * unbiased;
        }
    } else {
        m = running_mean[channel];
        variance = running_var[channel];
    }

    mean[channel] = m;
    invstd[channel] = rsqrtf(variance + eps);
}

__global__ void normalize_kernel(
    const float* __restrict__ input,
    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 count,
    int channels)
{
    int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (; index < count; index += stride) {
        const int channel = index % channels;
        const float scaled = input[index] * scale[channel];
        const float normalized = (scaled - mean[channel]) * invstd[channel];
        output[index] = normalized * weight[channel] + bias[channel];
    }
}

std::vector<torch::Tensor> scaled_bn_cuda(
    torch::Tensor input,
    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 guard(input.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int rows = input.size(0);
    const int channels = input.size(1);
    const int parts = (rows + ROWS_PER_PART - 1) / ROWS_PER_PART;

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

    torch::Tensor sums;
    torch::Tensor squares;
    const float* sums_ptr = nullptr;
    const float* squares_ptr = nullptr;

    if (training) {
        sums = torch::empty({parts, channels}, input.options());
        squares = torch::empty_like(sums);
        sums_ptr = sums.data_ptr<float>();
        squares_ptr = squares.data_ptr<float>();

        partial_stats_kernel<<<
            dim3((channels + 31) / 32, parts), 256, 0, stream>>>(
                input.data_ptr<float>(),
                scale.data_ptr<float>(),
                sums.data_ptr<float>(),
                squares.data_ptr<float>(),
                rows,
                channels);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    float* running_mean_ptr = running_mean.numel()
        ? running_mean.data_ptr<float>() : nullptr;
    float* running_var_ptr = running_var.numel()
        ? running_var.data_ptr<float>() : nullptr;

    finish_stats_kernel<<<(channels + 255) / 256, 256, 0, stream>>>(
        sums_ptr,
        squares_ptr,
        running_mean_ptr,
        running_var_ptr,
        mean.data_ptr<float>(),
        invstd.data_ptr<float>(),
        rows,
        channels,
        parts,
        static_cast<float>(momentum),
        static_cast<float>(eps),
        training);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

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

    normalize_kernel<<<blocks, 256, 0, stream>>>(
        input.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>(),
        count,
        channels);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return {output, mean, invstd};
}
"""

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


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

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

        grad_bias = grad_output.sum(dim=0)
        grad_weight = (grad_output * normalized).sum(dim=0)

        if ctx.training:
            rows = input.shape[0]
            grad_scaled = (
                grad_output
                - grad_bias / rows
                - normalized * (grad_weight / rows)
            ) * (weight * invstd)
        else:
            grad_scaled = grad_output * (weight * invstd)

        grad_input = grad_scaled * scale
        grad_scale = (grad_scaled * input).sum(dim=0).reshape_as(scale)

        return (
            grad_input, 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):
        linear = self.gemm(x)
        bn = self.bn
        channels = bn.num_features

        scale_supported = (
            self.scale.shape == (channels,)
            or self.scale.shape == (1, channels)
        )
        tensors = [linear, self.scale, bn.weight, bn.bias]
        if bn.running_mean is not None:
            tensors.extend([bn.running_mean, bn.running_var])

        supported = (
            linear.is_cuda
            and linear.ndim == 2
            and linear.shape[0] > 1
            and channels > 0
            and scale_supported
            and bn.affine
            and all(
                t is not None
                and t.dtype == torch.float32
                and t.device == linear.device
                and t.is_contiguous()
                for t in tensors
            )
        )

        if not supported:
            return bn(linear * self.scale)

        training = bn.training or bn.running_mean is None
        momentum = 0.0 if bn.momentum is None else bn.momentum

        if bn.training and bn.track_running_stats:
            bn.num_batches_tracked.add_(1)
            if bn.momentum is None:
                momentum = 1.0 / float(bn.num_batches_tracked.item())

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

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