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

# Fused CUDA implementation:
#   1. GEMM + bias is computed with a single cuBLAS call (at::addmm).
#   2. The (subtract -> multiply -> ReLU) chain is fused into ONE custom
#      CUDA kernel (vectorized with float4) that runs in-place on the GEMM
#      output, eliminating 3 separate elementwise kernel launches and their
#      intermediate global-memory traffic.

cpp_source = (
    "torch::Tensor fused_linear_sub_mul_relu(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double sub_val, double mul_val);"
)

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

// Vectorized (float4) fused subtract -> multiply -> ReLU kernel.
__global__ void sub_mul_relu_vec4_kernel(const float4* __restrict__ in,
                                         float4* __restrict__ out,
                                         const float sub_val,
                                         const float mul_val,
                                         const int n4) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n4) {
        float4 v = in[idx];
        v.x = fmaxf((v.x - sub_val) * mul_val, 0.0f);
        v.y = fmaxf((v.y - sub_val) * mul_val, 0.0f);
        v.z = fmaxf((v.z - sub_val) * mul_val, 0.0f);
        v.w = fmaxf((v.w - sub_val) * mul_val, 0.0f);
        out[idx] = v;
    }
}

// Scalar fallback for the (unlikely) case where numel is not a multiple of 4.
__global__ void sub_mul_relu_kernel(const float* __restrict__ in,
                                    float* __restrict__ out,
                                    const float sub_val,
                                    const float mul_val,
                                    const int n) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        out[idx] = fmaxf((in[idx] - sub_val) * mul_val, 0.0f);
    }
}

torch::Tensor fused_linear_sub_mul_relu(torch::Tensor x,
                                        torch::Tensor weight,
                                        torch::Tensor bias,
                                        double sub_val,
                                        double mul_val) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight 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");

    const c10::cuda::CUDAGuard device_guard(x.device());

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    auto b_c = bias.contiguous();

    // Single cuBLAS call: out = x @ weight^T + bias (bias fused via beta path).
    auto out = at::addmm(b_c, x_c, w_c.t());

    const int64_t n = out.numel();
    float* ptr = out.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int block = 256;
    if ((n & 3) == 0) {
        const int n4 = static_cast<int>(n >> 2);
        const int grid = (n4 + block - 1) / block;
        sub_mul_relu_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(ptr),
            reinterpret_cast<float4*>(ptr),
            static_cast<float>(sub_val),
            static_cast<float>(mul_val),
            n4);
    } else {
        const int grid = static_cast<int>((n + block - 1) / block);
        sub_mul_relu_kernel<<<grid, block, 0, stream>>>(
            ptr, ptr,
            static_cast<float>(sub_val),
            static_cast<float>(mul_val),
            static_cast<int>(n));
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

fused_ops = load_inline(
    name="fused_linear_sub_mul_relu_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_linear_sub_mul_relu"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM (with fused bias) followed by a single
    custom CUDA kernel that fuses subtract -> multiply -> ReLU.
    """
    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):
        return fused_ops.fused_linear_sub_mul_relu(
            x,
            self.linear.weight,
            self.linear.bias,
            self.subtract_value,
            self.multiply_value,
        )


batch_size = 1024
in_features = 8192
out_features = 8192
subtract_value = 2.0
multiply_value = 1.5


def get_inputs():
    return [torch.rand(batch_size, in_features)]


def get_init_inputs():
    return [in_features, out_features, subtract_value, multiply_value]
```