```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> scale_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 = 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 <algorithm>
#include <vector>

constexpr int ROW_TILE = 256;

__global__ void partial_statistics(
    const float* __restrict__ x,
    const float* __restrict__ scale,
    float* __restrict__ partial_mean,
    float* __restrict__ partial_m2,
    int rows, int features)
{
    __shared__ float sums[8][32];
    __shared__ float squares[8][32];

    const int lane = threadIdx.x;
    const int group = threadIdx.y;
    const int f = blockIdx.x * 32 + lane;
    const int begin = blockIdx.y * ROW_TILE;
    const int count = min(ROW_TILE, rows - begin);

    float center = 0.0f;
    float sum = 0.0f;
    float square = 0.0f;

    if (f < features) {
        const float s = scale[f];
        center = __fmul_rn(x[(int64_t)begin * features + f], s);

        #pragma unroll
        for (int r = group; r < ROW_TILE; r += 8) {
            if (r < count) {
                const float value =
                    __fmul_rn(x[(int64_t)(begin + r) * features + f], s);
                const float delta = value - center;
                sum += delta;
                square = fmaf(delta, delta, square);
            }
        }
    }

    sums[group][lane] = sum;
    squares[group][lane] = square;
    __syncthreads();

    if (group == 0 && f < features) {
        float total = 0.0f;
        float total_square = 0.0f;
        #pragma unroll
        for (int g = 0; g < 8; ++g) {
            total += sums[g][lane];
            total_square += squares[g][lane];
        }

        const float offset = total / (float)count;
        const int64_t index = (int64_t)blockIdx.y * features + f;
        partial_mean[index] = center + offset;
        partial_m2[index] =
            fmaxf(fmaf(-total, offset, total_square), 0.0f);
    }
}

__global__ void finish_statistics(
    const float* __restrict__ partial_mean,
    const float* __restrict__ partial_m2,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ mean_out,
    float* __restrict__ invstd_out,
    int rows, int features, int tiles, float momentum, float eps)
{
    const int f = blockIdx.x * blockDim.x + threadIdx.x;
    if (f >= features)
        return;

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

    for (int tile = 0; tile < tiles; ++tile) {
        const int n = min(ROW_TILE, rows - tile * ROW_TILE);
        const int total = count + n;
        const int64_t index = (int64_t)tile * features + f;
        const float delta = partial_mean[index] - mean;
        const float ratio = (float)n / (float)total;
        mean = fmaf(delta, ratio, mean);
        m2 += partial_m2[index] + delta * delta * ((float)count * ratio);
        count = total;
    }

    mean_out[f] = mean;
    invstd_out[f] = rsqrtf(m2 / (float)rows + eps);
    running_mean[f] =
        (1.0f - momentum) * running_mean[f] + momentum * mean;
    running_var[f] =
        (1.0f - momentum) * running_var[f] +
        momentum * (m2 / (float)(rows - 1));
}

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

__device__ __forceinline__ float normalize_value(
    float x, float scale, float mean, float invstd,
    float weight, float bias)
{
    const float scaled = __fmul_rn(x, scale);
    const float normalized = (scaled - mean) * invstd;
    return fmaf(normalized, weight, bias);
}

__global__ void normalize4(
    const float4* __restrict__ x,
    const float4* __restrict__ scale,
    const float4* __restrict__ weight,
    const float4* __restrict__ bias,
    const float4* __restrict__ mean,
    const float4* __restrict__ invstd,
    float4* __restrict__ output,
    int64_t vectors, int feature_vectors)
{
    for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
         i < vectors; i += (int64_t)blockDim.x * gridDim.x) {
        const int f = i % feature_vectors;
        const float4 v = x[i];
        const float4 s = scale[f];
        const float4 w = weight[f];
        const float4 b = bias[f];
        const float4 m = mean[f];
        const float4 r = invstd[f];

        float4 out;
        out.x = normalize_value(v.x, s.x, m.x, r.x, w.x, b.x);
        out.y = normalize_value(v.y, s.y, m.y, r.y, w.y, b.y);
        out.z = normalize_value(v.z, s.z, m.z, r.z, w.z, b.z);
        out.w = normalize_value(v.w, s.w, m.w, r.w, w.w, b.w);
        output[i] = out;
    }
}

__global__ void normalize_scalar(
    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)
{
    for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
         i < elements; i += (int64_t)blockDim.x * gridDim.x) {
        const int f = i % features;
        output[i] = normalize_value(
            x[i], scale[f], mean[f], invstd[f], weight[f], bias[f]);
    }
}

std::vector<torch::Tensor> scale_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 guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();
    const int rows = x.size(0);
    const int features = x.size(1);
    const int tiles = (rows + ROW_TILE - 1) / ROW_TILE;

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

    if (training) {
        auto pm = torch::empty({tiles, features}, x.options());
        auto pv = torch::empty_like(pm);

        partial_statistics<<<
            dim3((features + 31) / 32, tiles), dim3(32, 8), 0, stream>>>(
                x.data_ptr<float>(), scale.data_ptr<float>(),
                pm.data_ptr<float>(), pv.data_ptr<float>(),
                rows, features);
        C10_CUDA_KERNEL_LAUNCH_CHECK();

        finish_statistics<<<(features + 255) / 256, 256, 0, stream>>>(
            pm.data_ptr<float>(), pv.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            mean.data_ptr<float>(), invstd.data_ptr<float>(),
            rows, features, tiles, (float)momentum, (float)eps);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        inference_statistics<<<(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, (float)eps);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    const bool vectorized =
        features % 4 == 0 &&
        reinterpret_cast<uintptr_t>(scale.data_ptr<float>()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(weight.data_ptr<float>()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(bias.data_ptr<float>()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(x.data_ptr<float>()) % 16 == 0;

    if (vectorized) {
        const int64_t vectors = x.numel() / 4;
        const int blocks =
            (int)std::min<int64_t>((vectors + 255) / 256, 4096);

        normalize4<<<blocks, 256, 0, stream>>>(
            reinterpret_cast<const float4*>(x.data_ptr<float>()),
            reinterpret_cast<const float4*>(scale.data_ptr<float>()),
            reinterpret_cast<const float4*>(weight.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            reinterpret_cast<const float4*>(mean.data_ptr<float>()),
            reinterpret_cast<const float4*>(invstd.data_ptr<float>()),
            reinterpret_cast<float4*>(output.data_ptr<float>()),
            vectors, features / 4);
    } else {
        const int blocks =
            (int)std::min<int64_t>((x.numel() + 255) / 256, 4096);

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

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

_ops = load_inline(
    name="fused_scale_batchnorm_vectorized_v4",
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=["scale_bn_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _ScaleBatchNorm(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, x, scale, weight, bias,
        running_mean, running_var, training, momentum, eps
    ):
        output, mean, invstd = _ops.scale_bn_cuda(
            x, scale, weight, bias,
            running_mean, running_var, training, momentum, eps
        )
        ctx.training = training
        ctx.eps = eps
        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
        scaled = x * scale

        if ctx.training and torch.is_grad_enabled():
            variance, mean = torch.var_mean(scaled, dim=0, correction=0)
            invstd = torch.rsqrt(variance + ctx.eps)

        normalized = (scaled - mean) * invstd
        sum_grad = grad_output.sum(0)
        sum_grad_normalized = (grad_output * normalized).sum(0)

        if ctx.training:
            count = x.shape[0]
            grad_scaled = (
                grad_output
                - sum_grad / count
                - normalized * (sum_grad_normalized / count)
            ) * (weight * invstd)
        else:
            grad_scaled = grad_output * (weight * invstd)

        return (
            grad_scaled * scale if ctx.needs_input_grad[0] else None,
            (grad_scaled * x).sum(0) if ctx.needs_input_grad[1] else None,
            sum_grad_normalized if ctx.needs_input_grad[2] else None,
            sum_grad if ctx.needs_input_grad[3] else 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):
        # Use the same single GEMM and precision policy as nn.Linear.
        x = self.gemm(x)
        bn = self.bn

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

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

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

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