```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(weight) -> Swish
# The GEMM is kept as-is since cuBLAS is already near-optimal for it.
# Everything after the GEMM is memory-bound, so fusing it into a single
# kernel eliminates multiple full read/write passes over the activation.
fused_gn_act_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstdint>

#define FULL_MASK 0xffffffffu

// One warp handles one (batch, group) pair. Each lane holds one or more
// channels of that group, warp reductions produce mean/var, then all
// elementwise ops are applied in registers before a single store.
__global__ void fused_gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ weight,
    float* __restrict__ out,
    const int batch_size,
    const int num_channels,
    const int num_groups,
    const int channels_per_group,
    const float eps)
{
    const int lane = threadIdx.x & 31;
    const int warp_id = threadIdx.x >> 5;
    const int warps_per_block = blockDim.x >> 5;
    const int group_global = blockIdx.x * warps_per_block + warp_id;
    const int total_groups = batch_size * num_groups;
    if (group_global >= total_groups) return;

    const int b = group_global / num_groups;
    const int g = group_global - b * num_groups;

    const int64_t base = (int64_t)b * num_channels + (int64_t)g * channels_per_group;
    const float* x_ptr = x + base;
    float* out_ptr = out + base;

    const int iters = (channels_per_group + 31) >> 5;

    // Load channels of this group into registers, coalesced.
    float vals[8];
    float sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        if (i < iters) {
            const int idx = (i << 5) + lane;
            if (idx < channels_per_group) {
                vals[i] = x_ptr[idx];
                sum += vals[i];
            }
        }
    }

    // Warp reduction for mean.
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_down_sync(FULL_MASK, sum, offset);
    sum = __shfl_sync(FULL_MASK, sum, 0);
    const float mean = sum / (float)channels_per_group;

    // Warp reduction for variance (biased, as in GroupNorm).
    float var_sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        if (i < iters) {
            const int idx = (i << 5) + lane;
            if (idx < channels_per_group) {
                const float d = vals[i] - mean;
                var_sum += d * d;
            }
        }
    }
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        var_sum += __shfl_down_sync(FULL_MASK, var_sum, offset);
    var_sum = __shfl_sync(FULL_MASK, var_sum, 0);
    const float rstd = rsqrtf(var_sum / (float)channels_per_group + eps);

    // GroupNorm affine -> Swish -> Multiply -> Swish, fully in registers.
    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        if (i < iters) {
            const int idx = (i << 5) + lane;
            if (idx < channels_per_group) {
                const int ch = g * channels_per_group + idx;
                const float y = (vals[i] - mean) * rstd * gamma[ch] + beta[ch];
                const float s = y / (1.0f + expf(-y));   // swish
                const float m = s * weight[ch];          // multiply
                out_ptr[idx] = m / (1.0f + expf(-m));    // swish
            }
        }
    }
}

torch::Tensor fused_gn_swish_mul_swish_cuda(
    torch::Tensor x,
    torch::Tensor gamma,
    torch::Tensor beta,
    torch::Tensor weight,
    int64_t num_groups,
    double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (batch, channels)");

    auto x_c = x.contiguous();
    auto out = torch::empty_like(x_c);

    const int batch_size = (int)x_c.size(0);
    const int num_channels = (int)x_c.size(1);
    const int channels_per_group = num_channels / (int)num_groups;
    TORCH_CHECK(num_channels % (int)num_groups == 0, "channels must divide groups");
    TORCH_CHECK(channels_per_group <= 256, "channels_per_group too large");

    const int total_groups = batch_size * (int)num_groups;
    const int warps_per_block = 8;
    const int block_size = warps_per_block * 32;
    const int num_blocks = (total_groups + warps_per_block - 1) / warps_per_block;

    fused_gn_swish_mul_swish_kernel<<<num_blocks, block_size>>>(
        x_c.data_ptr<float>(),
        gamma.data_ptr<float>(),
        beta.data_ptr<float>(),
        weight.data_ptr<float>(),
        out.data_ptr<float>(),
        batch_size,
        num_channels,
        (int)num_groups,
        channels_per_group,
        (float)eps);

    return out;
}
"""

fused_gn_act_cpp_source = (
    "torch::Tensor fused_gn_swish_mul_swish_cuda(torch::Tensor x, torch::Tensor gamma, "
    "torch::Tensor beta, torch::Tensor weight, int64_t num_groups, double eps);"
)

fused_gn_act = load_inline(
    name="fused_gn_act",
    cpp_sources=fused_gn_act_cpp_source,
    cuda_sources=fused_gn_act_cuda_source,
    functions=["fused_gn_swish_mul_swish_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: GEMM stays on cuBLAS, while GroupNorm + Swish +
    Multiply + Swish are fused into a single custom 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.fused_op = fused_gn_act

    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_op.fused_gn_swish_mul_swish_cuda(
            x,
            self.group_norm.weight,
            self.group_norm.bias,
            self.multiply_weight,
            self.num_groups,
            self.group_norm.eps,
        )
        return x
```