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

# Fused kernel: MaxPool1d(kernel_size=2, stride=2) + sum(dim=1) + scaling.
# For each batch row: out[b] = scale * sum_k( max(x[b, 2k], x[b, 2k+1]) )
# The matmul (the FLOP-heavy op) is left to the highly-optimized cuBLAS
# via F.linear; everything after it is fused into a single lightweight
# memory-bound kernel, eliminating intermediate tensors and kernel launches.
fused_ops_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

constexpr int BLOCK_SIZE = 512;

// Vectorized path (requires out_features % 4 == 0): each float4 load
// covers two adjacent max-pool pairs.
__global__ void fused_maxpool_sum_scale_vec_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    int out_features,
    float scale_factor) {

    const int batch = blockIdx.x;
    const float4* row = reinterpret_cast<const float4*>(x + (long)batch * out_features);

    const int num_pairs = out_features >> 1;
    const int num_vec = num_pairs >> 1;

    float local_sum = 0.0f;
    for (int k = threadIdx.x; k < num_vec; k += BLOCK_SIZE) {
        float4 v = row[k];
        local_sum += fmaxf(v.x, v.y) + fmaxf(v.z, v.w);
    }
    // Handle a trailing pair if num_pairs is odd.
    if ((num_pairs & 1) && threadIdx.x == 0) {
        const float* srow = x + (long)batch * out_features;
        const int k = num_pairs - 1;
        local_sum += fmaxf(srow[2 * k], srow[2 * k + 1]);
    }

    __shared__ float shared[BLOCK_SIZE];
    const int tid = threadIdx.x;
    shared[tid] = local_sum;
    __syncthreads();

    for (int s = BLOCK_SIZE >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        out[batch] = shared[0] * scale_factor;
    }
}

// Scalar fallback path for arbitrary out_features.
__global__ void fused_maxpool_sum_scale_scalar_kernel(
    const float* __restrict__ x,
    float* __restrict__ out,
    int out_features,
    float scale_factor) {

    const int batch = blockIdx.x;
    const float* row = x + (long)batch * out_features;
    const int num_pairs = out_features >> 1;

    float local_sum = 0.0f;
    for (int k = threadIdx.x; k < num_pairs; k += BLOCK_SIZE) {
        local_sum += fmaxf(row[2 * k], row[2 * k + 1]);
    }

    __shared__ float shared[BLOCK_SIZE];
    const int tid = threadIdx.x;
    shared[tid] = local_sum;
    __syncthreads();

    for (int s = BLOCK_SIZE >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            shared[tid] += shared[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        out[batch] = shared[0] * scale_factor;
    }
}

torch::Tensor fused_maxpool_sum_scale_cuda(torch::Tensor x, double scale_factor) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (batch, features)");

    auto x_c = x.contiguous();
    const int batch_size = x_c.size(0);
    const int out_features = x_c.size(1);

    auto out = torch::empty({batch_size}, x_c.options());

    if (batch_size == 0 || out_features == 0) {
        out.zero_();
        return out;
    }

    if (out_features % 4 == 0) {
        fused_maxpool_sum_scale_vec_kernel<<<batch_size, BLOCK_SIZE>>>(
            x_c.data_ptr<float>(), out.data_ptr<float>(), out_features,
            static_cast<float>(scale_factor));
    } else {
        fused_maxpool_sum_scale_scalar_kernel<<<batch_size, BLOCK_SIZE>>>(
            x_c.data_ptr<float>(), out.data_ptr<float>(), out_features,
            static_cast<float>(scale_factor));
    }

    return out;
}
"""

fused_ops_cpp_source = (
    "torch::Tensor fused_maxpool_sum_scale_cuda(torch::Tensor x, double scale_factor);"
)

fused_ops = load_inline(
    name="fused_maxpool_sum_scale",
    cpp_sources=fused_ops_cpp_source,
    cuda_sources=fused_ops_source,
    functions=["fused_maxpool_sum_scale_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model: cuBLAS matmul (F.linear) followed by a single fused
    CUDA kernel performing max-pool(k=2) + sum + scale in one pass.
    """
    def __init__(self, in_features, out_features, kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.max_pool = nn.MaxPool1d(kernel_size)
        self.kernel_size = kernel_size
        self.scale_factor = scale_factor
        self.fused_ops = fused_ops

    def forward(self, x):
        x = F.linear(x, self.matmul.weight, self.matmul.bias)
        if (
            self.kernel_size == 2
            and x.is_cuda
            and x.dtype == torch.float32
        ):
            return self.fused_ops.fused_maxpool_sum_scale_cuda(
                x, float(self.scale_factor)
            )
        # Generic fallback for any other configuration.
        x = self.max_pool(x.unsqueeze(1)).squeeze(1)
        x = torch.sum(x, dim=1)
        x = x * self.scale_factor
        return x
```