```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 GroupNorm + Swish + per-channel Multiply + Swish CUDA kernel.
# The GEMM (nn.Linear) is kept as-is since cuBLAS is already near-optimal.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

// Specialized kernel: channels-per-group == 32 -> exactly one warp per (sample, group).
__global__ void fused_gn_swish_mul_swish_cpg32_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ w,
    float* __restrict__ out,
    const int total_groups,
    const int G,
    const int C,
    const float eps)
{
    const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
    if (warp_id >= total_groups) return;
    const int lane = threadIdx.x & 31;

    const int n = warp_id / G;
    const int g = warp_id - n * G;
    const int base = n * C + (g << 5);
    const int c = (g << 5) + lane;

    const float v = x[base + lane];

    // Group mean via warp reduction.
    float sum = v;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_xor_sync(0xffffffffu, sum, offset);
    const float mean = sum * (1.0f / 32.0f);

    // Group variance (two-pass for numerical stability).
    const float d = v - mean;
    float varsum = d * d;
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        varsum += __shfl_xor_sync(0xffffffffu, varsum, offset);
    const float var = fmaxf(varsum * (1.0f / 32.0f), 0.0f);
    const float rstd = rsqrtf(var + eps);

    // GroupNorm affine.
    float y = d * rstd * gamma[c] + beta[c];
    // Swish: y * sigmoid(y) = y / (1 + exp(-y))
    y = y / (1.0f + expf(-y));
    // Per-channel multiply.
    y = y * w[c];
    // Swish again.
    y = y / (1.0f + expf(-y));

    out[base + lane] = y;
}

// Generic fallback: one warp per (sample, group), arbitrary channels-per-group.
__global__ void fused_gn_swish_mul_swish_generic_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ w,
    float* __restrict__ out,
    const int total_groups,
    const int G,
    const int C,
    const int cpg,
    const float eps)
{
    const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
    if (warp_id >= total_groups) return;
    const int lane = threadIdx.x & 31;

    const int n = warp_id / G;
    const int g = warp_id - n * G;
    const int base = n * C + g * cpg;

    float sum = 0.0f;
    for (int i = lane; i < cpg; i += 32) sum += x[base + i];
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_xor_sync(0xffffffffu, sum, offset);
    const float mean = sum / (float)cpg;

    float varsum = 0.0f;
    for (int i = lane; i < cpg; i += 32) {
        const float d = x[base + i] - mean;
        varsum += d * d;
    }
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        varsum += __shfl_xor_sync(0xffffffffu, varsum, offset);
    const float var = fmaxf(varsum / (float)cpg, 0.0f);
    const float rstd = rsqrtf(var + eps);

    for (int i = lane; i < cpg; i += 32) {
        const int c = g * cpg + i;
        float y = (x[base + i] - mean) * rstd * gamma[c] + beta[c];
        y = y / (1.0f + expf(-y));
        y = y * w[c];
        y = y / (1.0f + expf(-y));
        out[base + i] = y;
    }
}

torch::Tensor fused_gn_swish_mul_swish(
    torch::Tensor x,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor w,
    int64_t num_groups,
    double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (N, C)");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto gc = gamma.contiguous();
    auto bc = beta.contiguous();
    auto wc = w.contiguous();
    auto out = torch::empty_like(xc);

    const int N = (int)xc.size(0);
    const int C = (int)xc.size(1);
    const int G = (int)num_groups;
    const int cpg = C / G;
    const int total_groups = N * G;

    const int threads = 256;
    const int warps_per_block = threads / 32;
    const int blocks = (total_groups + warps_per_block - 1) / warps_per_block;

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

    if (cpg == 32) {
        fused_gn_swish_mul_swish_cpg32_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>(),
            total_groups, G, C, (float)eps);
    } else {
        fused_gn_swish_mul_swish_generic_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>(),
            total_groups, G, C, cpg, (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 w, int64_t num_groups, double eps);"
)

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


class ModelNew(nn.Module):
    """
    Optimized model: GEMM via cuBLAS (nn.Linear) followed by a single fused
    CUDA kernel performing GroupNorm + Swish + Multiply + Swish.
    """
    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

    def forward(self, x):
        # GEMM (batch_size, in_features) -> (batch_size, out_features)
        x = F.linear(x, self.gemm.weight, self.gemm.bias)

        if x.is_cuda and x.dtype == torch.float32:
            # Fused GroupNorm + Swish + Multiply + Swish in one kernel.
            return fused_ops.fused_gn_swish_mul_swish(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.num_groups,
                self.group_norm.eps,
            )

        # Fallback for non-CUDA / non-float32 inputs.
        x = self.group_norm(x)
        x = x * torch.sigmoid(x)
        x = x * self.multiply_weight
        x = x * torch.sigmoid(x)
        return x
```