```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

# Custom CUDA kernel that fuses the whole element-wise epilogue:
#   out = relu((x - subtract_value) * multiply_value)
# The GEMM (the dominant cost) is left to cuBLAS via F.linear, and the
# subtract + multiply + ReLU chain is collapsed into a single vectorized
# (float4) kernel pass over the data.
fused_epilogue_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

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

__global__ void sub_mul_relu_scalar_kernel(const float* __restrict__ in,
                                           float* __restrict__ out,
                                           float sub, float mul,
                                           int offset, int n) {
    int idx = offset + 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.is_contiguous(), "input must be contiguous");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "input must be float32");

    auto out = torch::empty_like(x);
    const int n = (int)x.numel();
    if (n == 0) {
        return out;
    }

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

    const int n4 = n / 4;
    if (n4 > 0) {
        const int blocks = (n4 + block_size - 1) / block_size;
        sub_mul_relu_vec4_kernel<<<blocks, block_size, 0, stream>>>(
            reinterpret_cast<const float4*>(x.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            (float)sub, (float)mul, n4);
    }

    const int rem_start = n4 * 4;
    if (rem_start < n) {
        const int rem = n - rem_start;
        const int blocks = (rem + block_size - 1) / block_size;
        sub_mul_relu_scalar_kernel<<<blocks, block_size, 0, stream>>>(
            x.data_ptr<float>(), out.data_ptr<float>(),
            (float)sub, (float)mul, rem_start, 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_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: matmul + bias stays on cuBLAS (F.linear), while the
    subtract, scalar multiply, and ReLU are fused into a single custom CUDA
    kernel, eliminating three separate kernel launches and the associated
    extra global-memory round trips.
    """
    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)

    def forward(self, x):
        x = F.linear(x, self.linear.weight, self.linear.bias)
        if x.is_cuda:
            return fused_epilogue.sub_mul_relu_cuda(
                x, self.subtract_value, self.multiply_value
            )
        # CPU fallback for robustness
        return torch.relu((x - self.subtract_value) * self.multiply_value)
```