```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_CPP = 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,
    torch::Tensor batches,
    bool training,
    double momentum,
    double eps);
"""

_CUDA = 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 ROW_TILE = 256;

__device__ __forceinline__ void merge_stats(
    float& mean, float& m2, int& count,
    float other_mean, float other_m2, int other_count) {
    if (other_count == 0) return;
    if (count == 0) {
        mean = other_mean;
        m2 = other_m2;
        count = other_count;
        return;
    }
    int total = count + other_count;
    float delta = other_mean - mean;
    float ratio = float(other_count) / float(total);
    m2 += other_m2 + delta * delta * (float(count) * ratio);
    mean += delta * ratio;
    count = total;
}

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

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int col = blockIdx.x * 32 + tx;
    int first = blockIdx.y * ROW_TILE;
    int end = min(first + ROW_TILE, rows);

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

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

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

    for (int stride = 4; stride > 0; stride >>= 1) {
        if (ty < stride) {
            merge_stats(
                mean, m2, count,
                sm[ty + stride][tx],
                sv[ty + stride][tx],
                sn[ty + stride][tx]);
            sm[ty][tx] = mean;
            sv[ty][tx] = m2;
            sn[ty][tx] = count;
        }
        __syncthreads();
    }

    if (ty == 0 && col < cols) {
        int64_t index = (int64_t)blockIdx.y * cols + col;
        partial_mean[index] = mean;
        partial_m2[index] = m2;
    }
}

__global__ void finish_stats(
    const float* __restrict__ partial_mean,
    const float* __restrict__ partial_m2,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    const int64_t* __restrict__ batches,
    float* __restrict__ mean_out,
    float* __restrict__ invstd_out,
    int rows, int cols, int parts,
    double momentum, double eps) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

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

    for (int part = 0; part < parts; ++part) {
        int n = min(ROW_TILE, rows - part * ROW_TILE);
        int64_t index = (int64_t)part * cols + col;
        merge_stats(mean, m2, count,
                    partial_mean[index], partial_m2[index], n);
    }

    mean_out[col] = mean;
    invstd_out[col] = rsqrtf(float(double(m2 / float(rows)) + eps));

    float factor = momentum < 0.0
        ? float(1.0 / double(batches[0] + 1))
        : float(momentum);
    running_mean[col] =
        (1.0f - factor) * running_mean[col] + factor * mean;
    running_var[col] =
        (1.0f - factor) * running_var[col]
        + factor * (m2 / float(rows - 1));
}

__global__ void evaluation_stats(
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float* __restrict__ mean,
    float* __restrict__ invstd,
    int cols, double eps) {
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col < cols) {
        mean[col] = running_mean[col];
        invstd[col] = rsqrtf(float(double(running_var[col]) + eps));
    }
}

__global__ void apply_bn(
    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__ out,
    int64_t* __restrict__ batches,
    int64_t size, int cols, bool training) {
    int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    if (index == 0 && training) {
        batches[0] += 1;
    }

    for (; index < size; index += stride) {
        int col = index % cols;
        float z = x[index] * scale[col];
        float normalized = (z - mean[col]) * invstd[col];
        out[index] = 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,
    torch::Tensor batches,
    bool training,
    double momentum,
    double eps) {
    c10::cuda::CUDAGuard guard(x.device());
    auto stream = at::cuda::getCurrentCUDAStream();

    int rows = x.size(0);
    int cols = x.size(1);
    auto out = torch::empty_like(x);
    auto mean = torch::empty({cols}, x.options());
    auto invstd = torch::empty({cols}, x.options());

    if (training) {
        int parts = (rows + ROW_TILE - 1) / ROW_TILE;
        auto partial_mean = torch::empty({parts, cols}, x.options());
        auto partial_m2 = torch::empty({parts, cols}, x.options());

        partial_stats<<<dim3((cols + 31) / 32, parts),
                        dim3(32, 8), 0, stream>>>(
            x.data_ptr<float>(), scale.data_ptr<float>(),
            partial_mean.data_ptr<float>(),
            partial_m2.data_ptr<float>(), rows, cols);
        C10_CUDA_KERNEL_LAUNCH_CHECK();

        finish_stats<<<(cols + 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>(),
            batches.data_ptr<int64_t>(),
            mean.data_ptr<float>(), invstd.data_ptr<float>(),
            rows, cols, parts, momentum, eps);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        evaluation_stats<<<(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, eps);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    int64_t size = x.numel();
    int blocks = int(std::min<int64_t>((size + 255) / 256, 65535));
    apply_bn<<<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>(),
        out.data_ptr<float>(), batches.data_ptr<int64_t>(),
        size, cols, training);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

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

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="fused_scale_batchnorm_welford_v1",
            cpp_sources=_CPP,
            cuda_sources=_CUDA,
            functions=["scaled_bn_cuda"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            verbose=False,
        )
    return _extension


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

    @staticmethod
    def backward(ctx, grad_out):
        x, scale, weight, mean, invstd = ctx.saved_tensors
        normalized = (x * scale - mean) * invstd
        sum_grad = grad_out.sum(dim=0)
        sum_grad_normalized = (grad_out * normalized).sum(dim=0)

        if ctx.training:
            n = x.shape[0]
            grad_z = (
                grad_out
                - sum_grad / n
                - normalized * (sum_grad_normalized / n)
            ) * (weight * invstd)
        else:
            grad_z = grad_out * (weight * invstd)

        grad_x = grad_z * scale if ctx.needs_input_grad[0] else None
        grad_scale = (
            (grad_z * x).sum(dim=0) if ctx.needs_input_grad[1] else None
        )
        grad_weight = (
            sum_grad_normalized if ctx.needs_input_grad[2] else None
        )
        grad_bias = sum_grad if ctx.needs_input_grad[3] else None

        return (
            grad_x, grad_scale, grad_weight, grad_bias,
            None, 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)

        fast_path = (
            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 self.scale.shape == (x.shape[1],)
            and self.bn.affine
            and self.bn.track_running_stats
            and all(
                t is not None
                and t.dtype == torch.float32
                and t.device == x.device
                and t.is_contiguous()
                for t in (
                    self.scale,
                    self.bn.weight,
                    self.bn.bias,
                    self.bn.running_mean,
                    self.bn.running_var,
                )
            )
        )

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

        momentum = -1.0 if self.bn.momentum is None else self.bn.momentum
        return _ScaledBatchNorm.apply(
            x,
            self.scale,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.bn.num_batches_tracked,
            self.bn.training,
            momentum,
            self.bn.eps,
        )
```