You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # 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
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # 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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

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

class Model(nn.Module):
    """
    Model that performs a matrix multiplication, subtraction, multiplication, and ReLU activation.
    """
    def __init__(self, in_features, out_features, subtract_value, multiply_value):
        super(Model, self).__init__()
        self.linear = nn.Linear(in_features, out_features)
        self.subtract_value = subtract_value
        self.multiply_value = multiply_value

    def forward(self, x):
        x = self.linear(x)
        x = x - self.subtract_value
        x = x * self.multiply_value
        x = torch.relu(x)
        return x

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]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 4 of 6 for this task. Your previous solution (turn 3) was:
```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Keep TF32 enabled so any torch-matmul fallback matches the reference's
# numerics/speed regime.
try:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# Fully fused implementation of:
#     y = relu((x @ W^T + bias - sub) * mul)
# Using the exact identity:
#     y = relu(mul * (x @ W^T) + mul * (bias - sub))
# the whole forward collapses into ONE cublasLt GEMM with a RELU_BIAS
# epilogue (alpha = mul, bias' = mul * (bias - sub), beta = 0).
#
# Primary path: FP16 tensor-core GEMM with FP32 accumulation and FP32 output
# (same 10-bit input mantissa as the TF32 reference path, but ~2x the math
# throughput). The GEMM algorithm is autotuned once per shape by timing the
# top cublasLt heuristics, then cached. The fp16 copy of the (constant)
# weight and bias' are cached across calls and only rebuilt if the parameter
# storage/version changes, so a steady-state forward is:
#     cast x to fp16 (~35 us) + one fused GEMM.
#
# Fallbacks (all numerically validated at import):
#   1) TF32 cublasLt fused GEMM (turn-2 path)
#   2) torch matmul + custom fused (sub-mul-relu) elementwise kernel
#   3) plain torch
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <vector>
#include <algorithm>

#define CUBLASLT_CHECK(expr)                                            \
  do {                                                                  \
    cublasStatus_t _status = (expr);                                    \
    TORCH_CHECK(_status == CUBLAS_STATUS_SUCCESS,                       \
                "cublasLt call failed with status ", (int)_status);     \
  } while (0)

static cublasLtHandle_t lt_handle() {
  static cublasLtHandle_t handle = nullptr;
  if (handle == nullptr) {
    CUBLASLT_CHECK(cublasLtCreate(&handle));
  }
  return handle;
}

__global__ void bias_prime_kernel(const float* __restrict__ bias,
                                  float* __restrict__ bias_prime,
                                  const float sub,
                                  const float mul,
                                  const int n) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < n) {
    bias_prime[idx] = (bias[idx] - sub) * mul;
  }
}

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

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

// In-place fused (x - sub) * mul -> relu (used by the non-Lt fallback).
torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul) {
  TORCH_CHECK(x.is_cuda(), "Input must be a CUDA tensor");
  TORCH_CHECK(x.scalar_type() == torch::kFloat32, "Input must be float32");
  TORCH_CHECK(x.is_contiguous(), "Input must be contiguous");

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

  const float sub_f = static_cast<float>(sub);
  const float mul_f = static_cast<float>(mul);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const int block_size = 256;

  if ((n % 4) == 0 &&
      (reinterpret_cast<uintptr_t>(x.data_ptr<float>()) % 16) == 0) {
    const long n4 = n / 4;
    long grid = (n4 + block_size - 1) / block_size;
    if (grid > 65535L) grid = 65535L;
    sub_mul_relu_vec4_kernel<<<(unsigned int)grid, block_size, 0, stream>>>(
        reinterpret_cast<float4*>(x.data_ptr<float>()), sub_f, mul_f, n4);
  } else {
    long grid = (n + block_size - 1) / block_size;
    if (grid > 65535L) grid = 65535L;
    sub_mul_relu_scalar_kernel<<<(unsigned int)grid, block_size, 0, stream>>>(
        x.data_ptr<float>(), sub_f, mul_f, n);
  }
  return x;
}

// ---------------------------------------------------------------------------
// FP16-in / FP32-out fused GEMM with RELU_BIAS epilogue, autotuned per shape.
// ---------------------------------------------------------------------------
namespace {
struct Fp16GemmPlan {
  cublasLtMatmulAlgo_t algo;
  size_t workspace = 0;
  int64_t M = -1, N = -1, K = -1;
  bool valid = false;
};
Fp16GemmPlan g_fp16_plan;
}  // namespace

torch::Tensor linear_fp16_relu_bias(torch::Tensor xh, torch::Tensor wh,
                                    torch::Tensor bias_prime, double alpha_d) {
  TORCH_CHECK(xh.is_cuda() && wh.is_cuda() && bias_prime.is_cuda(),
              "All tensors must be CUDA tensors");
  TORCH_CHECK(xh.scalar_type() == torch::kHalf &&
                  wh.scalar_type() == torch::kHalf,
              "x and weight must be float16");
  TORCH_CHECK(bias_prime.scalar_type() == torch::kFloat32,
              "bias_prime must be float32");
  TORCH_CHECK(xh.dim() == 2 && wh.dim() == 2, "x and weight must be 2D");

  auto x = xh.contiguous();
  auto w = wh.contiguous();
  auto bp = bias_prime.contiguous();

  const int64_t M = x.size(0);
  const int64_t K = x.size(1);
  const int64_t N = w.size(0);
  TORCH_CHECK(w.size(1) == K, "weight shape mismatch");
  TORCH_CHECK(bp.numel() == N, "bias shape mismatch");

  auto out = torch::empty({M, N}, x.options().dtype(torch::kFloat32));
  if (M == 0 || N == 0) {
    return out;
  }
  if (K == 0) {
    return torch::relu(bp).unsqueeze(0).expand({M, N}).contiguous();
  }

  cublasLtHandle_t handle = lt_handle();
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  cublasLtMatmulDesc_t desc = nullptr;
  CUBLASLT_CHECK(cublasLtMatmulDescCreate(&desc, CUBLAS_COMPUTE_32F,
                                          CUDA_R_32F));

  // Row-major operands expressed as their column-major transposes:
  //   out^T = W * x^T
  // W is row-major (N x K) == column-major (K x N); op(A)=T gives (N x K).
  // x is row-major (M x K) == column-major (K x M); op(B)=N gives (K x M).
  // out is row-major (M x N) == column-major (N x M).
  cublasOperation_t trans_a = CUBLAS_OP_T;
  cublasOperation_t trans_b = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)));

  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      desc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
  const float* bias_ptr = bp.data_ptr<float>();
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));

  cublasLtMatrixLayout_t a_desc = nullptr, b_desc = nullptr, d_desc = nullptr;
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&a_desc, CUDA_R_16F, K, N, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&b_desc, CUDA_R_16F, K, M, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&d_desc, CUDA_R_32F, N, M, N));

  float alpha = static_cast<float>(alpha_d);
  float beta = 0.0f;
  const void* a_ptr = w.data_ptr();
  const void* b_ptr = x.data_ptr();
  void* d_ptr = out.data_ptr<float>();

  if (!g_fp16_plan.valid || g_fp16_plan.M != M || g_fp16_plan.N != N ||
      g_fp16_plan.K != K) {
    // One-time autotune: time the top cublasLt heuristics and keep the best.
    cublasLtMatmulPreference_t pref = nullptr;
    CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&pref));
    const size_t max_ws = 256ull * 1024ull * 1024ull;
    CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
        pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_ws,
        sizeof(max_ws)));

    const int max_results = 32;
    std::vector<cublasLtMatmulHeuristicResult_t> heur(max_results);
    int count = 0;
    CUBLASLT_CHECK(cublasLtMatmulAlgoGetHeuristic(
        handle, desc, a_desc, b_desc, d_desc, d_desc, pref, max_results,
        heur.data(), &count));
    TORCH_CHECK(count > 0, "cublasLt found no fp16 GEMM algorithm");

    cudaEvent_t ev0, ev1;
    cudaEventCreate(&ev0);
    cudaEventCreate(&ev1);

    float best_ms = 1e30f;
    int best = -1;
    auto ws_opts = x.options().dtype(torch::kUInt8);
    for (int i = 0; i < count; ++i) {
      const size_t ws = heur[i].workspaceSize;
      if (ws > max_ws) continue;
      auto wsbuf = torch::empty({(int64_t)std::max<size_t>(ws, 1)}, ws_opts);
      void* ws_ptr = wsbuf.data_ptr();

      cublasStatus_t st = cublasLtMatmul(
          handle, desc, &alpha, a_ptr, a_desc, b_ptr, b_desc, &beta, d_ptr,
          d_desc, d_ptr, d_desc, &heur[i].algo, ws_ptr, ws, stream);
      if (st != CUBLAS_STATUS_SUCCESS) continue;
      st = cublasLtMatmul(handle, desc, &alpha, a_ptr, a_desc, b_ptr, b_desc,
                          &beta, d_ptr, d_desc, d_ptr, d_desc, &heur[i].algo,
                          ws_ptr, ws, stream);
      if (st != CUBLAS_STATUS_SUCCESS) continue;

      cudaEventRecord(ev0, stream);
      const int reps = 5;
      for (int r = 0; r < reps; ++r) {
        cublasLtMatmul(handle, desc, &alpha, a_ptr, a_desc, b_ptr, b_desc,
                       &beta, d_ptr, d_desc, d_ptr, d_desc, &heur[i].algo,
                       ws_ptr, ws, stream);
      }
      cudaEventRecord(ev1, stream);
      cudaEventSynchronize(ev1);
      float ms = 0.0f;
      cudaEventElapsedTime(&ms, ev0, ev1);
      if (ms < best_ms) {
        best_ms = ms;
        best = i;
      }
    }

    cudaEventDestroy(ev0);
    cudaEventDestroy(ev1);
    cublasLtMatmulPreferenceDestroy(pref);

    TORCH_CHECK(best >= 0, "all fp16 GEMM algorithms failed");
    g_fp16_plan.algo = heur[best].algo;
    g_fp16_plan.workspace = heur[best].workspaceSize;
    g_fp16_plan.M = M;
    g_fp16_plan.N = N;
    g_fp16_plan.K = K;
    g_fp16_plan.valid = true;
  }

  const size_t ws = g_fp16_plan.workspace;
  auto wsbuf = torch::empty({(int64_t)std::max<size_t>(ws, 1)},
                            x.options().dtype(torch::kUInt8));

  cublasStatus_t status = cublasLtMatmul(
      handle, desc, &alpha, a_ptr, a_desc, b_ptr, b_desc, &beta, d_ptr, d_desc,
      d_ptr, d_desc, &g_fp16_plan.algo, wsbuf.data_ptr(), ws, stream);

  cublasLtMatrixLayoutDestroy(a_desc);
  cublasLtMatrixLayoutDestroy(b_desc);
  cublasLtMatrixLayoutDestroy(d_desc);
  cublasLtMatmulDescDestroy(desc);

  TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
              "fp16 cublasLtMatmul failed with status ", (int)status);
  return out;
}

// ---------------------------------------------------------------------------
// TF32 fused GEMM fallback (single heuristic), same math as the fp16 path.
// ---------------------------------------------------------------------------
torch::Tensor linear_sub_mul_relu_cuda(torch::Tensor x, torch::Tensor weight,
                                       torch::Tensor bias,
                                       double subtract_value,
                                       double multiply_value) {
  TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
              "All tensors must be CUDA tensors");
  TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                  weight.scalar_type() == torch::kFloat32 &&
                  bias.scalar_type() == torch::kFloat32,
              "All tensors must be float32");
  TORCH_CHECK(x.dim() >= 1, "x must have at least 1 dimension");

  const int64_t K = x.size(-1);
  const int64_t N = weight.size(0);
  TORCH_CHECK(weight.dim() == 2 && weight.size(1) == K, "weight shape mismatch");
  TORCH_CHECK(bias.numel() == N, "bias shape mismatch");

  std::vector<int64_t> out_sizes(x.sizes().begin(), x.sizes().end());
  out_sizes.back() = N;

  auto x2d = x.reshape({-1, K}).contiguous();
  auto w = weight.contiguous();
  auto b = bias.contiguous();

  const int64_t M = x2d.size(0);
  auto out = torch::empty({M, N}, x2d.options());
  if (M == 0 || N == 0) {
    return out.reshape(out_sizes);
  }

  const float sub_f = static_cast<float>(subtract_value);
  const float mul_f = static_cast<float>(multiply_value);

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  auto bias_prime = torch::empty({N}, x2d.options());
  const int threads = 256;
  const int blocks = (int)((N + threads - 1) / threads);
  bias_prime_kernel<<<blocks, threads, 0, stream>>>(
      b.data_ptr<float>(), bias_prime.data_ptr<float>(), sub_f, mul_f, (int)N);

  if (K == 0) {
    auto y = torch::relu(bias_prime).expand({M, N}).contiguous();
    return y.reshape(out_sizes);
  }

  cublasLtHandle_t handle = lt_handle();

  cublasLtMatmulDesc_t op_desc = nullptr;
  cublasLtMatrixLayout_t a_desc = nullptr, b_desc = nullptr, d_desc = nullptr;
  cublasLtMatmulPreference_t pref = nullptr;

  CUBLASLT_CHECK(cublasLtMatmulDescCreate(&op_desc, CUBLAS_COMPUTE_32F_FAST_TF32,
                                          CUDA_R_32F));

  cublasOperation_t trans_a = CUBLAS_OP_T;
  cublasOperation_t trans_b = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)));

  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
  const float* bias_ptr = bias_prime.data_ptr<float>();
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));

  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&a_desc, CUDA_R_32F, K, N, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&b_desc, CUDA_R_32F, K, M, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&d_desc, CUDA_R_32F, N, M, N));

  const size_t max_workspace = 32ull * 1024ull * 1024ull;
  CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&pref));
  CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
      pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace,
      sizeof(max_workspace)));

  cublasLtMatmulHeuristicResult_t heuristic;
  int algo_count = 0;
  CUBLASLT_CHECK(cublasLtMatmulAlgoGetHeuristic(
      handle, op_desc, a_desc, b_desc, d_desc, d_desc, pref, 1, &heuristic,
      &algo_count));
  TORCH_CHECK(algo_count > 0, "cublasLt found no suitable GEMM algorithm");

  size_t workspace_size = heuristic.workspaceSize;
  if (workspace_size > max_workspace) workspace_size = max_workspace;
  auto workspace = torch::empty(
      {(int64_t)(workspace_size > 0 ? workspace_size : 1)},
      x2d.options().dtype(torch::kUInt8));

  float alpha = mul_f;
  float beta = 0.0f;

  cublasStatus_t status = cublasLtMatmul(
      handle, op_desc, &alpha,
      w.data_ptr<float>(), a_desc,
      x2d.data_ptr<float>(), b_desc,
      &beta,
      out.data_ptr<float>(), d_desc,
      out.data_ptr<float>(), d_desc,
      &heuristic.algo, workspace.data_ptr(), workspace_size, stream);

  cublasLtMatmulPreferenceDestroy(pref);
  cublasLtMatrixLayoutDestroy(a_desc);
  cublasLtMatrixLayoutDestroy(b_desc);
  cublasLtMatrixLayoutDestroy(d_desc);
  cublasLtMatmulDescDestroy(op_desc);

  TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
              "cublasLtMatmul failed with status ", (int)status);

  return out.reshape(out_sizes);
}
"""

_cpp_source = """
torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul);
torch::Tensor linear_sub_mul_relu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double subtract_value, double multiply_value);
torch::Tensor linear_fp16_relu_bias(torch::Tensor xh, torch::Tensor wh, torch::Tensor bias_prime, double alpha);
"""

try:
    _ext = load_inline(
        name="fused_linear_fp16_relu_v3",
        cpp_sources=_cpp_source,
        cuda_sources=_cuda_source,
        functions=[
            "sub_mul_relu_cuda",
            "linear_sub_mul_relu_cuda",
            "linear_fp16_relu_bias",
        ],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--use_fast_math"],
        extra_ldflags=["-lcublasLt", "-lcublas"],
    )
except Exception:
    _ext = None


def _self_test_fp16():
    """Validate the fused fp16 path against the plain torch computation.

    Uses the architecture's real shape/distribution (where a tight tolerance
    is required) plus a harder randn case (looser tolerance) to guard against
    any epilogue-semantics or environment surprises. The real-shape run also
    triggers the one-time GEMM autotune outside of the timed region.
    """
    if _ext is None or not torch.cuda.is_available():
        return False
    try:
        sub, mul = 2.0, 1.5
        M, K, N = 1024, 8192, 8192
        x = torch.rand(M, K, device="cuda", dtype=torch.float32)
        w = (torch.rand(N, K, device="cuda", dtype=torch.float32) * 2.0 - 1.0) / (K ** 0.5)
        b = (torch.rand(N, device="cuda", dtype=torch.float32) * 2.0 - 1.0) / (K ** 0.5)
        ref = torch.relu((F.linear(x, w, b) - sub) * mul)
        bp = ((b - sub) * mul).contiguous()
        got = _ext.linear_fp16_relu_bias(x.half(), w.half(), bp, mul)
        torch.cuda.synchronize()
        if not torch.allclose(ref, got, rtol=1e-3, atol=1e-3):
            return False
        del x, w, b, ref, got, bp

        Ms, Ks, Ns = 192, 1536, 768
        w = torch.randn(Ns, Ks, device="cuda", dtype=torch.float32) / (Ks ** 0.5)
        b = torch.randn(Ns, device="cuda", dtype=torch.float32) * 0.1
        bp = ((b - sub) * mul).contiguous()
        for x in (torch.rand(Ms, Ks, device="cuda", dtype=torch.float32),
                  torch.randn(Ms, Ks, device="cuda", dtype=torch.float32)):
            ref = torch.relu((F.linear(x, w, b) - sub) * mul)
            got = _ext.linear_fp16_relu_bias(x.half(), w.half(), bp, mul)
            torch.cuda.synchronize()
            if not torch.allclose(ref, got, rtol=2e-2, atol=2e-2):
                return False
        return True
    except Exception:
        return False


def _self_test_tf32():
    """Validate the fused TF32 cublasLt path numerically."""
    if _ext is None or not torch.cuda.is_available():
        return False
    try:
        x = torch.randn(8, 16, device="cuda", dtype=torch.float32)
        w = torch.randn(12, 16, device="cuda", dtype=torch.float32)
        b = torch.randn(12, device="cuda", dtype=torch.float32)
        sub, mul = 0.75, 1.25
        ref = torch.relu((F.linear(x, w, b) - sub) * mul)
        got = _ext.linear_sub_mul_relu_cuda(x, w, b, sub, mul)
        torch.cuda.synchronize()
        return bool(torch.allclose(ref, got, rtol=1e-2, atol=1e-2))
    except Exception:
        return False


_FP16_OK = _self_test_fp16()
_TF32_OK = _self_test_tf32()


class ModelNew(nn.Module):
    """
    Optimized model: the entire forward (linear + subtract + multiply + ReLU)
    is executed as one fused GEMM with a ReLU-bias epilogue, using:
        relu((x@W^T + b - s) * m) == relu(m * (x@W^T) + m*(b - s))
    Primary path uses an autotuned FP16-tensor-core / FP32-accumulate /
    FP32-output cublasLt GEMM (numerically equivalent-class to the TF32
    reference but ~2x math throughput), with the fp16 weight and folded bias
    cached across calls. Robust fallbacks: fused TF32 GEMM, then torch
    matmul + fused elementwise kernel, then plain torch.
    """
    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)
        self._w_half = None
        self._bias_prime = None
        self._cache_key = None
        self._use_fp16 = _FP16_OK
        self._use_tf32 = _TF32_OK

        # Warm the fused path (and its one-time autotune for the expected
        # shape) here so it does not land inside the timed region.
        if self._use_fp16 and torch.cuda.is_available():
            try:
                dx = torch.zeros(1024, in_features, device="cuda",
                                 dtype=torch.float16)
                dw = torch.zeros(out_features, in_features, device="cuda",
                                 dtype=torch.float16)
                db = torch.zeros(out_features, device="cuda",
                                 dtype=torch.float32)
                _ext.linear_fp16_relu_bias(dx, dw, db, self.multiply_value)
                torch.cuda.synchronize()
            except Exception:
                self._use_fp16 = False

    def forward(self, x):
        w = self.linear.weight
        if (self._use_fp16 and x.is_cuda and w.is_cuda
                and x.dtype == torch.float32 and w.dtype == torch.float32):
            try:
                b = self.linear.bias
                key = (w.data_ptr(), w._version, b._version)
                if key != self._cache_key:
                    self._w_half = w.detach().contiguous().half()
                    self._bias_prime = (
                        (b.detach() - self.subtract_value)
                        * self.multiply_value).contiguous()
                    self._cache_key = key
                K = x.shape[-1]
                xh = x.contiguous().half()
                out = _ext.linear_fp16_relu_bias(
                    xh.view(-1, K), self._w_half, self._bias_prime,
                    self.multiply_value)
                return out.view(x.shape[:-1] + (self.linear.out_features,))
            except Exception:
                self._use_fp16 = False

        if (self._use_tf32 and x.is_cuda and w.is_cuda
                and x.dtype == torch.float32 and w.dtype == torch.float32):
            try:
                return _ext.linear_sub_mul_relu_cuda(
                    x, w, self.linear.bias,
                    self.subtract_value, self.multiply_value)
            except Exception:
                self._use_tf32 = False

        y = F.linear(x, w, self.linear.bias)
        if (_ext is not None and y.is_cuda and y.dtype == torch.float32
                and y.is_contiguous()):
            return _ext.sub_mul_relu_cuda(
                y, self.subtract_value, self.multiply_value)
        return torch.relu((y - self.subtract_value) * self.multiply_value)
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 1.26 ms; the PyTorch reference (with TF32 enabled) takes 1.27 ms; speedup 1.008x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| void cutlass::Kernel2<cutlass_80_tensorop_s1688gemm_256x128_32x3_tn_align4>(cutlass_80_ten | 1061.5 | 1 |
| void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor<float, float,  | 51.2 | 1 |
| void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespace)::launc | 50.3 | 1 |
| void at::native::vectorized_elementwise_kernel<4, at::native::CUDAFunctorOnSelf_add<float> | 40.4 | 1 |
| Memset (Device) | 1.2 | 1 |

The fastest correct solution so far (1.124x speedup) is:
```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# The PyTorch reference is timed with TF32 enabled, so make sure TF32 is on
# for any fallback torch matmul path as well. The main path below uses
# cublasLt with CUBLAS_COMPUTE_32F_FAST_TF32 explicitly, independent of
# global flags.
try:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.set_float32_matmul_precision("high")
except Exception:
    pass

# Fully fused implementation:
#   y = relu((x @ W^T + bias - sub) * mul)
# Algebraically:
#   y = relu( mul * (x @ W^T) + mul*(bias - sub) )
# So we run a single cublasLt GEMM (TF32 tensor cores) with:
#   alpha = mul, epilogue = RELU_BIAS, bias' = mul*(bias - sub)
# i.e. the matmul + bias + subtract + multiply + ReLU all execute in ONE
# kernel launch. A tiny kernel precomputes bias' (out_features elements).
_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublasLt.h>

#define CUBLASLT_CHECK(expr)                                            \
  do {                                                                  \
    cublasStatus_t _status = (expr);                                    \
    TORCH_CHECK(_status == CUBLAS_STATUS_SUCCESS,                       \
                "cublasLt call failed with status ", (int)_status);     \
  } while (0)

__global__ void bias_prime_kernel(const float* __restrict__ bias,
                                  float* __restrict__ bias_prime,
                                  const float sub,
                                  const float mul,
                                  const int n) {
  int idx = blockIdx.x * blockDim.x + threadIdx.x;
  if (idx < n) {
    bias_prime[idx] = (bias[idx] - sub) * mul;
  }
}

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

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

// In-place fused (x - sub) * mul -> relu epilogue (fallback path).
torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul) {
  TORCH_CHECK(x.is_cuda(), "Input must be a CUDA tensor");
  TORCH_CHECK(x.scalar_type() == torch::kFloat32, "Input must be float32");
  TORCH_CHECK(x.is_contiguous(), "Input must be contiguous");

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

  const float sub_f = static_cast<float>(sub);
  const float mul_f = static_cast<float>(mul);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  const int block_size = 256;

  if ((n % 4) == 0 &&
      (reinterpret_cast<uintptr_t>(x.data_ptr<float>()) % 16) == 0) {
    const long n4 = n / 4;
    long grid = (n4 + block_size - 1) / block_size;
    if (grid > 65535L) grid = 65535L;
    sub_mul_relu_vec4_kernel<<<(unsigned int)grid, block_size, 0, stream>>>(
        reinterpret_cast<float4*>(x.data_ptr<float>()), sub_f, mul_f, n4);
  } else {
    long grid = (n + block_size - 1) / block_size;
    if (grid > 65535L) grid = 65535L;
    sub_mul_relu_scalar_kernel<<<(unsigned int)grid, block_size, 0, stream>>>(
        x.data_ptr<float>(), sub_f, mul_f, n);
  }
  return x;
}

// Fully fused: relu(mul * (x @ W^T) + mul*(bias - sub)) via one TF32 GEMM.
torch::Tensor linear_sub_mul_relu_cuda(torch::Tensor x, torch::Tensor weight,
                                       torch::Tensor bias,
                                       double subtract_value,
                                       double multiply_value) {
  TORCH_CHECK(x.is_cuda() && weight.is_cuda() && bias.is_cuda(),
              "All tensors must be CUDA tensors");
  TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                  weight.scalar_type() == torch::kFloat32 &&
                  bias.scalar_type() == torch::kFloat32,
              "All tensors must be float32");
  TORCH_CHECK(x.dim() >= 1, "x must have at least 1 dimension");

  const int64_t K = x.size(-1);
  const int64_t N = weight.size(0);
  TORCH_CHECK(weight.dim() == 2 && weight.size(1) == K, "weight shape mismatch");
  TORCH_CHECK(bias.numel() == N, "bias shape mismatch");

  std::vector<int64_t> out_sizes(x.sizes().begin(), x.sizes().end());
  out_sizes.back() = N;

  auto x2d = x.reshape({-1, K}).contiguous();
  auto w = weight.contiguous();
  auto b = bias.contiguous();

  const int64_t M = x2d.size(0);
  auto out = torch::empty({M, N}, x2d.options());
  if (M == 0 || N == 0) {
    return out.reshape(out_sizes);
  }

  const float sub_f = static_cast<float>(subtract_value);
  const float mul_f = static_cast<float>(multiply_value);

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  auto bias_prime = torch::empty({N}, x2d.options());
  const int threads = 256;
  const int blocks = (int)((N + threads - 1) / threads);
  bias_prime_kernel<<<blocks, threads, 0, stream>>>(
      b.data_ptr<float>(), bias_prime.data_ptr<float>(), sub_f, mul_f, (int)N);

  if (K == 0) {
    auto y = torch::relu(bias_prime).expand({M, N}).contiguous();
    return y.reshape(out_sizes);
  }

  static cublasLtHandle_t handle = nullptr;
  if (handle == nullptr) {
    CUBLASLT_CHECK(cublasLtCreate(&handle));
  }

  cublasLtMatmulDesc_t op_desc = nullptr;
  cublasLtMatrixLayout_t a_desc = nullptr, b_desc = nullptr, d_desc = nullptr;
  cublasLtMatmulPreference_t pref = nullptr;

  CUBLASLT_CHECK(cublasLtMatmulDescCreate(&op_desc, CUBLAS_COMPUTE_32F_FAST_TF32,
                                          CUDA_R_32F));

  // Row-major operands expressed as their column-major transposes:
  //   out^T = W * x^T
  // W is row-major (N x K) == column-major (K x N); op(A)=T gives (N x K).
  // x is row-major (M x K) == column-major (K x M); op(B)=N gives (K x M).
  // out is row-major (M x N) == column-major (N x M).
  cublasOperation_t trans_a = CUBLAS_OP_T;
  cublasOperation_t trans_b = CUBLAS_OP_N;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a, sizeof(trans_a)));
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)));

  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
  const float* bias_ptr = bias_prime.data_ptr<float>();
  CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(
      op_desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias_ptr, sizeof(bias_ptr)));

  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&a_desc, CUDA_R_32F, K, N, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&b_desc, CUDA_R_32F, K, M, K));
  CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&d_desc, CUDA_R_32F, N, M, N));

  const size_t max_workspace = 32ull * 1024ull * 1024ull;
  CUBLASLT_CHECK(cublasLtMatmulPreferenceCreate(&pref));
  CUBLASLT_CHECK(cublasLtMatmulPreferenceSetAttribute(
      pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace,
      sizeof(max_workspace)));

  cublasLtMatmulHeuristicResult_t heuristic;
  int algo_count = 0;
  CUBLASLT_CHECK(cublasLtMatmulAlgoGetHeuristic(
      handle, op_desc, a_desc, b_desc, d_desc, d_desc, pref, 1, &heuristic,
      &algo_count));
  TORCH_CHECK(algo_count > 0, "cublasLt found no suitable GEMM algorithm");

  size_t workspace_size = heuristic.workspaceSize;
  if (workspace_size > max_workspace) workspace_size = max_workspace;
  auto workspace = torch::empty(
      {(int64_t)(workspace_size > 0 ? workspace_size : 1)},
      x2d.options().dtype(torch::kUInt8));

  float alpha = mul_f;
  float beta = 0.0f;

  cublasStatus_t status = cublasLtMatmul(
      handle, op_desc, &alpha,
      w.data_ptr<float>(), a_desc,
      x2d.data_ptr<float>(), b_desc,
      &beta,
      out.data_ptr<float>(), d_desc,
      out.data_ptr<float>(), d_desc,
      &heuristic.algo, workspace.data_ptr(), workspace_size, stream);

  cublasLtMatmulPreferenceDestroy(pref);
  cublasLtMatrixLayoutDestroy(a_desc);
  cublasLtMatrixLayoutDestroy(b_desc);
  cublasLtMatrixLayoutDestroy(d_desc);
  cublasLtMatmulDescDestroy(op_desc);

  TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS,
              "cublasLtMatmul failed with status ", (int)status);

  return out.reshape(out_sizes);
}
"""

_cpp_source = """
torch::Tensor sub_mul_relu_cuda(torch::Tensor x, double sub, double mul);
torch::Tensor linear_sub_mul_relu_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, double subtract_value, double multiply_value);
"""

try:
    _ext = load_inline(
        name="fused_linear_sub_mul_relu_v2",
        cpp_sources=_cpp_source,
        cuda_sources=_cuda_source,
        functions=["sub_mul_relu_cuda", "linear_sub_mul_relu_cuda"],
        verbose=False,
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--use_fast_math"],
        extra_ldflags=["-lcublasLt", "-lcublas"],
    )
except Exception:
    _ext = None


def _lt_path_ok():
    """Verify the fused cublasLt path numerically against the plain torch
    computation on a tiny problem; guards against any epilogue-semantics or
    environment surprises."""
    if _ext is None or not torch.cuda.is_available():
        return False
    try:
        x = torch.randn(8, 16, device="cuda", dtype=torch.float32)
        w = torch.randn(12, 16, device="cuda", dtype=torch.float32)
        b = torch.randn(12, device="cuda", dtype=torch.float32)
        sub, mul = 0.75, 1.25
        ref = torch.relu((F.linear(x, w, b) - sub) * mul)
        got = _ext.linear_sub_mul_relu_cuda(x, w, b, sub, mul)
        torch.cuda.synchronize()
        return bool(torch.allclose(ref, got, rtol=1e-2, atol=1e-2))
    except Exception:
        return False


_USE_LT = _lt_path_ok()


class ModelNew(nn.Module):
    """
    Optimized model: the entire forward (linear + subtract + multiply + ReLU)
    is executed as a single fused TF32 cublasLt GEMM with a ReLU-bias
    epilogue, using the identity:
        relu((x@W^T + b - s) * m) == relu(m * (x@W^T) + m*(b - s))
    Falls back to a fused elementwise kernel (or plain torch) if the fused
    GEMM path is unavailable.
    """
    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):
        if _USE_LT and x.is_cuda and x.dtype == torch.float32:
            return _ext.linear_sub_mul_relu_cuda(
                x, self.linear.weight, self.linear.bias,
                self.subtract_value, self.multiply_value)

        y = F.linear(x, self.linear.weight, self.linear.bias)
        if (_ext is not None and y.is_cuda and y.dtype == torch.float32
                and y.is_contiguous()):
            return _ext.sub_mul_relu_cuda(
                y, self.subtract_value, self.multiply_value)
        return torch.relu((y - self.subtract_value) * self.multiply_value)
```

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.
