```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 CUDA kernel: GroupNorm + Swish + Multiply + Swish
# The GEMM is kept as cuBLAS (F.linear) since it is already near-peak;
# all subsequent memory-bound ops are fused into a single kernel.
# One warp handles one (batch, group) pair: group statistics are computed
# with warp shuffles, then the elementwise chain is applied in-register.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ b,
    const float* __restrict__ mw,
    float* __restrict__ out,
    int C, int groups, float eps)
{
    const int group_size = C / groups;
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int warps_per_block = blockDim.x >> 5;
    const int gw = blockIdx.x * warps_per_block + warp;
    const int batch = gw / groups;
    const int g = gw - batch * groups;
    const int base = batch * C + g * group_size;

    // Pass 1: group mean / variance (fully coalesced loads)
    float sum = 0.f, sq = 0.f;
    for (int i = lane; i < group_size; i += 32) {
        float v = x[base + i];
        sum += v;
        sq += v * v;
    }
    #pragma unroll
    for (int m = 16; m > 0; m >>= 1) {
        sum += __shfl_xor_sync(0xffffffffu, sum, m);
        sq  += __shfl_xor_sync(0xffffffffu, sq,  m);
    }
    const float mean = sum / (float)group_size;
    float var = sq / (float)group_size - mean * mean;
    var = fmaxf(var, 0.f);
    const float rstd = rsqrtf(var + eps);

    // Pass 2: normalize -> swish -> multiply -> swish (values still in L1)
    for (int i = lane; i < group_size; i += 32) {
        const int idx = base + i;
        const int c = g * group_size + i;
        float v = (x[idx] - mean) * rstd * w[c] + b[c];
        v = v / (1.f + expf(-v));   // x * sigmoid(x)
        v = v * mw[c];
        v = v / (1.f + expf(-v));   // x * sigmoid(x)
        out[idx] = v;
    }
}

torch::Tensor gn_swish_mul_swish(torch::Tensor x, torch::Tensor w, torch::Tensor b,
                                 torch::Tensor mw, int64_t groups, double eps)
{
    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);
    const int C = xc.size(1);
    const long total = (long)xc.size(0) * (long)groups;
    const int threads = 256;
    const int wpb = threads / 32;
    const long blocks = (total + wpb - 1) / wpb;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gn_swish_mul_swish_kernel<<<(int)blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(), w.data_ptr<float>(), b.data_ptr<float>(),
        mw.data_ptr<float>(), out.data_ptr<float>(),
        C, (int)groups, (float)eps);
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor gn_swish_mul_swish(torch::Tensor x, torch::Tensor w, torch::Tensor b, "
    "torch::Tensor mw, int64_t groups, double eps);"
)

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


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM followed by a single fused CUDA kernel that
    performs GroupNorm + Swish + Multiply + Swish in one pass.
    """
    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 = float(self.group_norm.eps)
        self.fused_ops = fused_ops

    def forward(self, x):
        # (batch_size, in_features) -> (batch_size, out_features)
        x = F.linear(x, self.gemm.weight, self.gemm.bias)
        # Fused: GroupNorm -> Swish -> Multiply -> Swish
        x = self.fused_ops.gn_swish_mul_swish(
            x,
            self.group_norm.weight,
            self.group_norm.bias,
            self.multiply_weight,
            self.num_groups,
            self.eps,
        )
        return x
```