```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:
#   1. Runs the large GEMM (linear) through cuBLAS (at::addmm), which is
#      already near-peak for a 128x16384x16384 matmul.
#   2. Fuses the two element-wise ops (torch.min(x, constant) and
#      x - constant) into a single in-place epilogue kernel, using
#      vectorized float4 loads/stores to saturate memory bandwidth.
# This removes one full read+write pass over the 8MB activation tensor,
# one kernel launch, and one allocation compared to the eager version.
fused_linear_min_sub_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>
#include <cstdint>

__global__ void min_sub_epilogue_vec4_kernel(float4* __restrict__ data,
                                             const float c,
                                             const int n4) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n4) {
        float4 v = data[idx];
        v.x = fminf(v.x, c) - c;
        v.y = fminf(v.y, c) - c;
        v.z = fminf(v.z, c) - c;
        v.w = fminf(v.w, c) - c;
        data[idx] = v;
    }
}

__global__ void min_sub_epilogue_scalar_kernel(float* __restrict__ data,
                                               const float c,
                                               const long n) {
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        data[idx] = fminf(data[idx], c) - c;
    }
}

torch::Tensor linear_min_sub_cuda(torch::Tensor x,
                                  torch::Tensor weight,
                                  torch::Tensor bias,
                                  double c) {
    const at::cuda::OptionalCUDAGuard device_guard(at::device_of(x));

    // GEMM + bias via cuBLAS: y = x @ weight^T + bias
    auto y = at::addmm(bias, x, weight.t());

    // Fused epilogue: y = min(y, c) - c, done in-place.
    const long n = y.numel();
    const float cf = static_cast<float>(c);
    float* ptr = y.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;

    if ((n % 4 == 0) && ((reinterpret_cast<uintptr_t>(ptr) & 0xF) == 0)) {
        const int n4 = static_cast<int>(n / 4);
        const int blocks = (n4 + threads - 1) / threads;
        min_sub_epilogue_vec4_kernel<<<blocks, threads, 0, stream>>>(
            reinterpret_cast<float4*>(ptr), cf, n4);
    } else {
        const long blocks = (n + threads - 1) / threads;
        min_sub_epilogue_scalar_kernel<<<blocks, threads, 0, stream>>>(
            ptr, cf, n);
    }

    return y;
}
"""

fused_linear_min_sub_cpp_source = (
    "torch::Tensor linear_min_sub_cuda(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, double c);"
)

fused_linear_min_sub = load_inline(
    name="fused_linear_min_sub",
    cpp_sources=fused_linear_min_sub_cpp_source,
    cuda_sources=fused_linear_min_sub_cuda_source,
    functions=["linear_min_sub_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS GEMM + bias with a fused CUDA epilogue that
    performs min(x, constant) - constant in a single in-place kernel.
    """
    def __init__(self, in_features, out_features, constant):
        super(ModelNew, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.constant = nn.Parameter(torch.tensor(constant))
        # Cache the scalar value so forward avoids any host<->device sync.
        self._constant_value = float(constant)
        self._op = fused_linear_min_sub

    def forward(self, x):
        return self._op.linear_min_sub_cuda(
            x,
            self.linear.weight,
            self.linear.bias,
            self._constant_value,
        )
```