```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,
    torch::Tensor counter,
    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;
constexpr int TILE_COLS = 32;

// Each block reduces 256 rows for 32 adjacent features. Subtracting a
// local reference before accumulating avoids cancellation in E[x²]-E[x]².
__global__ void partial_stats_kernel(
    const float* __restrict__ x,
    const float* __restrict__ scale,
    float* __restrict__ partial_mean,
    float* __restrict__ partial_m2,
    int rows,
    int cols)
{
    __shared__ float sums[256];
    __shared__ float squares[256];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int group = tid >> 5;
    const int col = blockIdx.x * TILE_COLS + lane;
    const int start = blockIdx.y * TILE_ROWS;
    const int count = min(TILE_ROWS, rows - start);

    float reference = 0.0f;
    float multiplier = 0.0f;
    if (col < cols) {
        multiplier = scale[col];
        reference = x[(int64_t)start * cols + col] * multiplier;
    }

    float sum = 0.0f;
    float square = 0.0f;
    if (col < cols) {
        for (int offset = group; offset < count; offset += 8) {
            const float value =
                x[(int64_t)(start + offset) * cols + col] * multiplier;
            const float delta = value - reference;
            sum += delta;
            square = fmaf(delta, delta, square);
        }
    }

    sums[tid] = sum;
    squares[tid] = square;
    __syncthreads();

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

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

__global__ void finalize_stats_kernel(
    const float* __restrict__ partial_mean,
    const float* __restrict__ partial_m2,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int64_t* counter,
    float* __restrict__ mean,
    float* __restrict__ inverse_std,
    int rows,
    int cols,
    bool training,
    float momentum,
    float eps)
{
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

    if (!training) {
        mean[col] = running_mean[col];
        inverse_std[col] = rsqrtf(running_var[col] + eps);
        return;
    }

    float combined_mean = 0.0f;
    float combined_m2 = 0.0f;
    int combined_count = 0;
    const int chunks = (rows + TILE_ROWS - 1) / TILE_ROWS;

    for (int chunk = 0; chunk < chunks; ++chunk) {
        const int count = min(TILE_ROWS, rows - chunk * TILE_ROWS);
        const int64_t index = (int64_t)chunk * cols + col;
        const float local_mean = partial_mean[index];
        const float local_m2 = partial_m2[index];
        const int new_count = combined_count + count;
        const float delta = local_mean - combined_mean;

        combined_m2 += local_m2 +
            delta * delta *
            ((float)combined_count * count / new_count);
        combined_mean += delta * ((float)count / new_count);
        combined_count = new_count;
    }

    mean[col] = combined_mean;
    inverse_std[col] = rsqrtf(combined_m2 / rows + eps);

    running_mean[col] =
        (1.0f - momentum) * running_mean[col] + momentum * combined_mean;
    running_var[col] =
        (1.0f - momentum) * running_var[col] +
        momentum * (combined_m2 / (rows - 1));

    if (col == 0) {
        counter[0] += 1;
    }
}

__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__ inverse_std,
    float* __restrict__ output,
    int64_t size,
    int cols)
{
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
         index < size; index += stride) {
        const int col = index % cols;
        const float scaled = x[index] * scale[col];
        const float normalized = (scaled - mean[col]) * inverse_std[col];
        output[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 counter,
    bool training,
    double momentum,
    double eps)
{
    const c10::cuda::CUDAGuard device_guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

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

    torch::Tensor partial_mean;
    torch::Tensor partial_m2;
    const float* partial_mean_ptr = nullptr;
    const float* partial_m2_ptr = nullptr;

    if (training) {
        const int chunks = (rows + TILE_ROWS - 1) / TILE_ROWS;
        partial_mean = torch::empty({chunks, cols}, x.options());
        partial_m2 = torch::empty({chunks, cols}, x.options());
        partial_mean_ptr = partial_mean.data_ptr<float>();
        partial_m2_ptr = partial_m2.data_ptr<float>();

        const dim3 grid((cols + TILE_COLS - 1) / TILE_COLS, chunks);
        partial_stats_kernel<<<grid, 256, 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();
    }

    finalize_stats_kernel<<<(cols + 255) / 256, 256, 0, stream>>>(
        partial_mean_ptr,
        partial_m2_ptr,
        running_mean.data_ptr<float>(),
        running_var.data_ptr<float>(),
        counter.data_ptr<int64_t>(),
        mean.data_ptr<float>(),
        inverse_std.data_ptr<float>(),
        rows, cols, training, (float)momentum, (float)eps);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int64_t size = x.numel();
    const int blocks = (int)std::min<int64_t>((size + 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>(),
        inverse_std.data_ptr<float>(),
        output.data_ptr<float>(),
        size, cols);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

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

_scaled_bn = load_inline(
    name="fused_linear_scale_batchnorm_cuda_v1",
    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, counter,
        training, momentum, eps,
    ):
        output, mean, inverse_std = _scaled_bn.scaled_bn_cuda(
            x, scale, weight, bias,
            running_mean, running_var, counter,
            training, momentum, eps,
        )
        ctx.training = training
        ctx.save_for_backward(x, scale, weight, mean, inverse_std)
        return output

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

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

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

        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, 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 = (
            x, self.scale, self.bn.weight, self.bn.bias,
            self.bn.running_mean, self.bn.running_var,
        )
        supported = (
            x.is_cuda
            and x.ndim == 2
            and x.shape[0] > 1
            and x.shape[1] > 0
            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.device == x.device
                and t.dtype == torch.float32
                and t.is_contiguous()
                for t in tensors
            )
            and (x.shape[0] + 255) // 256 <= 65535
        )

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

        momentum = self.bn.momentum
        if momentum is None:
            momentum = (
                1.0 / (self.bn.num_batches_tracked.item() + 1)
                if self.bn.training else 0.0
            )

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