```python
# 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

# Custom CUDA kernels:
#  1) compute_affine_coef_kernel: folds BatchNorm (running or batch stats) and the
#     scalar "scale" parameter into a per-channel affine transform z = a * x + b.
#  2) fused_affine_softmax_kernel: one thread block per row applies the affine
#     transform (values cached in registers), then computes a numerically stable
#     softmax with two warp/shared-memory reductions. Reads the GEMM output once
#     and writes the final probabilities once.
# The GEMM itself stays with cuBLAS (nn.Linear), which is already near-optimal
# for fp32; everything after it is fused into a single pass over the data.
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <math.h>

#define BLOCK_SIZE 256
#define NUM_WARPS (BLOCK_SIZE / 32)
#define MAX_ITERS 8

__global__ void compute_affine_coef_kernel(
    const float* __restrict__ mean,
    const float* __restrict__ var,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ scale,
    float* __restrict__ a,
    float* __restrict__ b,
    float eps,
    int n)
{
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j < n) {
        float s = scale[0];
        float inv_std = 1.0f / sqrtf(var[j] + eps);
        float g = gamma[j];
        a[j] = s * g * inv_std;
        b[j] = s * (beta[j] - mean[j] * g * inv_std);
    }
}

__device__ __forceinline__ float warp_reduce_max(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, offset));
    }
    return v;
}

__device__ __forceinline__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_xor_sync(0xffffffffu, v, offset);
    }
    return v;
}

// One block per row. Requires nvec == iters * BLOCK_SIZE with iters <= MAX_ITERS.
__global__ void __launch_bounds__(BLOCK_SIZE) fused_affine_softmax_kernel(
    const float4* __restrict__ x,
    const float4* __restrict__ a,
    const float4* __restrict__ b,
    float4* __restrict__ out,
    int nvec,
    int iters)
{
    const int row = blockIdx.x;
    const float4* __restrict__ x_row = x + (size_t)row * nvec;
    float4* __restrict__ out_row = out + (size_t)row * nvec;

    __shared__ float smax[NUM_WARPS];
    __shared__ float ssum[NUM_WARPS];
    __shared__ float srow_max;
    __shared__ float srow_sum;

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int wid = tid >> 5;

    float4 z[MAX_ITERS];

    // Pass 1: z = a * x + b, keep z in registers, track row max.
    float m = -INFINITY;
    #pragma unroll
    for (int it = 0; it < MAX_ITERS; ++it) {
        if (it < iters) {
            int idx = tid + it * BLOCK_SIZE;
            float4 xv = x_row[idx];
            float4 av = a[idx];
            float4 bv = b[idx];
            float4 zv;
            zv.x = fmaf(av.x, xv.x, bv.x);
            zv.y = fmaf(av.y, xv.y, bv.y);
            zv.z = fmaf(av.z, xv.z, bv.z);
            zv.w = fmaf(av.w, xv.w, bv.w);
            z[it] = zv;
            m = fmaxf(m, fmaxf(fmaxf(zv.x, zv.y), fmaxf(zv.z, zv.w)));
        }
    }
    m = warp_reduce_max(m);
    if (lane == 0) smax[wid] = m;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < NUM_WARPS) ? smax[lane] : -INFINITY;
        v = warp_reduce_max(v);
        if (lane == 0) srow_max = v;
    }
    __syncthreads();
    const float row_max = srow_max;

    // Pass 2: exponentiate (from registers) and accumulate partition sum.
    float s = 0.0f;
    #pragma unroll
    for (int it = 0; it < MAX_ITERS; ++it) {
        if (it < iters) {
            float4 zv = z[it];
            zv.x = expf(zv.x - row_max);
            zv.y = expf(zv.y - row_max);
            zv.z = expf(zv.z - row_max);
            zv.w = expf(zv.w - row_max);
            z[it] = zv;
            s += (zv.x + zv.y) + (zv.z + zv.w);
        }
    }
    s = warp_reduce_sum(s);
    if (lane == 0) ssum[wid] = s;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < NUM_WARPS) ? ssum[lane] : 0.0f;
        v = warp_reduce_sum(v);
        if (lane == 0) srow_sum = v;
    }
    __syncthreads();
    const float inv_sum = 1.0f / srow_sum;

    // Pass 3: normalize and store.
    #pragma unroll
    for (int it = 0; it < MAX_ITERS; ++it) {
        if (it < iters) {
            int idx = tid + it * BLOCK_SIZE;
            float4 zv = z[it];
            zv.x *= inv_sum;
            zv.y *= inv_sum;
            zv.z *= inv_sum;
            zv.w *= inv_sum;
            out_row[idx] = zv;
        }
    }
}

torch::Tensor fused_bn_scale_softmax(
    torch::Tensor x,
    torch::Tensor mean,
    torch::Tensor var,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor scale,
    double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    TORCH_CHECK(x.is_contiguous() && mean.is_contiguous() && var.is_contiguous() &&
                gamma.is_contiguous() && beta.is_contiguous() && scale.is_contiguous(),
                "all tensors must be contiguous");

    const int batch = x.size(0);
    const int n = x.size(1);
    TORCH_CHECK(n % (4 * BLOCK_SIZE) == 0, "feature dim must be divisible by 1024");
    const int nvec = n / 4;
    const int iters = nvec / BLOCK_SIZE;
    TORCH_CHECK(iters >= 1 && iters <= MAX_ITERS, "unsupported feature dim");

    auto a = torch::empty_like(mean);
    auto b = torch::empty_like(mean);
    auto out = torch::empty_like(x);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int coef_threads = 256;
    const int coef_blocks = (n + coef_threads - 1) / coef_threads;
    compute_affine_coef_kernel<<<coef_blocks, coef_threads, 0, stream>>>(
        mean.data_ptr<float>(),
        var.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        scale.data_ptr<float>(),
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        static_cast<float>(eps),
        n);

    fused_affine_softmax_kernel<<<batch, BLOCK_SIZE, 0, stream>>>(
        reinterpret_cast<const float4*>(x.data_ptr<float>()),
        reinterpret_cast<const float4*>(a.data_ptr<float>()),
        reinterpret_cast<const float4*>(b.data_ptr<float>()),
        reinterpret_cast<float4*>(out.data_ptr<float>()),
        nvec,
        iters);

    return out;
}
"""

_cpp_source = (
    "torch::Tensor fused_bn_scale_softmax(torch::Tensor x, torch::Tensor mean, "
    "torch::Tensor var, torch::Tensor gamma, torch::Tensor beta, "
    "torch::Tensor scale, double eps);"
)

_fused_ops = load_inline(
    name="fused_bn_scale_softmax_ext",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_bn_scale_softmax"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

_BLOCK_SIZE = 256
_MAX_ITERS = 8


class ModelNew(nn.Module):
    """
    Optimized version of Model. The Linear (GEMM) layer is left to cuBLAS, while
    BatchNorm + scalar scaling + Softmax are fused into a single custom CUDA
    kernel (per-channel affine coefficients are folded together first).
    """

    def __init__(self, in_features, out_features, bn_eps=1e-5, bn_momentum=0.1, scale_shape=(1,)):
        super(ModelNew, self).__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.out_features = out_features

    def forward(self, x):
        x = F.linear(x, self.gemm.weight, self.gemm.bias)
        bn = self.bn
        f = self.out_features

        fused_ok = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 2
            and bn.affine
            and bn.track_running_stats
            and bn.momentum is not None
            and self.scale.numel() == 1
            and (f % (4 * _BLOCK_SIZE)) == 0
            and 1 <= (f // (4 * _BLOCK_SIZE)) <= _MAX_ITERS
        )

        if fused_ok:
            if bn.training:
                # Batch statistics (biased variance), exactly like BatchNorm1d.
                var, mean = torch.var_mean(x, dim=0, unbiased=False)
                with torch.no_grad():
                    bn.num_batches_tracked.add_(1)
                    n = x.size(0)
                    m = bn.momentum
                    bn.running_mean.mul_(1.0 - m).add_(mean, alpha=m)
                    bn.running_var.mul_(1.0 - m).add_(var, alpha=m * (float(n) / float(n - 1)))
                mean = mean.contiguous()
                var = var.contiguous()
            else:
                mean = bn.running_mean
                var = bn.running_var

            return _fused_ops.fused_bn_scale_softmax(
                x, mean, var, bn.weight, bn.bias, self.scale, float(bn.eps)
            )

        # Generic fallback: original operator sequence.
        x = bn(x)
        x = self.scale * x
        return F.softmax(x, dim=1)
```