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


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

torch::Tensor fused_forward(
    torch::Tensor input,
    torch::Tensor linear_weight,
    torch::Tensor linear_bias,
    torch::Tensor bn_weight,
    torch::Tensor bn_bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor scale,
    bool training,
    double eps,
    double momentum);
"""

_cuda_source = 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 <cuda_fp16.h>
#include <stdint.h>
#include <math.h>

constexpr int TILE_ROWS = 128;
constexpr int THREADS = 256;

struct __align__(8) Half4 {
    __half2 lo;
    __half2 hi;
};

__global__ void convert_operands_vectorized(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    __half* __restrict__ half_input,
    __half* __restrict__ half_weight,
    int64_t input_vectors,
    int64_t total_vectors)
{
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (; i < total_vectors; i += stride) {
        float4 v;
        Half4* destination;

        if (i < input_vectors) {
            v = reinterpret_cast<const float4*>(input)[i];
            destination = reinterpret_cast<Half4*>(half_input) + i;
        } else {
            int64_t j = i - input_vectors;
            v = reinterpret_cast<const float4*>(weight)[j];
            destination = reinterpret_cast<Half4*>(half_weight) + j;
        }

        Half4 result;
        result.lo = __floats2half2_rn(v.x, v.y);
        result.hi = __floats2half2_rn(v.z, v.w);
        *destination = result;
    }
}

__global__ void convert_operands_scalar(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    __half* __restrict__ half_input,
    __half* __restrict__ half_weight,
    int64_t input_elements,
    int64_t total_elements)
{
    int64_t i = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;

    for (; i < total_elements; i += stride) {
        if (i < input_elements) {
            half_input[i] = __float2half_rn(input[i]);
        } else {
            int64_t j = i - input_elements;
            half_weight[j] = __float2half_rn(weight[j]);
        }
    }
}

__global__ void partial_stats(
    const float* __restrict__ x,
    const float* __restrict__ linear_bias,
    float* __restrict__ partial,
    int rows, int cols, int parts)
{
    __shared__ float means[8][32];
    __shared__ float m2s[8][32];

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

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

    if (col < cols) {
        float b = linear_bias[col];
        for (int row = start + ty; row < end; row += 8) {
            float v = __fadd_rn(x[(int64_t)row * cols + col], b);
            ++count;
            float delta = v - mean;
            mean += delta / count;
            m2 = fmaf(delta, v - mean, m2);
        }
    }

    means[ty][tx] = mean;
    m2s[ty][tx] = m2;
    __syncthreads();

    if (ty == 0 && col < cols) {
        float total_mean = 0.0f;
        float total_m2 = 0.0f;
        int total_count = 0;

        #pragma unroll
        for (int lane = 0; lane < 8; ++lane) {
            int n = max(0, (end - start - lane + 7) / 8);
            if (n > 0) {
                int combined = total_count + n;
                float delta = means[lane][tx] - total_mean;
                total_m2 += m2s[lane][tx]
                    + delta * delta * ((float)total_count * n / combined);
                total_mean += delta * ((float)n / combined);
                total_count = combined;
            }
        }

        int64_t index = (int64_t)blockIdx.y * cols + col;
        partial[index] = total_mean;
        partial[(int64_t)parts * cols + index] = total_m2;
    }
}

__global__ void training_coefficients(
    const float* __restrict__ partial,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ coefficients,
    int rows, int cols, int parts,
    float eps, float momentum)
{
    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(TILE_ROWS, rows - part * TILE_ROWS);
        int combined = count + n;
        int64_t index = (int64_t)part * cols + col;
        float delta = partial[index] - mean;

        m2 += partial[(int64_t)parts * cols + index]
            + delta * delta * ((float)count * n / combined);
        mean += delta * ((float)n / combined);
        count = combined;
    }

    coefficients[col] = mean;
    coefficients[cols + col] =
        weight[col] * rsqrtf(fmaxf(m2 / rows, 0.0f) + eps);
    coefficients[2 * cols + col] = bias[col];

    running_mean[col] =
        (1.0f - momentum) * running_mean[col] + momentum * mean;
    running_var[col] =
        (1.0f - momentum) * running_var[col]
        + momentum * (m2 / (rows - 1));
}

__global__ void inference_coefficients(
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ running_mean,
    const float* __restrict__ running_var,
    float* __restrict__ coefficients,
    int cols, float eps)
{
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col >= cols) return;

    coefficients[col] = running_mean[col];
    coefficients[cols + col] = weight[col] * rsqrtf(running_var[col] + eps);
    coefficients[2 * cols + col] = bias[col];
}

template<bool MAXIMUM>
__device__ __forceinline__ float block_reduce(float value)
{
    __shared__ float scratch[8];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, value, offset);
        if (MAXIMUM) value = fmaxf(value, other);
        else value += other;
    }

    if (lane == 0) scratch[warp] = value;
    __syncthreads();

    if (warp == 0) {
        value = lane < 8 ? scratch[lane] : (MAXIMUM ? -INFINITY : 0.0f);

        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            float other = __shfl_down_sync(0xffffffff, value, offset);
            if (MAXIMUM) value = fmaxf(value, other);
            else value += other;
        }

        if (lane == 0) scratch[0] = value;
    }

    __syncthreads();
    float result = scratch[0];
    __syncthreads();
    return result;
}

template<int ITEMS>
__global__ void normalized_softmax(
    const float* __restrict__ x,
    const float* __restrict__ linear_bias,
    const float* __restrict__ coefficients,
    const float* __restrict__ scale,
    float* __restrict__ out,
    int cols, bool scalar_scale)
{
    int tid = threadIdx.x;
    int64_t base = (int64_t)blockIdx.x * cols;
    float scalar = scalar_scale ? scale[0] : 1.0f;

    float values[ITEMS];
    float maximum = -INFINITY;

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int col = tid + i * THREADS;
        float value = -INFINITY;

        if (col < cols) {
            float z = __fadd_rn(x[base + col], linear_bias[col]);
            value = fmaf(
                z - coefficients[col],
                coefficients[cols + col],
                coefficients[2 * cols + col]
            ) * (scalar_scale ? scalar : scale[col]);
        }

        values[i] = value;
        maximum = fmaxf(maximum, value);
    }

    maximum = block_reduce<true>(maximum);
    float sum = 0.0f;

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        values[i] = __expf(values[i] - maximum);
        sum += values[i];
    }

    float inverse_sum = 1.0f / block_reduce<false>(sum);

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        int col = tid + i * THREADS;
        if (col < cols)
            out[base + col] = values[i] * inverse_sum;
    }
}

torch::Tensor fused_forward(
    torch::Tensor input,
    torch::Tensor linear_weight,
    torch::Tensor linear_bias,
    torch::Tensor bn_weight,
    torch::Tensor bn_bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor scale,
    bool training,
    double eps,
    double momentum)
{
    const c10::cuda::CUDAGuard guard(input.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    int rows = input.size(0);
    int inner = input.size(1);
    int cols = linear_weight.size(0);

    auto x = torch::empty({rows, cols}, input.options());
    auto out = torch::empty_like(x);
    if (rows == 0) return out;

    auto half_input = torch::empty(
        {rows, inner}, input.options().dtype(torch::kFloat16));
    auto half_weight = torch::empty(
        {cols, inner}, input.options().dtype(torch::kFloat16));

    int64_t nx = input.numel();
    int64_t nw = linear_weight.numel();

    const float* input_ptr = input.data_ptr<float>();
    const float* weight_ptr = linear_weight.data_ptr<float>();
    __half* half_input_ptr =
        reinterpret_cast<__half*>(half_input.data_ptr<at::Half>());
    __half* half_weight_ptr =
        reinterpret_cast<__half*>(half_weight.data_ptr<at::Half>());

    bool vectorized =
        nx % 4 == 0 && nw % 4 == 0 &&
        reinterpret_cast<uintptr_t>(input_ptr) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(weight_ptr) % 16 == 0;

    if (vectorized) {
        int64_t count = (nx + nw) / 4;
        int blocks = static_cast<int>(
            std::min<int64_t>((count + 255) / 256, 65535));
        convert_operands_vectorized<<<blocks, 256, 0, stream>>>(
            input_ptr, weight_ptr, half_input_ptr, half_weight_ptr,
            nx / 4, count);
    } else {
        int64_t count = nx + nw;
        int blocks = static_cast<int>(
            std::min<int64_t>((count + 255) / 256, 65535));
        convert_operands_scalar<<<blocks, 256, 0, stream>>>(
            input_ptr, weight_ptr, half_input_ptr, half_weight_ptr,
            nx, count);
    }

    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    TORCH_CHECK(
        cublasSetStream(handle, stream) == CUBLAS_STATUS_SUCCESS,
        "cublasSetStream failed");

    const float alpha = 1.0f;
    const float beta = 0.0f;

    // FP16 tensor-core operands with FP32 accumulation and FP32 output.
    cublasStatus_t status = cublasGemmEx(
        handle,
        CUBLAS_OP_T, CUBLAS_OP_N,
        cols, rows, inner,
        &alpha,
        half_weight_ptr, CUDA_R_16F, inner,
        half_input_ptr, CUDA_R_16F, inner,
        &beta,
        x.data_ptr<float>(), CUDA_R_32F, cols,
        CUBLAS_COMPUTE_32F,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP);

    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
                "cublasGemmEx failed: ", static_cast<int>(status));

    auto coefficients = torch::empty({3, cols}, input.options());

    if (training) {
        int parts = (rows + TILE_ROWS - 1) / TILE_ROWS;
        auto partial = torch::empty({2, parts, cols}, input.options());

        partial_stats<<<
            dim3((cols + 31) / 32, parts), dim3(32, 8), 0, stream
        >>>(
            x.data_ptr<float>(), linear_bias.data_ptr<float>(),
            partial.data_ptr<float>(), rows, cols, parts);

        training_coefficients<<<(cols + 255) / 256, 256, 0, stream>>>(
            partial.data_ptr<float>(),
            bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            coefficients.data_ptr<float>(),
            rows, cols, parts, (float)eps, (float)momentum);
    } else {
        inference_coefficients<<<(cols + 255) / 256, 256, 0, stream>>>(
            bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
            coefficients.data_ptr<float>(), cols, (float)eps);
    }

    bool scalar_scale = scale.numel() == 1;

    #define LAUNCH(N) \
        normalized_softmax<N><<<rows, THREADS, 0, stream>>>( \
            x.data_ptr<float>(), linear_bias.data_ptr<float>(), \
            coefficients.data_ptr<float>(), scale.data_ptr<float>(), \
            out.data_ptr<float>(), cols, scalar_scale)

    if      (cols <= 256)   { LAUNCH(1); }
    else if (cols <= 512)   { LAUNCH(2); }
    else if (cols <= 1024)  { LAUNCH(4); }
    else if (cols <= 2048)  { LAUNCH(8); }
    else if (cols <= 4096)  { LAUNCH(16); }
    else if (cols <= 8192)  { LAUNCH(32); }
    else                   { LAUNCH(64); }

    #undef LAUNCH

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ops = load_inline(
    name="half_gemm_fp32_bn_scale_softmax_v3",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublas"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(
        self,
        in_features,
        out_features,
        bn_eps=1e-5,
        bn_momentum=0.1,
        scale_shape=(1,),
    ):
        super().__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.bn = nn.BatchNorm1d(
            out_features, eps=bn_eps, momentum=bn_momentum
        )
        self.scale = nn.Parameter(torch.ones(scale_shape))
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        cols = self.gemm.out_features
        tensors = (
            self.gemm.weight,
            self.gemm.bias,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.scale,
        )

        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and x.size(1) == self.gemm.in_features
            and x.size(1) > 0
            and 0 < cols <= 16384
            and not torch.is_autocast_enabled()
            and self.bn.affine
            and self.bn.track_running_stats
            and (not self.bn.training or x.size(0) > 1)
            and tuple(self.scale.shape)
            in {(), (1,), (cols,), (1, 1), (1, cols)}
            and all(
                t is not None
                and t.device == x.device
                and t.dtype == torch.float32
                and t.is_contiguous()
                for t in tensors
            )
        )

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

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

        return _ops.fused_forward(
            x,
            self.gemm.weight,
            self.gemm.bias,
            self.bn.weight,
            self.bn.bias,
            self.bn.running_mean,
            self.bn.running_var,
            self.scale,
            self.bn.training,
            self.bn.eps,
            momentum,
        )
```