```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
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernel that fuses the entire post-GEMM epilogue:
#   out = relu((x @ W^T + bias - subtract_value) * multiply_value)
# The GEMM itself is kept in cuBLAS (via torch.matmul) since it is the
# compute-bound part, while bias-add, subtraction, multiplication and ReLU
# (4 separate PyTorch kernels / memory passes) are fused into a single
# bandwidth-bound kernel.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <algorithm>

namespace {

__global__ void fused_bias_sub_mul_relu_scalar_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const long total,
    const int n,
    const float sub_val,
    const float mul_val)
{
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < total; i += stride) {
        const int col = (int)(i % n);
        const float v = (x[i] + bias[col] - sub_val) * mul_val;
        out[i] = fmaxf(v, 0.0f);
    }
}

__global__ void fused_bias_sub_mul_relu_vec4_kernel(
    const float4* __restrict__ x,
    const float4* __restrict__ bias,
    float4* __restrict__ out,
    const long total4,
    const int n4,
    const float sub_val,
    const float mul_val)
{
    const long stride = (long)gridDim.x * blockDim.x;
    for (long i = (long)blockIdx.x * blockDim.x + threadIdx.x; i < total4; i += stride) {
        const int col4 = (int)(i % n4);
        const float4 xv = x[i];
        const float4 bv = bias[col4];
        float4 r;
        r.x = fmaxf((xv.x + bv.x - sub_val) * mul_val, 0.0f);
        r.y = fmaxf((xv.y + bv.y - sub_val) * mul_val, 0.0f);
        r.z = fmaxf((xv.z + bv.z - sub_val) * mul_val, 0.0f);
        r.w = fmaxf((xv.w + bv.w - sub_val) * mul_val, 0.0f);
        out[i] = r;
    }
}

} // namespace

torch::Tensor fused_epilogue(torch::Tensor x, torch::Tensor bias, double sub_val, double mul_val) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");

    auto x_c = x.contiguous();
    auto bias_c = bias.contiguous();
    auto out = at::empty_like(x_c);

    const long total = (long)x_c.numel();
    if (total == 0) {
        return out;
    }
    const int n = (int)x_c.size(-1);

    const int block_size = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (n % 4 == 0) {
        const long total4 = total >> 2;
        const int n4 = n >> 2;
        const long blocks = (total4 + block_size - 1) / block_size;
        const int num_blocks = (int)std::min<long>(std::max<long>(blocks, 1L), 262144L);
        fused_bias_sub_mul_relu_vec4_kernel<<<num_blocks, block_size, 0, stream>>>(
            reinterpret_cast<const float4*>(x_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias_c.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            total4, n4, (float)sub_val, (float)mul_val);
    } else {
        const long blocks = (total + block_size - 1) / block_size;
        const int num_blocks = (int)std::min<long>(std::max<long>(blocks, 1L), 1048576L);
        fused_bias_sub_mul_relu_scalar_kernel<<<num_blocks, block_size, 0, stream>>>(
            x_c.data_ptr<float>(),
            bias_c.data_ptr<float>(),
            out.data_ptr<float>(),
            total, n, (float)sub_val, (float)mul_val);
    }
    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_epilogue(torch::Tensor x, torch::Tensor bias, double sub_val, double mul_val);"
)

fused_epilogue_mod = load_inline(
    name="fused_linear_epilogue",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_epilogue"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The four epilogue operators (bias add, subtraction, multiplication and
    ReLU) are fused into a single custom CUDA kernel, eliminating three
    extra kernel launches and three full read/write passes over the
    [batch_size, out_features] activation tensor. The GEMM stays in cuBLAS.
    """
    def __init__(self, in_features, out_features, subtract_value, multiply_value):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.subtract_value = subtract_value
        self.multiply_value = multiply_value
        self.fused_epilogue = fused_epilogue_mod

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32:
            # cuBLAS GEMM (no bias), then one fused custom kernel for
            # bias + subtract + multiply + ReLU.
            y = torch.matmul(x, self.linear.weight.t())
            return self.fused_epilogue.fused_epilogue(
                y,
                self.linear.bias,
                float(self.subtract_value),
                float(self.multiply_value),
            )
        # Fallback for non-CUDA tensors.
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        return torch.relu(x)
```