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


_cpp = r"""
#include <torch/extension.h>
#include <vector>

torch::Tensor compensated_mm_cuda(torch::Tensor x, torch::Tensor weight);

std::vector<torch::Tensor> fused_bn_cuda(
    torch::Tensor z, torch::Tensor linear_bias, 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 <cuda_fp16.h>
#include <cublas_v2.h>
#include <algorithm>
#include <cstdint>
#include <vector>

constexpr int ROW_TILE = 256;
constexpr float RESIDUAL_SCALE = 4096.0f;

__global__ void split_half(
    const float* __restrict__ input,
    __half* __restrict__ high,
    __half* __restrict__ low,
    int64_t count)
{
    for (int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
         i < count; i += (int64_t)blockDim.x * gridDim.x) {
        const float v = input[i];
        const __half h = __float2half_rn(v);
        high[i] = h;
        low[i] = __float2half_rn(
            (v - __half2float(h)) * RESIDUAL_SCALE);
    }
}

torch::Tensor compensated_mm_cuda(torch::Tensor x, torch::Tensor weight)
{
    const c10::cuda::CUDAGuard guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();
    const int m = x.size(0);
    const int k = x.size(1);
    const int n = weight.size(0);

    const auto half_options = x.options().dtype(torch::kFloat16);
    auto xh = torch::empty(x.sizes(), half_options);
    auto xl = torch::empty_like(xh);
    auto wh = torch::empty(weight.sizes(), half_options);
    auto wl = torch::empty_like(wh);
    auto output = torch::empty({m, n}, x.options());

    const int xb = (int)std::min<int64_t>(
        (x.numel() + 255) / 256, 4096);
    const int wb = (int)std::min<int64_t>(
        (weight.numel() + 255) / 256, 4096);

    split_half<<<xb, 256, 0, stream>>>(
        x.data_ptr<float>(),
        reinterpret_cast<__half*>(xh.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(xl.data_ptr<at::Half>()),
        x.numel());
    split_half<<<wb, 256, 0, stream>>>(
        weight.data_ptr<float>(),
        reinterpret_cast<__half*>(wh.data_ptr<at::Half>()),
        reinterpret_cast<__half*>(wl.data_ptr<at::Half>()),
        weight.numel());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const auto handle = at::cuda::getCurrentCUDABlasHandle();
    cublasMath_t old_math;
    TORCH_CHECK(
        cublasGetMathMode(handle, &old_math) == CUBLAS_STATUS_SUCCESS,
        "Could not read cuBLAS math mode");

    const auto math_mode = static_cast<cublasMath_t>(
        CUBLAS_TENSOR_OP_MATH |
        CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION);
    TORCH_CHECK(
        cublasSetMathMode(handle, math_mode) == CUBLAS_STATUS_SUCCESS,
        "Could not set cuBLAS math mode");

    const float correction = 1.0f / RESIDUAL_SCALE;
    const float zero = 0.0f;
    const float one = 1.0f;

    auto gemm = [&](const torch::Tensor& a, const torch::Tensor& b,
                    const float* alpha, const float* beta) {
        return cublasGemmEx(
            handle, CUBLAS_OP_T, CUBLAS_OP_N,
            n, m, k,
            alpha,
            b.data_ptr<at::Half>(), CUDA_R_16F, k,
            a.data_ptr<at::Half>(), CUDA_R_16F, k,
            beta,
            output.data_ptr<float>(), CUDA_R_32F, n,
            CUBLAS_COMPUTE_32F,
            CUBLAS_GEMM_DEFAULT_TENSOR_OP);
    };

    // Accumulate the two small correction products before the main product.
    auto status = gemm(xh, wl, &correction, &zero);
    if (status == CUBLAS_STATUS_SUCCESS)
        status = gemm(xl, wh, &correction, &one);
    if (status == CUBLAS_STATUS_SUCCESS)
        status = gemm(xh, wh, &one, &one);

    const auto restored = cublasSetMathMode(handle, old_math);
    TORCH_CHECK(
        status == CUBLAS_STATUS_SUCCESS,
        "Compensated GEMM failed: ", static_cast<int>(status));
    TORCH_CHECK(
        restored == CUBLAS_STATUS_SUCCESS,
        "Could not restore cuBLAS math mode");
    return output;
}

__device__ __forceinline__ float scaled_value(
    float z, float linear_bias, float scale)
{
    return __fmul_rn(__fadd_rn(z, linear_bias), scale);
}

__global__ void partial_statistics(
    const float* __restrict__ z,
    const float* __restrict__ linear_bias,
    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 b = linear_bias[f];
        const float s = scale[f];
        center = scaled_value(z[(int64_t)begin * features + f], b, s);

        #pragma unroll
        for (int r = group; r < ROW_TILE; r += 8) {
            if (r < count) {
                const float value = scaled_value(
                    z[(int64_t)(begin + r) * features + f], b, 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 z, float linear_bias, float scale,
    float mean, float invstd, float weight, float bias)
{
    const float value = scaled_value(z, linear_bias, scale);
    const float normalized = (value - mean) * invstd;
    return fmaf(normalized, weight, bias);
}

__global__ void normalize4(
    const float4* __restrict__ z,
    const float4* __restrict__ linear_bias,
    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 = z[i];
        const float4 l = linear_bias[f];
        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, l.x, s.x, m.x, r.x, w.x, b.x);
        out.y = normalize_value(v.y, l.y, s.y, m.y, r.y, w.y, b.y);
        out.z = normalize_value(v.z, l.z, s.z, m.z, r.z, w.z, b.z);
        out.w = normalize_value(v.w, l.w, s.w, m.w, r.w, w.w, b.w);
        output[i] = out;
    }
}

__global__ void normalize_scalar(
    const float* __restrict__ z,
    const float* __restrict__ linear_bias,
    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(
            z[i], linear_bias[f], scale[f],
            mean[f], invstd[f], weight[f], bias[f]);
    }
}

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

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

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

        partial_statistics<<<
            dim3((features + 31) / 32, tiles), dim3(32, 8), 0, stream>>>(
                z.data_ptr<float>(), linear_bias.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>(linear_bias.data_ptr<float>()) % 16 == 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;

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

        normalize4<<<blocks, 256, 0, stream>>>(
            reinterpret_cast<const float4*>(z.data_ptr<float>()),
            reinterpret_cast<const float4*>(linear_bias.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>((z.numel() + 255) / 256, 4096);

        normalize_scalar<<<blocks, 256, 0, stream>>>(
            z.data_ptr<float>(), linear_bias.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>(),
            z.numel(), features);
    }

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

_ops = load_inline(
    name="compensated_half_gemm_vector_bn_v5",
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=["compensated_mm_cuda", "fused_bn_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class _CompensatedMM(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight):
        ctx.save_for_backward(x, weight)
        return _ops.compensated_mm_cuda(x, weight)

    @staticmethod
    def backward(ctx, grad_output):
        x, weight = ctx.saved_tensors
        grad_x = grad_output @ weight if ctx.needs_input_grad[0] else None
        grad_weight = (
            grad_output.transpose(0, 1) @ x
            if ctx.needs_input_grad[1] else None
        )
        return grad_x, grad_weight


class _FusedBatchNorm(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, z, linear_bias, scale, weight, bias,
        running_mean, running_var, training, momentum, eps
    ):
        output, mean, invstd = _ops.fused_bn_cuda(
            z, linear_bias, scale, weight, bias,
            running_mean, running_var, training, momentum, eps
        )
        ctx.training = training
        ctx.eps = eps
        ctx.save_for_backward(z, linear_bias, scale, weight, mean, invstd)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        z, linear_bias, scale, weight, mean, invstd = ctx.saved_tensors
        linear_output = z + linear_bias
        scaled = linear_output * 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 = z.shape[0]
            grad_scaled = (
                grad_output
                - sum_grad / count
                - normalized * (sum_grad_normalized / count)
            ) * (weight * invstd)
        else:
            grad_scaled = grad_output * (weight * invstd)

        grad_linear = grad_scaled * scale
        return (
            grad_linear if ctx.needs_input_grad[0] else None,
            grad_linear.sum(0) if ctx.needs_input_grad[1] else None,
            (grad_scaled * linear_output).sum(0)
            if ctx.needs_input_grad[2] else None,
            sum_grad_normalized if ctx.needs_input_grad[3] else None,
            sum_grad if ctx.needs_input_grad[4] 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):
        linear = self.gemm
        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] == linear.in_features
            and linear.in_features > 0
            and linear.out_features > 0
            and self.scale.shape == (linear.out_features,)
            and bn.affine
            and bn.track_running_stats
            and not torch.is_autocast_enabled()
            and all(
                t is not None
                and t.device == x.device
                and t.dtype == torch.float32
                and t.is_contiguous()
                for t in (
                    linear.weight, linear.bias, self.scale,
                    bn.weight, bn.bias, bn.running_mean, bn.running_var,
                )
            )
        )

        if not use_custom:
            return bn(linear(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()

        z = _CompensatedMM.apply(x, linear.weight)
        return _FusedBatchNorm.apply(
            z, linear.bias, self.scale, bn.weight, bn.bias,
            bn.running_mean, bn.running_var,
            bn.training, factor, bn.eps,
        )
```