```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 batches,
    torch::Tensor scale, bool training, double eps,
    double momentum, bool cumulative);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <algorithm>
#include <cstdint>
#include <memory>

constexpr int THREADS = 256;
constexpr int TILE_ROWS = 128;
constexpr size_t WORKSPACE_LIMIT = 32ull * 1024 * 1024;

static void check_blas(cublasStatus_t status) {
    TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
                "cuBLASLt error: ", static_cast<int>(status));
}

struct GemmPlan {
    int rows, cols, inner, device;
    bool direct_float = false;
    size_t workspace_size = 0;
    cublasLtHandle_t handle = nullptr;
    cublasLtMatmulDesc_t operation = nullptr;
    cublasLtMatrixLayout_t a = nullptr, b = nullptr, c = nullptr;
    cublasLtMatmulAlgo_t algorithm{};

    void clear_descriptors() {
        if (a) cublasLtMatrixLayoutDestroy(a);
        if (b) cublasLtMatrixLayoutDestroy(b);
        if (c) cublasLtMatrixLayoutDestroy(c);
        if (operation) cublasLtMatmulDescDestroy(operation);
        a = b = c = nullptr;
        operation = nullptr;
    }

    bool configure(bool direct) {
        clear_descriptors();

        const auto compute = direct
            ? CUBLAS_COMPUTE_32F_FAST_16F
            : CUBLAS_COMPUTE_32F;
        const auto dtype = direct ? CUDA_R_32F : CUDA_R_16F;

        if (cublasLtMatmulDescCreate(
                &operation, compute, CUDA_R_32F)
            != CUBLAS_STATUS_SUCCESS)
            return false;

        cublasOperation_t trans_a = CUBLAS_OP_T;
        cublasOperation_t trans_b = CUBLAS_OP_N;
        cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;

        if (cublasLtMatmulDescSetAttribute(
                operation, CUBLASLT_MATMUL_DESC_TRANSA,
                &trans_a, sizeof(trans_a)) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatmulDescSetAttribute(
                operation, CUBLASLT_MATMUL_DESC_TRANSB,
                &trans_b, sizeof(trans_b)) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatmulDescSetAttribute(
                operation, CUBLASLT_MATMUL_DESC_EPILOGUE,
                &epilogue, sizeof(epilogue)) != CUBLAS_STATUS_SUCCESS)
            return false;

        if (cublasLtMatrixLayoutCreate(
                &a, dtype, inner, cols, inner) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(
                &b, dtype, inner, rows, inner) != CUBLAS_STATUS_SUCCESS ||
            cublasLtMatrixLayoutCreate(
                &c, CUDA_R_32F, cols, rows, cols) != CUBLAS_STATUS_SUCCESS)
            return false;

        cublasLtMatmulPreference_t preference = nullptr;
        check_blas(cublasLtMatmulPreferenceCreate(&preference));
        size_t limit = WORKSPACE_LIMIT;
        check_blas(cublasLtMatmulPreferenceSetAttribute(
            preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
            &limit, sizeof(limit)));

        cublasLtMatmulHeuristicResult_t candidates[8]{};
        int count = 0;
        const auto status = cublasLtMatmulAlgoGetHeuristic(
            handle, operation, a, b, c, c, preference,
            8, candidates, &count);
        cublasLtMatmulPreferenceDestroy(preference);

        if (status != CUBLAS_STATUS_SUCCESS) return false;
        for (int i = 0; i < count; ++i) {
            if (candidates[i].state == CUBLAS_STATUS_SUCCESS) {
                algorithm = candidates[i].algo;
                workspace_size = candidates[i].workspaceSize;
                direct_float = direct;
                return true;
            }
        }
        return false;
    }

    GemmPlan(int m, int n, int k, int dev)
        : rows(m), cols(n), inner(k), device(dev) {
        check_blas(cublasLtCreate(&handle));
        // Allow tensor-core conversion inside GEMM, eliminating the
        // explicit conversion pass on implementations supporting it.
        if (!configure(true)) {
            TORCH_CHECK(configure(false),
                        "No suitable cuBLASLt GEMM algorithm");
        }
    }

    ~GemmPlan() {
        clear_descriptors();
        if (handle) cublasLtDestroy(handle);
    }
};

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

__global__ void convert_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;
    const int64_t stride = (int64_t)blockDim.x * gridDim.x;
    for (; i < total_vectors; i += stride) {
        float4 v;
        Half4* dst;
        if (i < input_vectors) {
            v = reinterpret_cast<const float4*>(input)[i];
            dst = reinterpret_cast<Half4*>(half_input) + i;
        } else {
            const int64_t j = i - input_vectors;
            v = reinterpret_cast<const float4*>(weight)[j];
            dst = reinterpret_cast<Half4*>(half_weight) + j;
        }
        Half4 h;
        h.lo = __floats2half2_rn(v.x, v.y);
        h.hi = __floats2half2_rn(v.z, v.w);
        *dst = h;
    }
}

__global__ void convert_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;
    const 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
            half_weight[i - input_elements] =
                __float2half_rn(weight[i - input_elements]);
    }
}

__global__ void partial_stats(
    const float* __restrict__ x,
    float* __restrict__ partial,
    int64_t* __restrict__ batches,
    int rows, int cols, int parts)
{
    __shared__ float sums[8][32];
    __shared__ float squares[8][32];

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

    if (blockIdx.x == 0 && blockIdx.y == 0 && tx == 0 && ty == 0)
        batches[0] += 1;

    float anchor = 0.0f;
    float sum = 0.0f;
    float square = 0.0f;
    if (col < cols) {
        anchor = x[(int64_t)start * cols + col];
        for (int row = start + ty; row < end; row += 8) {
            const float d = x[(int64_t)row * cols + col] - anchor;
            sum += d;
            square = fmaf(d, d, square);
        }
    }

    sums[ty][tx] = sum;
    squares[ty][tx] = square;
    __syncthreads();

    if (ty == 0 && col < cols) {
        float total = 0.0f;
        float total_square = 0.0f;
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            total += sums[i][tx];
            total_square += squares[i][tx];
        }
        const float inverse_count = 1.0f / (end - start);
        const int64_t index = (int64_t)blockIdx.y * cols + col;
        partial[index] = anchor + total * inverse_count;
        partial[(int64_t)parts * cols + index] =
            fmaxf(fmaf(-total, total * inverse_count, total_square), 0.0f);
    }
}

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

    if (cumulative) momentum = (float)(1.0 / (double)batches[0]);
    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)
{
    const 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];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    #pragma unroll
    for (int offset = 16; offset; offset >>= 1) {
        const 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; offset >>= 1) {
            const float other =
                __shfl_down_sync(0xffffffff, value, offset);
            if (MAXIMUM) value = fmaxf(value, other);
            else value += other;
        }
        if (lane == 0) scratch[0] = value;
    }

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

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

    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        const int col = tid + i * THREADS;
        float value = -INFINITY;
        if (col < cols) {
            value = fmaf(
                x[base + col] - 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];
    }

    const float inverse_sum = 1.0f / block_reduce<false>(sum);
    #pragma unroll
    for (int i = 0; i < ITEMS; ++i) {
        const 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 batches,
    torch::Tensor scale, bool training, double eps,
    double momentum, bool cumulative)
{
    const c10::cuda::CUDAGuard guard(input.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

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

    static thread_local std::unique_ptr<GemmPlan> plan;
    if (!plan || plan->rows != rows || plan->cols != cols ||
        plan->inner != inner || plan->device != device)
        plan.reset(new GemmPlan(rows, cols, inner, device));

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

    torch::Tensor half_input, half_weight;
    const void* a_ptr = linear_weight.data_ptr<float>();
    const void* b_ptr = input.data_ptr<float>();

    if (!plan->direct_float) {
        half_input = torch::empty(
            {rows, inner}, input.options().dtype(torch::kFloat16));
        half_weight = torch::empty(
            {cols, inner}, input.options().dtype(torch::kFloat16));

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

        const int64_t nx = input.numel();
        const int64_t nw = linear_weight.numel();
        const int block_limit =
            8 * at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
        const 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) {
            const int64_t count = (nx + nw) / 4;
            const int blocks = static_cast<int>(
                std::min<int64_t>((count + 255) / 256, block_limit));
            convert_vectorized<<<blocks, 256, 0, stream>>>(
                input_ptr, weight_ptr, hi, hw, nx / 4, count);
        } else {
            const int64_t count = nx + nw;
            const int blocks = static_cast<int>(
                std::min<int64_t>((count + 255) / 256, block_limit));
            convert_scalar<<<blocks, 256, 0, stream>>>(
                input_ptr, weight_ptr, hi, hw, nx, count);
        }
        a_ptr = hw;
        b_ptr = hi;
    }

    const void* bias_ptr = linear_bias.data_ptr<float>();
    check_blas(cublasLtMatmulDescSetAttribute(
        plan->operation, CUBLASLT_MATMUL_DESC_BIAS_POINTER,
        &bias_ptr, sizeof(bias_ptr)));

    auto workspace = torch::empty(
        {static_cast<int64_t>(plan->workspace_size)},
        input.options().dtype(torch::kUInt8));

    const float alpha = 1.0f, beta = 0.0f;
    check_blas(cublasLtMatmul(
        plan->handle, plan->operation, &alpha,
        a_ptr, plan->a, b_ptr, plan->b, &beta,
        x.data_ptr<float>(), plan->c,
        x.data_ptr<float>(), plan->c,
        &plan->algorithm,
        plan->workspace_size ? workspace.data_ptr() : nullptr,
        plan->workspace_size, stream));

    auto coefficients = torch::empty({3, cols}, input.options());
    if (training) {
        const 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>(), partial.data_ptr<float>(),
            batches.data_ptr<int64_t>(), 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>(),
            batches.data_ptr<int64_t>(), coefficients.data_ptr<float>(),
            rows, cols, parts, (float)eps, (float)momentum, cumulative);
    } 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);
    }

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

    #define LAUNCH(N) \
        normalized_softmax<N><<<rows, THREADS, 0, stream>>>( \
            x.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="direct_tensorcore_gemm_fused_bn_softmax_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=["-lcublasLt"],
    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,
        )
        batches = self.bn.num_batches_tracked

        supported = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.ndim == 2
            and x.is_contiguous()
            and x.size(0) > 0
            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
            )
            and batches is not None
            and batches.device == x.device
            and batches.dtype == torch.int64
            and batches.numel() == 1
            and batches.is_contiguous()
        )

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

        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,
            batches,
            self.scale,
            self.bn.training,
            self.bn.eps,
            0.0 if self.bn.momentum is None else self.bn.momentum,
            self.bn.momentum is None,
        )
```