```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 fused CUDA op:
#   out = min(x @ weight^T + bias, constant) - constant
# The GEMM itself is delegated to cuBLAS (at::mm) since it is near-optimal,
# while the bias-add, min and subtract are fused into a single vectorized
# epilogue kernel (eliminates 3 extra full read/write passes over the output).
_fused_cpp_source = """
torch::Tensor fused_linear_min_sub(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, torch::Tensor constant);
"""

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

namespace {

constexpr int kThreads = 256;

// Fused epilogue: out = fmin(mm + bias, c) - c
// Vectorized with float4 when possible; scalar tail handles the remainder.
__global__ void bias_min_sub_kernel(const float* __restrict__ mm,
                                    const float* __restrict__ bias,
                                    const float* __restrict__ constant,
                                    float* __restrict__ out,
                                    const int nvec,
                                    const int total,
                                    const int out_features)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    const float c = constant[0];

    if (i < nvec) {
        const int base = i << 2;
        // out_features % 4 == 0 here, so a float4 never straddles a row.
        const int col = base % out_features;
        float4 v = *reinterpret_cast<const float4*>(mm + base);
        const float4 b = *reinterpret_cast<const float4*>(bias + col);
        v.x = fminf(v.x + b.x, c) - c;
        v.y = fminf(v.y + b.y, c) - c;
        v.z = fminf(v.z + b.z, c) - c;
        v.w = fminf(v.w + b.w, c) - c;
        *reinterpret_cast<float4*>(out + base) = v;
    } else {
        const int idx = (nvec << 2) + (i - nvec);
        if (idx < total) {
            const int col = idx % out_features;
            float v = mm[idx];
            v = fminf(v + bias[col], c) - c;
            out[idx] = v;
        }
    }
}

} // namespace

torch::Tensor fused_linear_min_sub(torch::Tensor x,
                                   torch::Tensor weight,
                                   torch::Tensor bias,
                                   torch::Tensor constant)
{
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    TORCH_CHECK(weight.scalar_type() == at::kFloat, "weight must be float32");
    TORCH_CHECK(bias.scalar_type() == at::kFloat, "bias must be float32");

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

    auto x_c = x.contiguous();
    auto w_c = weight.contiguous();
    auto b_c = bias.contiguous();
    auto c_c = constant.to(x_c.device(), at::kFloat).contiguous();

    // Flatten leading dims so arbitrary input shapes work like nn.Linear.
    const auto x_sizes = x_c.sizes().vec();
    auto x2 = (x_c.dim() == 2) ? x_c : x_c.reshape({-1, x_c.size(-1)});

    // GEMM via cuBLAS.
    auto out = at::mm(x2, w_c.t());

    const long total = out.numel();
    const int out_features = static_cast<int>(out.size(-1));
    const int nvec = (out_features % 4 == 0) ? static_cast<int>(total >> 2) : 0;
    const int tail = static_cast<int>(total) - (nvec << 2);
    const int nthreads = nvec + tail;
    const int blocks = (nthreads + kThreads - 1) / kThreads;

    if (blocks > 0) {
        auto stream = at::cuda::getCurrentCUDAStream();
        bias_min_sub_kernel<<<blocks, kThreads, 0, stream>>>(
            out.data_ptr<float>(),
            b_c.data_ptr<float>(),
            c_c.data_ptr<float>(),
            out.data_ptr<float>(),
            nvec,
            static_cast<int>(total),
            out_features);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    if (x_c.dim() != 2) {
        auto out_sizes = x_sizes;
        out_sizes.back() = out_features;
        return out.reshape(out_sizes);
    }
    return out;
}
"""

_fused_ops = load_inline(
    name="fused_linear_min_sub_ext",
    cpp_sources=_fused_cpp_source,
    cuda_sources=_fused_cuda_source,
    functions=["fused_linear_min_sub"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The linear layer's GEMM is computed with cuBLAS, and the bias-add,
    torch.min(constant) and constant subtraction are fused into a single
    custom CUDA epilogue kernel, removing extra memory traffic and kernel
    launches. Parameter names are kept identical to the reference model so
    state_dicts are interchangeable.
    """
    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):
        if x.is_cuda and x.dtype == torch.float32:
            return _fused_ops.fused_linear_min_sub(
                x, self.linear.weight, self.linear.bias, self.constant
            )
        # Fallback to the reference implementation on CPU.
        x = self.linear(x)
        x = torch.min(x, self.constant)
        x = x - self.constant
        return x
```