```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: fused (bias add + subtract + multiply + ReLU) epilogue.
# The heavy GEMM is left to cuBLAS (torch.mm), and the entire elementwise
# chain is fused into a single memory-bound kernel using vectorized float4
# loads/stores, eliminating 3 kernel launches and intermediate tensor traffic.
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void linear_epilogue_vec4_kernel(
    const float4* __restrict__ y,
    const float4* __restrict__ bias,
    float4* __restrict__ out,
    const int vec_per_row,
    const float sub,
    const float mul)
{
    const int row = blockIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (col < vec_per_row) {
        const int idx = row * vec_per_row + col;
        const float4 v = y[idx];
        const float4 b = bias[col];
        float4 r;
        r.x = fmaxf((v.x + b.x - sub) * mul, 0.0f);
        r.y = fmaxf((v.y + b.y - sub) * mul, 0.0f);
        r.z = fmaxf((v.z + b.z - sub) * mul, 0.0f);
        r.w = fmaxf((v.w + b.w - sub) * mul, 0.0f);
        out[idx] = r;
    }
}

__global__ void linear_epilogue_scalar_kernel(
    const float* __restrict__ y,
    const float* __restrict__ bias,
    float* __restrict__ out,
    const long long total,
    const int features,
    const float sub,
    const float mul)
{
    const long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        const int col = (int)(idx % (long long)features);
        out[idx] = fmaxf((y[idx] + bias[col] - sub) * mul, 0.0f);
    }
}

torch::Tensor linear_epilogue_cuda(torch::Tensor y, torch::Tensor bias, double sub, double mul) {
    TORCH_CHECK(y.is_cuda() && bias.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32 && bias.scalar_type() == torch::kFloat32,
                "only float32 is supported");
    TORCH_CHECK(y.is_contiguous() && bias.is_contiguous(), "inputs must be contiguous");

    auto out = torch::empty_like(y);
    const long long total = y.numel();
    const int features = (int)y.size(-1);
    const long long rows = total / features;
    const float s = (float)sub;
    const float m = (float)mul;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (features % 4 == 0 && rows <= 65535) {
        const int vec_per_row = features / 4;
        dim3 block(256);
        dim3 grid((vec_per_row + block.x - 1) / block.x, (unsigned int)rows);
        linear_epilogue_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(y.data_ptr<float>()),
            reinterpret_cast<const float4*>(bias.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            vec_per_row, s, m);
    } else {
        const int block = 256;
        const long long blocks = (total + block - 1) / block;
        linear_epilogue_scalar_kernel<<<blocks, block, 0, stream>>>(
            y.data_ptr<float>(),
            bias.data_ptr<float>(),
            out.data_ptr<float>(),
            total, features, s, m);
    }
    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor linear_epilogue_cuda(torch::Tensor y, torch::Tensor bias, double sub, double mul);"
)

# Compile the inline CUDA extension
fused_epilogue = load_inline(
    name="linear_fused_epilogue",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["linear_epilogue_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM + custom fused CUDA epilogue that performs
    bias-add, subtraction, multiplication, and ReLU in a single kernel.
    """
    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 = float(subtract_value)
        self.multiply_value = float(multiply_value)
        self.fused_epilogue = fused_epilogue

    def forward(self, x):
        if x.is_cuda and x.dtype == torch.float32 and x.dim() == 2:
            # GEMM via cuBLAS, then one fused kernel for the whole epilogue.
            y = torch.mm(x, self.linear.weight.t())
            return self.fused_epilogue.linear_epilogue_cuda(
                y, self.linear.bias, self.subtract_value, self.multiply_value
            )
        # Fallback for CPU / other dtypes / non-2D inputs.
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        return torch.relu(x)
```