```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
# One warp handles one (sample, group) pair. With group_size == 32 (our case:
# 8192 channels / 256 groups) each lane owns exactly one channel, statistics
# are computed with warp-shuffle reductions (two-pass for numerical accuracy),
# and the whole elementwise chain is applied from registers with a single
# global read and a single global write.
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

__global__ void fused_gn_swish_mul_swish_kernel(
    const float* __restrict__ x,
    const float* __restrict__ gn_weight,
    const float* __restrict__ gn_bias,
    const float* __restrict__ mul_weight,
    float* __restrict__ out,
    const int num_rows,
    const int C,
    const int num_groups,
    const float eps)
{
    const int group_size = C / num_groups;
    const int lane = threadIdx.x & 31;
    const int warp_id = ((int)blockIdx.x * (int)blockDim.x + (int)threadIdx.x) >> 5;
    const int total_groups = num_rows * num_groups;
    if (warp_id >= total_groups) return;

    const long base = (long)warp_id * (long)group_size;
    const float* xg = x + base;
    float* og = out + base;
    const int c0 = (warp_id % num_groups) * group_size;

    if (group_size == 32) {
        // Fast path: one element per lane, everything stays in registers.
        const float v = xg[lane];

        float sum = v;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, off);
        sum = __shfl_sync(0xffffffffu, sum, 0);
        const float mean = sum * (1.0f / 32.0f);

        const float d = v - mean;
        float vs = d * d;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            vs += __shfl_down_sync(0xffffffffu, vs, off);
        vs = __shfl_sync(0xffffffffu, vs, 0);
        const float rstd = rsqrtf(vs * (1.0f / 32.0f) + eps);

        const int c = c0 + lane;
        float h = d * rstd * gn_weight[c] + gn_bias[c];
        float s = 1.0f / (1.0f + expf(-h));
        h = h * s;                    // Swish
        h = h * mul_weight[c];        // Multiply
        s = 1.0f / (1.0f + expf(-h));
        og[lane] = h * s;             // Swish
        return;
    }

    // Generic path for any group size.
    float sum = 0.0f;
    for (int i = lane; i < group_size; i += 32) sum += xg[i];
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);
    sum = __shfl_sync(0xffffffffu, sum, 0);
    const float mean = sum / (float)group_size;

    float vs = 0.0f;
    for (int i = lane; i < group_size; i += 32) {
        const float d = xg[i] - mean;
        vs += d * d;
    }
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        vs += __shfl_down_sync(0xffffffffu, vs, off);
    vs = __shfl_sync(0xffffffffu, vs, 0);
    const float rstd = rsqrtf(vs / (float)group_size + eps);

    for (int i = lane; i < group_size; i += 32) {
        const int c = c0 + i;
        float h = (xg[i] - mean) * rstd * gn_weight[c] + gn_bias[c];
        float s = 1.0f / (1.0f + expf(-h));
        h = h * s;
        h = h * mul_weight[c];
        s = 1.0f / (1.0f + expf(-h));
        og[i] = h * s;
    }
}

torch::Tensor fused_gn_swish_mul_swish_cuda(
    torch::Tensor x,
    torch::Tensor gn_weight,
    torch::Tensor gn_bias,
    torch::Tensor mul_weight,
    int64_t num_groups,
    double eps)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.dtype() == torch::kFloat32, "x must be float32");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (N, C)");

    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(x));

    auto xc = x.contiguous();
    auto wc = gn_weight.contiguous();
    auto bc = gn_bias.contiguous();
    auto mc = mul_weight.contiguous();
    auto out = torch::empty_like(xc);

    const int C = (int)xc.size(1);
    const int num_rows = (int)xc.size(0);
    const int total_groups = num_rows * (int)num_groups;

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

    auto stream = at::cuda::getCurrentCUDAStream();
    fused_gn_swish_mul_swish_kernel<<<blocks, threads, 0, stream>>>(
        xc.data_ptr<float>(),
        wc.data_ptr<float>(),
        bc.data_ptr<float>(),
        mc.data_ptr<float>(),
        out.data_ptr<float>(),
        num_rows, C, (int)num_groups, (float)eps);

    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor fused_gn_swish_mul_swish_cuda("
    "torch::Tensor x, torch::Tensor gn_weight, torch::Tensor gn_bias, "
    "torch::Tensor mul_weight, int64_t num_groups, double eps);"
)

_fused_ext = 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_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class _FusedGnSwishMulSwish(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, gn_weight, gn_bias, mul_weight, num_groups, eps):
        out = _fused_ext.fused_gn_swish_mul_swish_cuda(
            x, gn_weight, gn_bias, mul_weight, num_groups, eps
        )
        ctx.save_for_backward(x, gn_weight, gn_bias, mul_weight)
        ctx.num_groups = num_groups
        ctx.eps = eps
        return out

    @staticmethod
    def backward(ctx, grad_out):
        x, gn_weight, gn_bias, mul_weight = ctx.saved_tensors
        with torch.enable_grad():
            xr = x.detach().requires_grad_(True)
            wr = gn_weight.detach().requires_grad_(True)
            br = gn_bias.detach().requires_grad_(True)
            mr = mul_weight.detach().requires_grad_(True)
            h = F.group_norm(xr, ctx.num_groups, wr, br, ctx.eps)
            h = h * torch.sigmoid(h)
            h = h * mr
            h = h * torch.sigmoid(h)
            gx, gw, gb, gm = torch.autograd.grad(h, [xr, wr, br, mr], grad_out)
        return gx, gw, gb, gm, None, None


class ModelNew(nn.Module):
    """
    Optimized model: keeps the cuBLAS GEMM (nn.Linear) and fuses
    GroupNorm + Swish + Multiply + Swish 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))

    def forward(self, x):
        # GEMM (cuBLAS) dominates runtime; keep it on the highly tuned path.
        x = self.gemm(x)
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 2:
            x = _FusedGnSwishMulSwish.apply(
                x,
                self.group_norm.weight,
                self.group_norm.bias,
                self.multiply_weight,
                self.group_norm.num_groups,
                float(self.group_norm.eps),
            )
        else:
            x = self.group_norm(x)
            x = x * torch.sigmoid(x)
            x = x * self.multiply_weight
            x = x * torch.sigmoid(x)
        return x
```