```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 epilogue: y = relu((x - sub) * mul)
# The three separate element-wise ops (sub, mul, relu) are fused into a single
# memory-bound kernel using vectorized float4 accesses for maximum bandwidth.
# The GEMM (the compute-dominant part) is left to the highly-optimized cuBLAS
# path via torch.addmm, and the epilogue is applied in-place on its output,
# eliminating 3 extra kernel launches and ~4 full tensor round-trips to HBM.
fused_epilogue_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

__global__ void sub_mul_relu_vec4_kernel(float4* __restrict__ data,
                                         const float sub,
                                         const float mul,
                                         const long n4) {
    const long stride = static_cast<long>(gridDim.x) * blockDim.x;
    for (long i = static_cast<long>(blockIdx.x) * blockDim.x + threadIdx.x;
         i < n4; i += stride) {
        float4 v = data[i];
        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);
        data[i] = v;
    }
}

__global__ void sub_mul_relu_scalar_kernel(float* __restrict__ data,
                                           const float sub,
                                           const float mul,
                                           const long start,
                                           const long n) {
    const long stride = static_cast<long>(gridDim.x) * blockDim.x;
    for (long i = start + static_cast<long>(blockIdx.x) * blockDim.x + threadIdx.x;
         i < n; i += stride) {
        const float v = data[i];
        data[i] = fmaxf((v - sub) * mul, 0.0f);
    }
}

torch::Tensor sub_mul_relu_inplace(torch::Tensor x, double sub, double mul) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

    const long n = x.numel();
    if (n == 0) {
        return x;
    }

    float* ptr = x.data_ptr<float>();
    const float s = static_cast<float>(sub);
    const float m = static_cast<float>(mul);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int threads = 256;
    long vec_end = 0;

    // Use 16-byte vectorized accesses when the pointer is aligned.
    if ((reinterpret_cast<uintptr_t>(ptr) & 0xF) == 0) {
        const long n4 = n >> 2;
        if (n4 > 0) {
            long blocks = (n4 + threads - 1) / threads;
            if (blocks > 65535L) blocks = 65535L;
            sub_mul_relu_vec4_kernel<<<static_cast<unsigned int>(blocks),
                                       threads, 0, stream>>>(
                reinterpret_cast<float4*>(ptr), s, m, n4);
        }
        vec_end = n4 << 2;
    }

    // Handle the tail (or the whole tensor if unaligned).
    if (vec_end < n) {
        const long rem = n - vec_end;
        long blocks = (rem + threads - 1) / threads;
        if (blocks > 65535L) blocks = 65535L;
        sub_mul_relu_scalar_kernel<<<static_cast<unsigned int>(blocks),
                                     threads, 0, stream>>>(
            ptr, s, m, vec_end, n);
    }

    return x;
}
"""

fused_epilogue_cpp_src = (
    "torch::Tensor sub_mul_relu_inplace(torch::Tensor x, double sub, double mul);"
)

# Compile the inline CUDA extension.
fused_epilogue = load_inline(
    name="fused_sub_mul_relu_epilogue",
    cpp_sources=fused_epilogue_cpp_src,
    cuda_sources=fused_epilogue_cuda_src,
    functions=["sub_mul_relu_inplace"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: matrix multiplication (cuBLAS via addmm) followed by a
    single fused CUDA kernel performing subtract -> multiply -> ReLU in-place.
    """
    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):
        lead_shape = x.shape[:-1]
        squeezed = False
        if x.dim() == 1:
            x = x.unsqueeze(0)
            squeezed = True

        x2 = x.contiguous().view(-1, x.shape[-1])

        # GEMM + bias (cuBLAS, bias fused through cublasLt epilogue).
        out = torch.addmm(self.linear.bias, x2, self.linear.weight.t())

        # Fused (subtract, multiply, ReLU) in a single custom CUDA kernel.
        out = fused_epilogue.sub_mul_relu_inplace(
            out, self.subtract_value, self.multiply_value
        )

        out = out.view(*lead_shape, out.shape[-1])
        if squeezed:
            out = out.squeeze(0)
        return out
```