You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

    ```
    # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, scales the result, and applies batch normalization.
    """
    def __init__(self, in_features, out_features, scale_shape, eps=1e-5, momentum=0.1):
        super(Model, self).__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)
        x = x * self.scale
        x = self.bn(x)
        return x

batch_size = 16384
in_features = 4096
out_features = 4096
scale_shape = (out_features,)

def get_inputs():
    return [torch.rand(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features, scale_shape]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 5 of 6 for this task. Your previous solution (turn 4) was:
```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,
        )
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 29.6 ms; the PyTorch reference (with TF32 enabled) takes 5.31 ms; speedup 0.179x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| ampere_sgemm_128x64_tn | 28917.3 | 1 |
| normalize4(float4 const*, float4 const*, float4 const*, float4 const*, float4 const*, floa | 383.5 | 1 |
| partial_statistics(float const*, float const*, float*, float*, int, int) | 198.1 | 1 |
| finish_statistics(float const*, float const*, float*, float*, float*, float*, int, int, in | 39.2 | 1 |
| void at::native::vectorized_elementwise_kernel<2, at::native::CUDAFunctorOnSelf_add<long>, | 2.6 | 1 |

The fastest correct solution so far (0.374x speedup) is:
```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 accurate_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 <cublas_v2.h>
#include <cuda_runtime.h>
#include <algorithm>
#include <vector>

constexpr int ROW_TILE = 256;

__device__ __forceinline__ float round_tf32(float x) {
    unsigned int u = __float_as_uint(x);
    if ((u & 0x7f800000u) != 0x7f800000u) {
        u += 0xfffu + ((u >> 13) & 1u);
        u &= 0xffffe000u;
    }
    return __uint_as_float(u);
}

__global__ void split_kernel(
    const float* __restrict__ input,
    float* __restrict__ high,
    float* __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 value = input[i];
        const float h = round_tf32(value);
        high[i] = h;
        low[i] = round_tf32(value - h);
    }
}

torch::Tensor accurate_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);

    auto xh = torch::empty_like(x);
    auto xl = torch::empty_like(x);
    auto wh = torch::empty_like(weight);
    auto wl = torch::empty_like(weight);
    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_kernel<<<xb, 256, 0, stream>>>(
        x.data_ptr<float>(), xh.data_ptr<float>(), xl.data_ptr<float>(),
        x.numel());
    split_kernel<<<wb, 256, 0, stream>>>(
        weight.data_ptr<float>(), wh.data_ptr<float>(), wl.data_ptr<float>(),
        weight.numel());
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    auto handle = at::cuda::getCurrentCUDABlasHandle();
    cublasMath_t old_math;
    TORCH_CHECK(cublasGetMathMode(handle, &old_math) == CUBLAS_STATUS_SUCCESS,
                "Could not get cuBLAS math mode");
    TORCH_CHECK(
        cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH)
            == CUBLAS_STATUS_SUCCESS,
        "Could not set cuBLAS math mode");

    const float alpha = 1.0f;
    const float zero = 0.0f;
    const float one = 1.0f;

    // XW^T ≈ X_high W_high^T + X_high W_low^T + X_low W_high^T.
    // Accumulate the small correction terms before the main product.
    auto gemm = [&](const torch::Tensor& a, const torch::Tensor& b,
                    const float* beta) {
        return cublasGemmEx(
            handle, CUBLAS_OP_T, CUBLAS_OP_N,
            n, m, k, &alpha,
            b.data_ptr<float>(), CUDA_R_32F, k,
            a.data_ptr<float>(), CUDA_R_32F, k,
            beta, output.data_ptr<float>(), CUDA_R_32F, n,
            CUBLAS_COMPUTE_32F_FAST_TF32,
            CUBLAS_GEMM_DEFAULT_TENSOR_OP);
    };

    auto status = gemm(xh, wl, &zero);
    if (status == CUBLAS_STATUS_SUCCESS)
        status = gemm(xl, wh, &one);
    if (status == CUBLAS_STATUS_SUCCESS)
        status = gemm(xh, wh, &one);

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

__device__ __forceinline__ float scaled_value(
    float z, float bias, float scale)
{
    return __fmul_rn(__fadd_rn(z, 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 feature = blockIdx.x * 32 + lane;
    const int begin = blockIdx.y * ROW_TILE;
    const int count = min(ROW_TILE, rows - begin);

    float center = 0.0f;
    float b = 0.0f;
    float s = 0.0f;
    if (feature < features) {
        b = linear_bias[feature];
        s = scale[feature];
        center = scaled_value(z[(int64_t)begin * features + feature], b, s);
    }

    float sum = 0.0f;
    float square = 0.0f;
    if (feature < features) {
        for (int r = group; r < count; r += 8) {
            const float value = scaled_value(
                z[(int64_t)(begin + r) * features + feature], 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 && feature < 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 + feature;
        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 feature = blockIdx.x * blockDim.x + threadIdx.x;
    if (feature >= 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 + feature;
        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[feature] = mean;
    invstd_out[feature] = rsqrtf(m2 / (float)rows + eps);
    running_mean[feature] =
        (1.0f - momentum) * running_mean[feature] + momentum * mean;
    running_var[feature] =
        (1.0f - momentum) * running_var[feature] +
        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);
    }
}

__global__ void normalize(
    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;
        const float value = scaled_value(z[i], linear_bias[f], scale[f]);
        const float normalized = (value - mean[f]) * invstd[f];
        output[i] = fmaf(normalized, 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);
    } 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 int blocks =
        (int)std::min<int64_t>((z.numel() + 255) / 256, 4096);
    normalize<<<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_tf32_fused_bn_v3",
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=["accurate_mm_cuda", "fused_bn_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class _AccurateMM(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight):
        ctx.save_for_backward(x, weight)
        return _ops.accurate_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 = _AccurateMM.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,
        )
```

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.
