```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 (x @ W^T + b) is dispatched to cuBLAS via at::addmm (optimal for
#      the large 128 x 16384 x 16384 matmul).
#   2. The subsequent `torch.min(y, constant)` and `- constant` pointwise ops
#      are fused into a single custom CUDA kernel that runs in-place on the
#      GEMM output, eliminating one full kernel launch and one intermediate
#      8MB tensor round-trip through HBM.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

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

__global__ void min_sub_scalar_vec4_kernel(float4* __restrict__ data,
                                           const float* __restrict__ c_ptr,
                                           const long n4) {
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n4) return;
    const float c = *c_ptr;
    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;
}

torch::Tensor linear_min_sub(torch::Tensor x, torch::Tensor weight,
                             torch::Tensor bias, torch::Tensor constant) {
    TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda() &&
                constant.is_cuda(), "All tensors must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "Expected float32 input");

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

    const long n = z.numel();
    if (n == 0) return z;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    float* ptr = z.data_ptr<float>();
    const float* c_ptr = constant.data_ptr<float>();

    const int block = 256;
    if ((n & 3) == 0) {
        // Vectorized path: one float4 (16 bytes) per thread.
        const long n4 = n >> 2;
        const long grid = (n4 + block - 1) / block;
        min_sub_scalar_vec4_kernel<<<grid, block, 0, stream>>>(
            reinterpret_cast<float4*>(ptr), c_ptr, n4);
    } else {
        const long grid = (n + block - 1) / block;
        min_sub_scalar_kernel<<<grid, block, 0, stream>>>(ptr, c_ptr, n);
    }

    return z;
}
"""

cpp_source = (
    "torch::Tensor linear_min_sub(torch::Tensor x, torch::Tensor weight, "
    "torch::Tensor bias, torch::Tensor constant);"
)

fused_linear_min_sub = load_inline(
    name="fused_linear_min_sub",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["linear_min_sub"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: GEMM via cuBLAS with the min/subtract epilogue fused
    into a single custom CUDA 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))

    def forward(self, x):
        return fused_linear_min_sub.linear_min_sub(
            x, self.linear.weight, self.linear.bias, self.constant
        )
```