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

# Fused kernel: GroupNorm + Swish + Multiply + Swish
# One warp handles one (batch_row, group) pair so statistics are computed
# with fast warp shuffles, and each element is read from global memory once.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

namespace {

constexpr unsigned int FULL_MASK = 0xffffffffu;

__global__ void fused_gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ mul_w,
    float* __restrict__ out,
    const int batch_size,
    const int num_groups,
    const int C,
    const int group_size,
    const float eps)
{
    const int warps_per_block = blockDim.x >> 5;
    const int warp_global = blockIdx.x * warps_per_block + (threadIdx.x >> 5);
    const int lane = threadIdx.x & 31;
    const int total_warps = batch_size * num_groups;
    if (warp_global >= total_warps) return;

    const int n = warp_global / num_groups;
    const int g = warp_global - n * num_groups;
    const int c0 = g * group_size;

    const float* __restrict__ row  = x   + static_cast<long long>(n) * C;
    float* __restrict__       orow = out + static_cast<long long>(n) * C;

    if (group_size <= 32) {
        // Fast path: one element per lane, kept in registers.
        const bool active = lane < group_size;
        float v = 0.f;
        if (active) v = row[c0 + lane];

        float sum = v;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(FULL_MASK, sum, off);
        sum = __shfl_sync(FULL_MASK, sum, 0);
        const float mean = sum / static_cast<float>(group_size);

        float d  = active ? (v - mean) : 0.f;
        float sq = d * d;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sq += __shfl_down_sync(FULL_MASK, sq, off);
        sq = __shfl_sync(FULL_MASK, sq, 0);
        const float var  = sq / static_cast<float>(group_size);
        const float rstd = rsqrtf(var + eps);

        if (active) {
            const int c = c0 + lane;
            float y = (v - mean) * rstd * __ldg(gamma + c) + __ldg(beta + c);
            float s = y / (1.f + __expf(-y));   // swish: y * sigmoid(y)
            float m = s * __ldg(mul_w + c);
            orow[c] = m / (1.f + __expf(-m));   // swish again
        }
    } else {
        // General path: strided loop over the group.
        float sum = 0.f;
        for (int i = lane; i < group_size; i += 32) sum += row[c0 + i];
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(FULL_MASK, sum, off);
        sum = __shfl_sync(FULL_MASK, sum, 0);
        const float mean = sum / static_cast<float>(group_size);

        float sq = 0.f;
        for (int i = lane; i < group_size; i += 32) {
            float d = row[c0 + i] - mean;
            sq += d * d;
        }
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sq += __shfl_down_sync(FULL_MASK, sq, off);
        sq = __shfl_sync(FULL_MASK, sq, 0);
        const float var  = sq / static_cast<float>(group_size);
        const float rstd = rsqrtf(var + eps);

        for (int i = lane; i < group_size; i += 32) {
            const int c = c0 + i;
            float v = row[c];
            float y = (v - mean) * rstd * __ldg(gamma + c) + __ldg(beta + c);
            float s = y / (1.f + __expf(-y));
            float m = s * __ldg(mul_w + c);
            orow[c] = m / (1.f + __expf(-m));
        }
    }
}

} // namespace

torch::Tensor fused_gn_swish_mul_swish(
    torch::Tensor x,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor mul_w,
    int64_t num_groups,
    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 (N, C)");

    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();
    auto wc = mul_w.contiguous();

    const int batch_size = static_cast<int>(xc.size(0));
    const int C          = static_cast<int>(xc.size(1));
    TORCH_CHECK(C % num_groups == 0, "C must be divisible by num_groups");
    const int group_size = C / static_cast<int>(num_groups);

    auto out = torch::empty_like(xc);

    const int total_warps     = batch_size * static_cast<int>(num_groups);
    const int threads         = 256;
    const int warps_per_block = threads / 32;
    const int blocks          = (total_warps + warps_per_block - 1) / warps_per_block;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    fused_gn_swish_mul_swish_kernel<<<blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(),
        gc.data_ptr<float>(),
        bc.data_ptr<float>(),
        wc.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        static_cast<int>(num_groups),
        C,
        group_size,
        static_cast<float>(eps));

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_gn_swish_mul_swish(torch::Tensor x, torch::Tensor gamma, "
    "torch::Tensor beta, torch::Tensor mul_w, int64_t num_groups, double eps);"
)

fused_op = load_inline(
    name="fused_gn_swish_mul_swish",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["fused_gn_swish_mul_swish"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: GEMM is delegated to cuBLAS (F.linear), and
    GroupNorm + Swish + Multiply + Swish are fused into a single CUDA kernel.
    """
    def __init__(self, in_features, out_features, num_groups, multiply_weight_shape):
        super(ModelNew, self).__init__()
        self.gemm = nn.Linear(in_features, out_features)
        self.group_norm = nn.GroupNorm(num_groups, out_features)
        self.multiply_weight = nn.Parameter(torch.randn(multiply_weight_shape))
        self.num_groups = num_groups
        self.eps = self.group_norm.eps
        self.fused_op = fused_op

    def forward(self, x):
        # (batch_size, in_features) -> (batch_size, out_features) via cuBLAS
        x = F.linear(x, self.gemm.weight, self.gemm.bias)
        if x.is_cuda and x.dim() == 2 and x.dtype == torch.float32:
            # Fused GroupNorm + Swish + Multiply + Swish
            return self.fused_op.fused_gn_swish_mul_swish(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.eps,
            )
        # Fallback for CPU / unexpected input configurations
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```