```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 (subtract -> multiply -> ReLU) into a single
# memory-bound pass over the GEMM output. The GEMM itself is dispatched to
# cuBLAS (via nn.Linear), which is already near-optimal for this size, so the
# main win comes from replacing three separate elementwise kernel launches
# (and their intermediate global-memory traffic) with one vectorized kernel.
fused_epilogue_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void sub_mul_relu_vec4_kernel(const float* __restrict__ in,
                                         float* __restrict__ out,
                                         float sub, float mul, int n4) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n4) {
        float4 v = reinterpret_cast<const float4*>(in)[idx];
        float4 r;
        r.x = fmaxf((v.x - sub) * mul, 0.0f);
        r.y = fmaxf((v.y - sub) * mul, 0.0f);
        r.z = fmaxf((v.z - sub) * mul, 0.0f);
        r.w = fmaxf((v.w - sub) * mul, 0.0f);
        reinterpret_cast<float4*>(out)[idx] = r;
    }
}

__global__ void sub_mul_relu_kernel(const float* __restrict__ in,
                                    float* __restrict__ out,
                                    float sub, float mul, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        out[idx] = fmaxf((in[idx] - sub) * mul, 0.0f);
    }
}

torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul) {
    TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "input must be float32");

    auto x_c = x.contiguous();
    auto out = torch::empty_like(x_c);
    const int n = static_cast<int>(x_c.numel());
    const float s = static_cast<float>(sub);
    const float m = static_cast<float>(mul);

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

    if (n % 4 == 0) {
        const int n4 = n / 4;
        const int grid = (n4 + block - 1) / block;
        sub_mul_relu_vec4_kernel<<<grid, block, 0, stream>>>(
            x_c.data_ptr<float>(), out.data_ptr<float>(), s, m, n4);
    } else {
        const int grid = (n + block - 1) / block;
        sub_mul_relu_kernel<<<grid, block, 0, stream>>>(
            x_c.data_ptr<float>(), out.data_ptr<float>(), s, m, n);
    }

    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul);"
)

fused_epilogue = load_inline(
    name="fused_sub_mul_relu",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_source,
    functions=["sub_mul_relu_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM (nn.Linear) followed by a single fused
    CUDA kernel performing subtract -> multiply -> ReLU in one pass.
    """

    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

    def forward(self, x):
        x = self.linear(x)
        if x.is_cuda and x.dtype == torch.float32:
            return self.fused_epilogue.sub_mul_relu_cuda(
                x, float(self.subtract_value), float(self.multiply_value)
            )
        # Fallback for non-CUDA / non-float32 inputs
        x = x - self.subtract_value
        x = x * self.multiply_value
        return torch.relu(x)
```